rest: make auth preserving redirects an option

lib/rest/rest.go

@@ -150,6 +150,7 @@ type Opts struct {
 	Trailer               *http.Header // set the request trailer
 	Close                 bool         // set to close the connection after this transaction
 	NoRedirect            bool         // if this is set then the client won't follow redirects
+	AuthRedirect          bool         // if this is set then the client will redirect with Auth
 }
 
 // Copy creates a copy of the options
@@ -210,6 +211,34 @@ func ClientWithNoRedirects(c *http.Client) *http.Client {
 	return &clientCopy
 }
 
+// ClientWithAuthRedirects makes a new http client which will re-apply Auth on redirects
+func ClientWithAuthRedirects(c *http.Client) *http.Client {
+	clientCopy := *c
+	clientCopy.CheckRedirect = func(req *http.Request, via []*http.Request) error {
+		if len(via) >= 10 {
+			return errors.New("stopped after 10 redirects")
+		} else if len(via) == 0 {
+			return nil
+		}
+		prevReq := via[len(via)-1]
+		resp := req.Response
+		if resp == nil {
+			return nil
+		}
+		// Look at previous response to see if it was a redirect and preserve auth if so
+		switch resp.StatusCode {
+		case http.StatusMovedPermanently, http.StatusFound, http.StatusSeeOther, http.StatusTemporaryRedirect, http.StatusPermanentRedirect:
+			// Reapply Auth (if any) from previous request on redirect
+			auth := prevReq.Header.Get("Authorization")
+			if auth != "" {
+				req.Header.Add("Authorization", auth)
+			}
+		}
+		return nil
+	}
+	return &clientCopy
+}
+
 // Call makes the call and returns the http.Response
 //
 // if err == nil then resp.Body will need to be closed unless
@@ -302,6 +331,8 @@ func (api *Client) Call(ctx context.Context, opts *Opts) (resp *http.Response, e
 	var c *http.Client
 	if opts.NoRedirect {
 		c = ClientWithNoRedirects(api.c)
+	} else if opts.AuthRedirect {
+		c = ClientWithAuthRedirects(api.c)
 	} else {
 		c = api.c
 	}