github.com/iikira/iikira-go-utils@v0.0.0-20230610031953-f2cb11cde33a/requester/multipartreader/multipartreader.go (about)

     1  // Package multipartreader helps you encode large files in MIME multipart format
     2  // without reading the entire content into memory.
     3  package multipartreader
     4  
     5  import (
     6  	"bytes"
     7  	"errors"
     8  	"fmt"
     9  	"github.com/iikira/iikira-go-utils/requester/rio"
    10  	"io"
    11  	"mime/multipart"
    12  	"strings"
    13  	"sync"
    14  	"sync/atomic"
    15  )
    16  
    17  type (
    18  	// MultipartReader MIME multipart format
    19  	MultipartReader struct {
    20  		length      int64
    21  		contentType string
    22  		boundary    string
    23  
    24  		formBody  string
    25  		parts     []*part
    26  		part64s   []*part64
    27  		formClose string
    28  
    29  		mu          sync.Mutex
    30  		closed      bool
    31  		multiReader io.Reader
    32  	}
    33  
    34  	part struct {
    35  		form      string
    36  		readerlen rio.ReaderLen
    37  	}
    38  
    39  	part64 struct {
    40  		form        string
    41  		readerlen64 rio.ReaderLen64
    42  	}
    43  )
    44  
    45  // NewMultipartReader 返回初始化的 *MultipartReader
    46  func NewMultipartReader() (mr *MultipartReader) {
    47  	builder := &strings.Builder{}
    48  	writer := multipart.NewWriter(builder)
    49  	mr = &MultipartReader{
    50  		contentType: writer.FormDataContentType(),
    51  		boundary:    writer.Boundary(),
    52  	}
    53  
    54  	mr.length += int64(builder.Len())
    55  	mr.formBody = builder.String()
    56  	return
    57  }
    58  
    59  // AddFormField 增加 form 表单
    60  func (mr *MultipartReader) AddFormField(fieldname string, readerlen rio.ReaderLen) {
    61  	if readerlen == nil {
    62  		return
    63  	}
    64  
    65  	mpart := &part{
    66  		form:      fmt.Sprintf("--%s\r\nContent-Disposition: form-data; name=\"%s\"\r\n\r\n", mr.boundary, fieldname),
    67  		readerlen: readerlen,
    68  	}
    69  	atomic.AddInt64(&mr.length, int64(len(mpart.form)+mpart.readerlen.Len()))
    70  	mr.parts = append(mr.parts, mpart)
    71  }
    72  
    73  func (mr *MultipartReader) AddFormFieldBytes(fieldname string, b []byte) {
    74  	mr.AddFormField(fieldname, bytes.NewReader(b))
    75  }
    76  
    77  func (mr *MultipartReader) AddFormFieldString(fieldname string, s string) {
    78  	mr.AddFormField(fieldname, strings.NewReader(s))
    79  }
    80  
    81  // AddFormFile 增加 form 文件表单
    82  func (mr *MultipartReader) AddFormFile(fieldname, filename string, readerlen64 rio.ReaderLen64) {
    83  	if readerlen64 == nil {
    84  		return
    85  	}
    86  
    87  	mpart64 := &part64{
    88  		form:        fmt.Sprintf("--%s\r\nContent-Disposition: form-data; name=\"%s\"; filename=\"%s\"\r\n\r\n", mr.boundary, fieldname, filename),
    89  		readerlen64: readerlen64,
    90  	}
    91  	atomic.AddInt64(&mr.length, int64(len(mpart64.form))+mpart64.readerlen64.Len())
    92  	mr.part64s = append(mr.part64s, mpart64)
    93  }
    94  
    95  //CloseMultipart 关闭multipartreader
    96  func (mr *MultipartReader) CloseMultipart() error {
    97  	mr.mu.Lock()
    98  	defer mr.mu.Unlock()
    99  	if mr.closed {
   100  		return errors.New("multipartreader already closed")
   101  	}
   102  
   103  	mr.formClose = "\r\n--" + mr.boundary + "--\r\n"
   104  	atomic.AddInt64(&mr.length, int64(len(mr.formClose)))
   105  
   106  	numReaders := 0
   107  	if mr.formBody != "" {
   108  		numReaders++
   109  	}
   110  	numReaders += 2*len(mr.parts) + 2*len(mr.part64s)
   111  	if mr.formClose != "" {
   112  		numReaders++
   113  	}
   114  
   115  	readers := make([]io.Reader, 0, numReaders)
   116  	readers = append(readers, strings.NewReader(mr.formBody))
   117  	for k := range mr.parts {
   118  		readers = append(readers, strings.NewReader(mr.parts[k].form), mr.parts[k].readerlen)
   119  	}
   120  	for k := range mr.part64s {
   121  		readers = append(readers, strings.NewReader(mr.part64s[k].form), mr.part64s[k].readerlen64)
   122  	}
   123  	readers = append(readers, strings.NewReader(mr.formClose))
   124  	mr.multiReader = io.MultiReader(readers...)
   125  
   126  	mr.closed = true
   127  	return nil
   128  }
   129  
   130  //ContentType 返回Content-Type
   131  func (mr *MultipartReader) ContentType() string {
   132  	return mr.contentType
   133  }
   134  
   135  func (mr *MultipartReader) Read(p []byte) (n int, err error) {
   136  	if !mr.closed {
   137  		return 0, errors.New("multipartreader not closed")
   138  	}
   139  	n, err = mr.multiReader.Read(p)
   140  	return n, err
   141  }
   142  
   143  // Len 返回表单内容总长度
   144  func (mr *MultipartReader) Len() int64 {
   145  	return atomic.LoadInt64(&mr.length)
   146  }