Merge pull request #159 from tooolbox/153-auth-refactor

Refactored buildHTTP to fix Auth with GetBody.

Author: casualjim

client/request.go

@@ -18,7 +18,6 @@ import (
 	"bytes"
 	"fmt"
 	"io"
-	"io/ioutil"
 	"log"
 	"mime/multipart"
 	"net/http"
@@ -42,6 +41,7 @@ func newRequest(method, pathPattern string, writer runtime.ClientRequestWriter)
 		header:      make(http.Header),
 		query:       make(url.Values),
 		timeout:     DefaultTimeout,
+		getBody:     getRequestBuffer,
 	}, nil
 }
 
@@ -67,6 +67,8 @@ type request struct {
 	payload    interface{}
 	timeout    time.Duration
 	buf        *bytes.Buffer
+
+	getBody func(r *request) []byte
 }
 
 var (
@@ -93,60 +95,34 @@ func (r *request) buildHTTP(mediaType, basePath string, producers map[string]run
 		return nil, err
 	}
 
-	if auth != nil {
-		if err := auth.AuthenticateRequest(r, registry); err != nil {
-			return nil, err
-		}
-	}
-
-	// create http request
-	var reinstateSlash bool
-	if r.pathPattern != "" && r.pathPattern != "/" && r.pathPattern[len(r.pathPattern)-1] == '/' {
-		reinstateSlash = true
-	}
-	urlPath := path.Join(basePath, r.pathPattern)
-	for k, v := range r.pathParams {
-		urlPath = strings.Replace(urlPath, "{"+k+"}", url.PathEscape(v), -1)
-	}
-	if reinstateSlash {
-		urlPath = urlPath + "/"
-	}
-
-	var body io.ReadCloser
+	// Our body must be an io.Reader.
+	// When we create the http.Request, if we pass it a
+	// bytes.Buffer then it will wrap it in an io.ReadCloser
+	// and set the content length automatically.
+	var body io.Reader
 	var pr *io.PipeReader
 	var pw *io.PipeWriter
 
 	r.buf = bytes.NewBuffer(nil)
 	if r.payload != nil || len(r.formFields) > 0 || len(r.fileFields) > 0 {
-		body = ioutil.NopCloser(r.buf)
+		body = r.buf
 		if r.isMultipart(mediaType) {
 			pr, pw = io.Pipe()
 			body = pr
 		}
 	}
-	req, err := http.NewRequest(r.method, urlPath, body)
-
-	if err != nil {
-		return nil, err
-	}
-
-	req.URL.RawQuery = r.query.Encode()
-	req.Header = r.header
 
 	// check if this is a form type request
 	if len(r.formFields) > 0 || len(r.fileFields) > 0 {
 		if !r.isMultipart(mediaType) {
-			req.Header.Set(runtime.HeaderContentType, mediaType)
+			r.header.Set(runtime.HeaderContentType, mediaType)
 			formString := r.formFields.Encode()
-			// set content length before writing to the buffer
-			req.ContentLength = int64(len(formString))
-			// write the form values as the body
 			r.buf.WriteString(formString)
-			return req, nil
+			goto DoneChoosingBodySource
 		}
 
 		mp := multipart.NewWriter(pw)
-		req.Header.Set(runtime.HeaderContentType, mangleContentType(mediaType, mp.Boundary()))
+		r.header.Set(runtime.HeaderContentType, mangleContentType(mediaType, mp.Boundary()))
 
 		go func() {
 			defer func() {
@@ -184,65 +160,121 @@ func (r *request) buildHTTP(mediaType, basePath string, producers map[string]run
 			}
 
 		}()
-		return req, nil
 
+		goto DoneChoosingBodySource
 	}
 
 	// if there is payload, use the producer to write the payload, and then
 	// set the header to the content-type appropriate for the payload produced
 	if r.payload != nil {
 		// TODO: infer most appropriate content type based on the producer used,
 		// and the `consumers` section of the spec/operation
-		req.Header.Set(runtime.HeaderContentType, mediaType)
+		r.header.Set(runtime.HeaderContentType, mediaType)
 		if rdr, ok := r.payload.(io.ReadCloser); ok {
-			req.Body = rdr
-
-			return req, nil
+			body = rdr
+			goto DoneChoosingBodySource
 		}
 
 		if rdr, ok := r.payload.(io.Reader); ok {
-			req.Body = ioutil.NopCloser(rdr)
+			body = rdr
+			goto DoneChoosingBodySource
+		}
 
-			return req, nil
+		producer := producers[mediaType]
+		if err := producer.Produce(r.buf, r.payload); err != nil {
+			return nil, err
 		}
+	}
 
-		req.GetBody = func() (io.ReadCloser, error) {
-			var b bytes.Buffer
-			producer := producers[mediaType]
-			if err := producer.Produce(&b, r.payload); err != nil {
-				return nil, err
-			}
+DoneChoosingBodySource:
 
-			return ioutil.NopCloser(&b), nil
-		}
+	if runtime.CanHaveBody(r.method) && body == nil && r.header.Get(runtime.HeaderContentType) == "" {
+		r.header.Set(runtime.HeaderContentType, mediaType)
+	}
 
-		// set the content length of the request or else a chunked transfer is
-		// declared, and this corrupts outgoing JSON payloads. the content's
-		// length must be set prior to the body being written per the spec at
-		// https://golang.org/pkg/net/http
+	if auth != nil {
+
+		// If we're not using r.buf as our http.Request's body,
+		// either the payload is an io.Reader or io.ReadCloser,
+		// or we're doing a multipart form/file.
 		//
-		//     If Body is present, Content-Length is <= 0 and TransferEncoding
-		//     hasn't been set to "identity", Write adds
-		//     "Transfer-Encoding: chunked" to the header. Body is closed
-		//     after it is sent.
+		// In those cases, if the AuthenticateRequest call asks for the body,
+		// we must read it into a buffer and provide that, then use that buffer
+		// as the body of our http.Request.
 		//
-		// to that end a temporary buffer, b, is created to produce the payload
-		// body, and then its size is used to set the request's content length
-		var b bytes.Buffer
-		producer := producers[mediaType]
-		if err := producer.Produce(&b, r.payload); err != nil {
-			return nil, err
+		// This is done in-line with the GetBody() request rather than ahead
+		// of time, because there's no way to know if the AuthenticateRequest
+		// will even ask for the body of the request.
+		//
+		// If for some reason the copy fails, there's no way to return that
+		// error to the GetBody() call, so return it afterwards.
+		//
+		// An error from the copy action is prioritized over any error
+		// from the AuthenticateRequest call, because the mis-read
+		// body may have interfered with the auth.
+		//
+		var copyErr error
+		if buf, ok := body.(*bytes.Buffer); body != nil && (!ok || buf != r.buf) {
+
+			var copied bool
+			r.getBody = func(r *request) []byte {
+
+				if copied {
+					return getRequestBuffer(r)
+				}
+
+				defer func() {
+					copied = true
+				}()
+
+				if _, copyErr = io.Copy(r.buf, body); copyErr != nil {
+					return nil
+				}
+
+				if closer, ok := body.(io.ReadCloser); ok {
+					if copyErr = closer.Close(); copyErr != nil {
+						return nil
+					}
+				}
+
+				body = r.buf
+				return getRequestBuffer(r)
+			}
 		}
-		req.ContentLength = int64(b.Len())
-		if _, err := r.buf.Write(b.Bytes()); err != nil {
-			return nil, err
+
+		authErr := auth.AuthenticateRequest(r, registry)
+
+		if copyErr != nil {
+			return nil, fmt.Errorf("error retrieving the response body: %v", copyErr)
+		}
+
+		if authErr != nil {
+			return nil, authErr
 		}
+
+	}
+
+	// create http request
+	var reinstateSlash bool
+	if r.pathPattern != "" && r.pathPattern != "/" && r.pathPattern[len(r.pathPattern)-1] == '/' {
+		reinstateSlash = true
+	}
+	urlPath := path.Join(basePath, r.pathPattern)
+	for k, v := range r.pathParams {
+		urlPath = strings.Replace(urlPath, "{"+k+"}", url.PathEscape(v), -1)
+	}
+	if reinstateSlash {
+		urlPath = urlPath + "/"
 	}
 
-	if runtime.CanHaveBody(req.Method) && req.Body == nil && req.Header.Get(runtime.HeaderContentType) == "" {
-		req.Header.Set(runtime.HeaderContentType, mediaType)
+	req, err := http.NewRequest(r.method, urlPath, body)
+	if err != nil {
+		return nil, err
 	}
 
+	req.URL.RawQuery = r.query.Encode()
+	req.Header = r.header
+
 	return req, nil
 }
 
@@ -266,6 +298,10 @@ func (r *request) GetPath() string {
 }
 
 func (r *request) GetBody() []byte {
+	return r.getBody(r)
+}
+
+func getRequestBuffer(r *request) []byte {
 	if r.buf == nil {
 		return nil
 	}