github.com/qiniu/go-sdk/v7@v7.27.0/internal/clientv2/multipart.go (about)

     1  package clientv2
     2  
     3  import (
     4  	"crypto/rand"
     5  	"fmt"
     6  	"io"
     7  	"mime/multipart"
     8  	"net/textproto"
     9  	"strings"
    10  	"sync"
    11  
    12  	"github.com/qiniu/go-sdk/v7/internal/context"
    13  	compatible_io "github.com/qiniu/go-sdk/v7/internal/io"
    14  )
    15  
    16  type (
    17  	keyValuePair struct {
    18  		key, value string
    19  	}
    20  	keyFilePair struct {
    21  		key, fileName, contentType string
    22  		stream                     compatible_io.ReadSeekCloser
    23  	}
    24  
    25  	MultipartForm struct {
    26  		values       []keyValuePair
    27  		files        []keyFilePair
    28  		ctx          context.Context
    29  		cancel       context.CancelCauseFunc
    30  		w            *io.PipeWriter
    31  		boundary     string
    32  		boundaryOnce sync.Once
    33  	}
    34  
    35  	multipartFormReader struct {
    36  		multipartWriter *multipart.Writer
    37  		form            *MultipartForm
    38  		r               *io.PipeReader
    39  		closed          bool
    40  	}
    41  )
    42  
    43  func (f *MultipartForm) SetValue(key, value string) *MultipartForm {
    44  	f.values = append(f.values, keyValuePair{key, value})
    45  	return f
    46  }
    47  
    48  func (f *MultipartForm) SetFile(key, fileName, contentType string, stream compatible_io.ReadSeekCloser) *MultipartForm {
    49  	f.files = append(f.files, keyFilePair{key, fileName, contentType, stream})
    50  	return f
    51  }
    52  
    53  func (f *MultipartForm) generateBoundary() string {
    54  	f.boundaryOnce.Do(func() {
    55  		var buf [30]byte
    56  		_, err := io.ReadFull(rand.Reader, buf[:])
    57  		if err != nil {
    58  			panic(err)
    59  		}
    60  		f.boundary = fmt.Sprintf("%x", buf[:])
    61  	})
    62  	return f.boundary
    63  }
    64  
    65  func newMultipartFormReader(form *MultipartForm) *multipartFormReader {
    66  	reader := &multipartFormReader{form: form}
    67  	reader.r, form.w = io.Pipe()
    68  	reader.multipartWriter = multipart.NewWriter(form.w)
    69  	reader.multipartWriter.SetBoundary(form.generateBoundary())
    70  
    71  	go func(multipartWriter *multipart.Writer, w *io.PipeWriter, ctx context.Context, cancel context.CancelCauseFunc) {
    72  		defer w.Close()
    73  		defer multipartWriter.Close()
    74  
    75  		for _, pair := range form.values {
    76  			select {
    77  			case <-ctx.Done():
    78  				return
    79  			default:
    80  				if err := multipartWriter.WriteField(pair.key, pair.value); err != nil {
    81  					cancel(err)
    82  					return
    83  				}
    84  			}
    85  		}
    86  		for _, pair := range form.files {
    87  			select {
    88  			case <-ctx.Done():
    89  				return
    90  			default:
    91  				if err := reader.createFormFile(pair.key, pair.fileName, pair.contentType, pair.stream); err != nil {
    92  					cancel(err)
    93  					return
    94  				}
    95  			}
    96  		}
    97  	}(reader.multipartWriter, form.w, form.ctx, form.cancel)
    98  
    99  	return reader
   100  }
   101  
   102  func (r *multipartFormReader) Read(p []byte) (int, error) {
   103  	select {
   104  	case <-r.form.ctx.Done():
   105  		return 0, context.Cause(r.form.ctx)
   106  	default:
   107  		return r.r.Read(p)
   108  	}
   109  }
   110  
   111  func (r *multipartFormReader) Close() error {
   112  	if r.closed {
   113  		return nil
   114  	}
   115  	r.closed = true
   116  	r.form.cancel(io.ErrClosedPipe)
   117  	err := r.r.Close()
   118  	for _, pair := range r.form.files {
   119  		if e := pair.stream.Close(); e != nil && err == nil {
   120  			err = e
   121  		}
   122  	}
   123  	return err
   124  }
   125  
   126  func (r *multipartFormReader) formDataContentType() string {
   127  	return r.multipartWriter.FormDataContentType()
   128  }
   129  
   130  func (r *multipartFormReader) createFormFile(fieldName, fileName, contentType string, stream compatible_io.ReadSeekCloser) error {
   131  	h := make(textproto.MIMEHeader)
   132  	h.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"; filename="%s"`, escapeQuotes(fieldName), escapeQuotes(fileName)))
   133  	h.Set("Content-Type", contentType)
   134  	if w, err := r.multipartWriter.CreatePart(h); err != nil {
   135  		return err
   136  	} else if _, err := io.Copy(w, stream); err != nil {
   137  		return err
   138  	}
   139  	return nil
   140  }
   141  
   142  func GetMultipartFormRequestBody(info *MultipartForm) GetRequestBody {
   143  	return func(o *RequestParams) (io.ReadCloser, error) {
   144  		if cancel := info.cancel; cancel != nil {
   145  			cancel(io.ErrClosedPipe)
   146  		}
   147  		info.ctx, info.cancel = context.WithCancelCause(context.Background())
   148  
   149  		if w := info.w; w != nil {
   150  			w.Close()
   151  			info.w = nil
   152  		}
   153  
   154  		for _, pair := range info.files {
   155  			if _, err := pair.stream.Seek(0, io.SeekStart); err != nil {
   156  				return nil, err
   157  			}
   158  		}
   159  		r := newMultipartFormReader(info)
   160  		o.Header.Set("Content-Type", r.formDataContentType())
   161  		return r, nil
   162  	}
   163  }
   164  
   165  var quoteEscaper = strings.NewReplacer("\\", "\\\\", `"`, "\\\"")
   166  
   167  func escapeQuotes(s string) string {
   168  	return quoteEscaper.Replace(s)
   169  }