fix: improve flow and error handling in the multipart request #1030 (#1101)

Author: jeevatkm

client.go

@@ -710,6 +710,7 @@ func (c *Client) R() *Request {
 		AllowNonIdempotentRetry:    c.allowNonIdempotentRetry,
 		HeaderAuthorizationKey:     c.headerAuthorizationKey,
 
+		mu:                  new(sync.Mutex),
 		client:              c,
 		baseURL:             c.baseURL,
 		multipartFields:     make([]*MultipartField, 0),
@@ -2372,36 +2373,42 @@ func (c *Client) execute(req *Request) (*Response, error) {
 
 	req.Time = time.Now()
 	resp, err := c.Client().Do(req.withTimeout())
+	// Cancel multipart context for io.Copy to stop reading/writing further
+	if req.isMultiPart && req.multipartCancelFunc != nil {
+		req.multipartCancelFunc()
+	}
 
 	response := &Response{Request: req, RawResponse: resp}
 	response.setReceivedAt()
 	if err != nil {
 		return response, err
 	}
-	if req.multipartErrChan != nil {
-		if err = <-req.multipartErrChan; err != nil {
-			return response, err
+	if req.isMultiPart && req.multipartErrChan != nil {
+		// read all multipart errors from channel
+		for err = range req.multipartErrChan {
+			response.CascadeError = wrapErrors(err, response.CascadeError)
 		}
 	}
+
 	if resp != nil {
 		if c.circuitBreaker != nil {
 			c.circuitBreaker.applyPolicies(resp)
 		}
 
 		response.Body = resp.Body
 		if err = response.wrapContentDecompresser(); err != nil {
-			return response, err
+			return response, response.wrapError(err, false)
 		}
 
 		response.wrapLimitReadCloser()
-	}
 
-	if !req.DoNotParseResponse {
-		if req.ResponseBodyUnlimitedReads || req.Debug {
-			response.wrapCopyReadCloser()
+		if !req.DoNotParseResponse {
+			if req.ResponseBodyUnlimitedReads || req.Debug {
+				response.wrapCopyReadCloser()
 
-			if err = response.readAll(); err != nil {
-				return response, err
+				if err = response.readAll(); err != nil {
+					return response, response.wrapError(err, false)
+				}
 			}
 		}
 	}
@@ -2415,8 +2422,7 @@ func (c *Client) execute(req *Request) (*Response, error) {
 		}
 	}
 
-	err = response.CascadeError
-	return response, err
+	return response, response.wrapError(nil, false)
 }
 
 // getting TLS client config if not exists then create one

middleware.go

@@ -7,6 +7,7 @@ package resty
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"mime"
@@ -238,7 +239,7 @@ func createRawRequest(c *Client, r *Request) (err error) {
 	}
 
 	// get the context reference back from underlying RawRequest
-	r.ctx = r.RawRequest.Context()
+	r.SetContext(r.RawRequest.Context())
 
 	// Assign close connection option
 	r.RawRequest.Close = r.CloseConnection
@@ -289,105 +290,138 @@ func addCredentials(c *Client, r *Request) error {
 	return nil
 }
 
-func handleMultipart(c *Client, r *Request) error {
-	for k, v := range c.FormData() {
-		if _, ok := r.FormData[k]; ok {
-			continue
-		}
-		r.FormData[k] = v[:]
-	}
-
-	mfLen := len(r.multipartFields)
-	if mfLen == 0 {
-		r.bodyBuf = acquireBuffer()
-		mw := multipart.NewWriter(r.bodyBuf)
+var multipartWriteField = func(w *multipart.Writer, name, value string) error {
+	return w.WriteField(name, value)
+}
 
-		// set boundary if it is provided by the user
-		if !isStringEmpty(r.multipartBoundary) {
-			if err := mw.SetBoundary(r.multipartBoundary); err != nil {
+var multipartWriteFormData = func(w *multipart.Writer, r *Request) error {
+	for k, v := range r.FormData {
+		for _, iv := range v {
+			if err := multipartWriteField(w, k, iv); err != nil {
 				return err
 			}
 		}
+	}
+	return nil
+}
 
-		if err := r.writeFormData(mw); err != nil {
-			return err
-		}
-
-		r.Header.Set(hdrContentTypeKey, mw.FormDataContentType())
-		closeq(mw)
+var multipartCreatePart = func(w *multipart.Writer, h textproto.MIMEHeader) (io.Writer, error) {
+	return w.CreatePart(h)
+}
 
+var multipartSetBoundary = func(w *multipart.Writer, r *Request) error {
+	if isStringEmpty(r.multipartBoundary) {
 		return nil
 	}
+	return w.SetBoundary(r.multipartBoundary)
+}
 
-	// multipart streaming
-	bodyReader, bodyWriter := io.Pipe()
-	mw := multipart.NewWriter(bodyWriter)
-	r.Body = bodyReader
-	r.multipartErrChan = make(chan error, 1)
+func handleMultipartFormData(r *Request) error {
+	r.bodyBuf = acquireBuffer()
+	mw := multipart.NewWriter(r.bodyBuf)
+	defer mw.Close()
 
-	// set boundary if it is provided by the user
-	if !isStringEmpty(r.multipartBoundary) {
-		if err := mw.SetBoundary(r.multipartBoundary); err != nil {
-			return err
-		}
+	// set custom multipart boundary if exists
+	if err := multipartSetBoundary(mw, r); err != nil {
+		return err
 	}
 
-	go func() {
-		defer close(r.multipartErrChan)
-		if err := createMultipart(mw, r); err != nil {
-			r.multipartErrChan <- err
-		}
-		closeq(mw)
-		closeq(bodyWriter)
-	}()
-
 	r.Header.Set(hdrContentTypeKey, mw.FormDataContentType())
-	return nil
-}
 
-var mpCreatePart = func(w *multipart.Writer, h textproto.MIMEHeader) (io.Writer, error) {
-	return w.CreatePart(h)
+	return multipartWriteFormData(mw, r)
 }
 
-func createMultipart(w *multipart.Writer, r *Request) error {
-	if err := r.writeFormData(w); err != nil {
-		return err
+func handleMultipart(c *Client, r *Request) error {
+	for k, v := range c.FormData() {
+		if _, ok := r.FormData[k]; ok {
+			continue
+		}
+		r.FormData[k] = v[:]
 	}
 
+	if len(r.multipartFields) == 0 {
+		return handleMultipartFormData(r)
+	}
+
+	// pre-process multipart fields to catch possible errors
 	for _, mf := range r.multipartFields {
-		if len(mf.Values) > 0 {
-			for _, v := range mf.Values {
-				w.WriteField(mf.Name, v)
-			}
+		if mf.isValues() {
 			continue
 		}
 
-		if err := mf.openFileIfRequired(); err != nil {
+		if err := mf.openFile(); err != nil {
 			return err
 		}
 
-		p := make([]byte, 512)
-		size, err := mf.Reader.Read(p)
-		if err != nil && err != io.EOF {
+		if err := mf.detectContentType(); err != nil {
 			return err
 		}
-		// auto detect content type if empty
-		if isStringEmpty(mf.ContentType) {
-			mf.ContentType = http.DetectContentType(p[:size])
-		}
+	}
 
-		partWriter, err := mpCreatePart(w, mf.createHeader())
-		if err != nil {
-			return err
+	// multipart streaming
+	br, bw := io.Pipe()
+	mw := multipart.NewWriter(bw)
+	r.Body = br
+
+	// set custom multipart boundary if exists
+	if err := multipartSetBoundary(mw, r); err != nil {
+		closeq(bw)
+		return err
+	}
+
+	r.Header.Set(hdrContentTypeKey, mw.FormDataContentType())
+
+	r.multipartErrChan = make(chan error, 1)
+	go func() {
+		defer close(r.multipartErrChan)
+		defer func() {
+			if err := mw.Close(); err != nil {
+				r.multipartErrChan <- err
+			}
+			if err := bw.Close(); err != nil {
+				r.multipartErrChan <- err
+			}
+		}()
+
+		if err := multipartWriteFormData(mw, r); err != nil {
+			r.multipartErrChan <- err
+			return
 		}
 
-		partWriter = mf.wrapProgressCallbackIfPresent(partWriter)
-		partWriter.Write(p[:size])
+		ctx, cancel := context.WithCancel(r.Context())
+		r.multipartCancelFunc = cancel
+		for _, mf := range r.multipartFields {
+			if mf.isValues() {
+				for _, v := range mf.Values {
+					if err := multipartWriteField(mw, mf.Name, v); err != nil {
+						r.multipartErrChan <- err
+						return
+					}
+				}
+				continue
+			}
 
-		if _, err = ioCopy(partWriter, mf.Reader); err != nil {
-			return err
+			partWriter, err := multipartCreatePart(mw, mf.createHeader())
+			if err != nil {
+				r.multipartErrChan <- err
+				return
+			}
+
+			partWriter = mf.wrapProgressCallbackIfPresent(partWriter)
+			if len(mf.tempBuf) > 0 {
+				if _, err = partWriter.Write(mf.tempBuf); err != nil {
+					r.multipartErrChan <- err
+					return
+				}
+			}
+
+			reader := &gracefulStopReader{ctx: ctx, r: mf.Reader}
+			if _, err = ioCopy(partWriter, reader); err != nil {
+				r.multipartErrChan <- err
+				return
+			}
 		}
-	}
+	}()
 
 	return nil
 }
@@ -482,7 +516,8 @@ func handleRequestBody(c *Client, r *Request) error {
 // based on registered HTTP response `Content-Type` decoder, see [Client.AddContentTypeDecoder];
 // if [Request.SetResult], [Request.SetResultError], or [Client.SetResultError] is used
 func AutoParseResponseMiddleware(c *Client, res *Response) (err error) {
-	if res.CascadeError != nil || res.Request.DoNotParseResponse {
+	if (res.CascadeError != nil && (res.Request.isMultiPart && res.StatusCode() == 0)) ||
+		res.Request.DoNotParseResponse {
 		return // move on
 	}