fix(#433): stream multipart uploads without buffering files

client.go

@@ -1742,6 +1742,9 @@ func (c *Client) roundTrip(r *Request) (resp *Response, err error) {
 
 	// setup header
 	contentLength := int64(len(r.Body))
+	if r.contentLength != 0 {
+		contentLength = r.contentLength
+	}
 
 	var reqBody io.ReadCloser
 	if r.GetBody != nil {
@@ -1750,6 +1753,10 @@ func (c *Client) roundTrip(r *Request) (resp *Response, err error) {
 			return
 		}
 	}
+	getBody := r.GetBody
+	if r.unReplayableBody != nil {
+		getBody = nil
+	}
 	req := &http.Request{
 		Method:        r.Method,
 		Header:        r.Headers.Clone(),
@@ -1760,7 +1767,7 @@ func (c *Client) roundTrip(r *Request) (resp *Response, err error) {
 		ProtoMinor:    1,
 		ContentLength: contentLength,
 		Body:          reqBody,
-		GetBody:       r.GetBody,
+		GetBody:       getBody,
 		Close:         r.close,
 	}
 	for _, cookie := range r.Cookies {

middleware.go

@@ -123,29 +123,95 @@ func writeMultipartFormFile(w *multipart.Writer, file *FileUpload, r *Request) e
 	return err
 }
 
-func writeMultiPart(r *Request, w *multipart.Writer) {
-	defer w.Close() // close multipart to write tailer boundary
+func writeMultiPart(r *Request, w *multipart.Writer) error {
 	if len(r.FormData) > 0 {
 		for k, vs := range r.FormData {
 			for _, v := range vs {
-				w.WriteField(k, v)
+				if err := w.WriteField(k, v); err != nil {
+					return err
+				}
 			}
 		}
 	} else if len(r.OrderedFormData) > 0 {
 		if len(r.OrderedFormData)%2 != 0 {
-			r.error = errBadOrderedFormData
-			return
+			return errBadOrderedFormData
 		}
 		maxIndex := len(r.OrderedFormData) - 2
 		for i := 0; i <= maxIndex; i += 2 {
 			key := r.OrderedFormData[i]
 			value := r.OrderedFormData[i+1]
-			w.WriteField(key, value)
+			if err := w.WriteField(key, value); err != nil {
+				return err
+			}
 		}
 	}
 	for _, file := range r.uploadFiles {
-		writeMultipartFormFile(w, file, r)
+		if err := writeMultipartFormFile(w, file, r); err != nil {
+			return err
+		}
 	}
+	return w.Close() // close multipart to write trailing boundary
+}
+
+type countingWriter struct {
+	n int64
+}
+
+func (w *countingWriter) Write(p []byte) (int, error) {
+	w.n += int64(len(p))
+	return len(p), nil
+}
+
+func multipartContentLength(r *Request, boundary string) (int64, bool, error) {
+	cw := new(countingWriter)
+	w := multipart.NewWriter(cw)
+	if err := w.SetBoundary(boundary); err != nil {
+		return 0, false, err
+	}
+	if len(r.FormData) > 0 {
+		for k, vs := range r.FormData {
+			for _, v := range vs {
+				if err := w.WriteField(k, v); err != nil {
+					return 0, false, err
+				}
+			}
+		}
+	} else if len(r.OrderedFormData) > 0 {
+		if len(r.OrderedFormData)%2 != 0 {
+			return 0, false, errBadOrderedFormData
+		}
+		for i := 0; i < len(r.OrderedFormData); i += 2 {
+			if err := w.WriteField(r.OrderedFormData[i], r.OrderedFormData[i+1]); err != nil {
+				return 0, false, err
+			}
+		}
+	}
+	for _, file := range r.uploadFiles {
+		if file.FileSize <= 0 || file.ContentType == "" {
+			return 0, false, nil
+		}
+		if _, err := w.CreatePart(createMultipartHeader(file, file.ContentType)); err != nil {
+			return 0, false, err
+		}
+		cw.n += file.FileSize
+	}
+	if err := w.Close(); err != nil {
+		return 0, false, err
+	}
+	return cw.n, true, nil
+}
+
+func multipartBody(r *Request, boundary string) io.ReadCloser {
+	pr, pw := io.Pipe()
+	go func() {
+		w := multipart.NewWriter(pw)
+		if err := w.SetBoundary(boundary); err != nil {
+			pw.CloseWithError(err)
+			return
+		}
+		pw.CloseWithError(writeMultiPart(r, w))
+	}()
+	return pr
 }
 
 func handleMultiPart(c *Client, r *Request) (err error) {
@@ -154,32 +220,28 @@ func handleMultiPart(c *Client, r *Request) (err error) {
 		b = c.multipartBoundaryFunc()
 	}
 
-	if r.forceChunkedEncoding {
-		pr, pw := io.Pipe()
-		r.GetBody = func() (io.ReadCloser, error) {
-			return pr, nil
+	w := multipart.NewWriter(io.Discard)
+	if len(b) > 0 {
+		if err = w.SetBoundary(b); err != nil {
+			return err
 		}
-		w := multipart.NewWriter(pw)
-		if len(b) > 0 {
-			w.SetBoundary(b)
-		}
-		r.SetContentType(w.FormDataContentType())
-		go func() {
-			writeMultiPart(r, w)
-			pw.Close() // close pipe writer so that pipe reader could get EOF, and stop upload
-		}()
-	} else {
-		buf := new(bytes.Buffer)
-		w := multipart.NewWriter(buf)
-		if len(b) > 0 {
-			w.SetBoundary(b)
+	}
+	boundary := w.Boundary()
+	r.SetContentType(w.FormDataContentType())
+	r.Body = nil
+	r.GetBody = func() (io.ReadCloser, error) {
+		return multipartBody(r, boundary), nil
+	}
+	r.contentLength = -1
+	if !r.forceChunkedEncoding {
+		var known bool
+		r.contentLength, known, err = multipartContentLength(r, boundary)
+		if err != nil {
+			return err
 		}
-		writeMultiPart(r, w)
-		r.GetBody = func() (io.ReadCloser, error) {
-			return io.NopCloser(bytes.NewReader(buf.Bytes())), nil
+		if !known {
+			r.contentLength = -1
 		}
-		r.Body = buf.Bytes()
-		r.SetContentType(w.FormDataContentType())
 	}
 	return
 }

request.go

@@ -39,6 +39,7 @@ type Request struct {
 	Method          string
 	Body            []byte
 	GetBody         GetContentFunc
+	contentLength   int64
 	// URL is an auto-generated field, and is nil in request middleware (OnBeforeRequest),
 	// consider using RawURL if you want, it's not nil in client middleware (WrapRoundTripFunc)
 	URL *urlpkg.URL
@@ -262,14 +263,16 @@ func (r *Request) SetQueryParamsFromStruct(v any) *Request {
 
 // SetFileReader set up a multipart form with a reader to upload file.
 func (r *Request) SetFileReader(paramName, filename string, reader io.Reader) *Request {
+	if rc, ok := reader.(io.ReadCloser); ok {
+		r.unReplayableBody = rc
+	} else {
+		r.unReplayableBody = io.NopCloser(reader)
+	}
 	r.SetFileUpload(FileUpload{
 		ParamName: paramName,
 		FileName:  filename,
 		GetFileContent: func() (io.ReadCloser, error) {
-			if rc, ok := reader.(io.ReadCloser); ok {
-				return rc, nil
-			}
-			return io.NopCloser(reader), nil
+			return r.unReplayableBody, nil
 		},
 	})
 	return r
@@ -278,8 +281,10 @@ func (r *Request) SetFileReader(paramName, filename string, reader io.Reader) *R
 // SetFileBytes set up a multipart form with given []byte to upload.
 func (r *Request) SetFileBytes(paramName, filename string, content []byte) *Request {
 	r.SetFileUpload(FileUpload{
-		ParamName: paramName,
-		FileName:  filename,
+		ParamName:   paramName,
+		FileName:    filename,
+		FileSize:    int64(len(content)),
+		ContentType: http.DetectContentType(content),
 		GetFileContent: func() (io.ReadCloser, error) {
 			return io.NopCloser(bytes.NewReader(content)), nil
 		},
@@ -307,24 +312,28 @@ func (r *Request) SetFile(paramName, filePath string) *Request {
 	}
 	fileInfo, err := os.Stat(filePath)
 	if err != nil {
+		file.Close()
 		r.client.log.Errorf("failed to stat file %s: %v", filePath, err)
 		r.appendError(err)
 		return r
 	}
+	cbuf := make([]byte, 512)
+	n, readErr := file.Read(cbuf)
+	file.Close()
+	if readErr != nil && readErr != io.EOF {
+		r.client.log.Errorf("failed to read %s: %v", filePath, readErr)
+		r.appendError(readErr)
+		return r
+	}
 	r.isMultiPart = true
 	return r.SetFileUpload(FileUpload{
 		ParamName: paramName,
 		FileName:  filepath.Base(filePath),
 		GetFileContent: func() (io.ReadCloser, error) {
-			if r.RetryAttempt > 0 {
-				file, err = os.Open(filePath)
-				if err != nil {
-					return nil, err
-				}
-			}
-			return file, nil
+			return os.Open(filePath)
 		},
-		FileSize: fileInfo.Size(),
+		FileSize:    fileInfo.Size(),
+		ContentType: http.DetectContentType(cbuf[:n]),
 	})
 }
 

request_test.go

@@ -923,6 +923,76 @@ func TestUploadMultipart(t *testing.T) {
 	tests.AssertContains(t, resp.String(), "value2", true)
 }
 
+func TestMultipartFileIsStreamedWithContentLength(t *testing.T) {
+	c := tc()
+	r := c.R().
+		SetFile("file", tests.GetTestFilePath("sample-image.png")).
+		SetFormData(map[string]string{"param": "value"})
+
+	err := handleMultiPart(c, r)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if r.Body != nil {
+		t.Fatal("multipart body was buffered in memory")
+	}
+	if r.contentLength <= r.uploadFiles[0].FileSize {
+		t.Fatalf("invalid multipart content length: %d", r.contentLength)
+	}
+
+	for attempt := 0; attempt < 2; attempt++ {
+		r.RetryAttempt = attempt
+		body, err := r.GetBody()
+		if err != nil {
+			t.Fatal(err)
+		}
+		n, err := io.Copy(io.Discard, body)
+		body.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+		if n != r.contentLength {
+			t.Fatalf("body length is %d, want %d", n, r.contentLength)
+		}
+	}
+}
+
+func TestMultipartReaderWithUnknownSizeIsStreamed(t *testing.T) {
+	c := tc()
+	r := c.R().SetFileReader("file", "file.txt", strings.NewReader("test"))
+
+	err := handleMultiPart(c, r)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if r.Body != nil {
+		t.Fatal("multipart body was buffered in memory")
+	}
+	if r.contentLength != -1 {
+		t.Fatalf("content length is %d, want -1", r.contentLength)
+	}
+	body, err := r.GetBody()
+	if err != nil {
+		t.Fatal(err)
+	}
+	n, err := io.Copy(io.Discard, body)
+	body.Close()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n <= int64(len("test")) {
+		t.Fatalf("multipart body is unexpectedly short: %d", n)
+	}
+}
+
+func TestMultipartReaderCannotBeRetried(t *testing.T) {
+	_, err := tc().R().
+		SetRetryCount(1).
+		SetFileReader("file", "file.txt", strings.NewReader("test")).
+		Post("/file-text")
+	tests.AssertEqual(t, errRetryableWithUnReplayableBody, err)
+}
+
 func TestFixPragmaCache(t *testing.T) {
 	resp, err := tc().EnableForceHTTP1().R().Get("/pragma")
 	assertSuccess(t, resp, err)