diff --git a/client.go b/client.go index 8113235..2c8acd0 100644 --- a/client.go +++ b/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 { diff --git a/middleware.go b/middleware.go index 54129d8..dde3f0e 100644 --- a/middleware.go +++ b/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 } diff --git a/request.go b/request.go index b5ac8ec..7435d1f 100644 --- a/request.go +++ b/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]), }) } diff --git a/request_test.go b/request_test.go index a1c649a..2937a5d 100644 --- a/request_test.go +++ b/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)