rest: fix retries reading the request body after the request has finished
Before this change the transport could still be reading the request body when Call returned, in two ways: - the goroutine writing a multipart upload form was never stopped or waited for, so after a failed request it kept reading the file body while the caller retried - the transport itself is documented to close the request body, possibly in a different goroutine after the request has finished, and hiding the body's Close method removed any way to wait for that Both of these raced with callers which seek to the start of a pooled buffer and retry the request, corrupting the retried body, and the abandoned multipart goroutine could read pooled pages after they had been returned to the pool. Call now waits for the transport to finish with the request body before returning, and the multipart form writer is stopped and waited for once the request has finished.
This commit is contained in:
+51
-3
@@ -322,7 +322,16 @@ func (api *Client) Call(ctx context.Context, opts *Opts) (resp *http.Response, e
|
||||
if len(opts.Parameters) > 0 {
|
||||
url += "?" + opts.Parameters.Encode()
|
||||
}
|
||||
body := readers.NoCloser(opts.Body)
|
||||
var bodyClosed chan struct{}
|
||||
body := opts.Body
|
||||
if _, hasClose := body.(io.Closer); hasClose {
|
||||
// Hide the body's Close method from the transport so it can't
|
||||
// close the caller's body, arranging for the transport's Close
|
||||
// to signal that it has finished with the body instead.
|
||||
bodyClosed = make(chan struct{})
|
||||
closed := bodyClosed
|
||||
body = readers.NoCloserNotify(body, func() { close(closed) })
|
||||
}
|
||||
// If length is set and zero then nil out the body to stop use
|
||||
// use of chunked encoding and insert a "Content-Length: 0"
|
||||
// header.
|
||||
@@ -331,6 +340,7 @@ func (api *Client) Call(ctx context.Context, opts *Opts) (resp *http.Response, e
|
||||
// files except 0 length files.
|
||||
if opts.ContentLength != nil && *opts.ContentLength == 0 {
|
||||
body = nil
|
||||
bodyClosed = nil
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, opts.Method, url, body)
|
||||
if err != nil {
|
||||
@@ -402,6 +412,14 @@ func (api *Client) Call(ctx context.Context, opts *Opts) (resp *http.Response, e
|
||||
}
|
||||
api.mu.RUnlock()
|
||||
resp, err = c.Do(req)
|
||||
if bodyClosed != nil {
|
||||
// Wait for the transport to close the request body before
|
||||
// returning. It always closes it, but is documented to
|
||||
// possibly do so in a different goroutine even after Do has
|
||||
// returned, and the caller may seek or close the body for a
|
||||
// retry as soon as we return.
|
||||
<-bodyClosed
|
||||
}
|
||||
api.mu.RLock()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -433,6 +451,25 @@ func CreateFormFile(w *multipart.Writer, fieldname, filename, contentType string
|
||||
return w.CreatePart(h)
|
||||
}
|
||||
|
||||
// multipartBody is the request body of a multipart upload as returned
|
||||
// by MultipartUpload.
|
||||
//
|
||||
// Close stops the goroutine writing the form and waits for it to
|
||||
// finish, so afterwards the file reader passed to MultipartUpload is
|
||||
// no longer being read from. It is safe to call multiple times.
|
||||
type multipartBody struct {
|
||||
*io.PipeReader
|
||||
done <-chan struct{}
|
||||
}
|
||||
|
||||
// Close aborts reading the form and waits for the goroutine writing
|
||||
// it to stop.
|
||||
func (b *multipartBody) Close() error {
|
||||
err := b.PipeReader.Close()
|
||||
<-b.done
|
||||
return err
|
||||
}
|
||||
|
||||
// MultipartUpload creates an io.Reader which produces an encoded a
|
||||
// multipart form upload from the params passed in and the passed in
|
||||
//
|
||||
@@ -443,6 +480,11 @@ func CreateFormFile(w *multipart.Writer, fieldname, filename, contentType string
|
||||
//
|
||||
// the int64 returned is the overhead in addition to the file contents, in case Content-Length is required
|
||||
//
|
||||
// Closing the returned reader stops the goroutine writing the form
|
||||
// and waits for it to finish, after which in is no longer being read
|
||||
// from. Callers which retry with a seekable in should close it after
|
||||
// each attempt.
|
||||
//
|
||||
// NB This doesn't allow setting the content type of the attachment
|
||||
func MultipartUpload(ctx context.Context, in io.Reader, params url.Values, contentName, fileName string, contentType string) (io.ReadCloser, string, int64, error) {
|
||||
bodyReader, bodyWriter := io.Pipe()
|
||||
@@ -529,7 +571,7 @@ func MultipartUpload(ctx context.Context, in io.Reader, params url.Values, conte
|
||||
_ = bodyWriter.Close()
|
||||
}()
|
||||
|
||||
return bodyReader, formContentType, multipartLength, nil
|
||||
return &multipartBody{PipeReader: bodyReader, done: quit}, formContentType, multipartLength, nil
|
||||
}
|
||||
|
||||
// CallJSON runs Call and decodes the body as a JSON object into response (if not nil)
|
||||
@@ -600,14 +642,20 @@ func (api *Client) callCodec(ctx context.Context, opts *Opts, request any, respo
|
||||
}
|
||||
opts = opts.Copy()
|
||||
|
||||
var mpBody io.ReadCloser
|
||||
var overhead int64
|
||||
opts.Body, opts.ContentType, overhead, err = MultipartUpload(ctx, opts.Body, params, opts.MultipartContentName, opts.MultipartFileName, opts.MultipartContentType)
|
||||
mpBody, opts.ContentType, overhead, err = MultipartUpload(ctx, opts.Body, params, opts.MultipartContentName, opts.MultipartFileName, opts.MultipartContentType)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
opts.Body = mpBody
|
||||
if opts.ContentLength != nil {
|
||||
*opts.ContentLength += overhead
|
||||
}
|
||||
// Stop the multipart writer reading the body when the call has
|
||||
// finished so the caller can seek or close the body (e.g. to
|
||||
// retry) as soon as we return.
|
||||
defer func() { _ = mpBody.Close() }()
|
||||
}
|
||||
resp, err = api.Call(ctx, opts)
|
||||
if err != nil {
|
||||
|
||||
Reference in New Issue
Block a user