From 6ad54e3df6085f905d742168f676620f80b422ef Mon Sep 17 00:00:00 2001
From: Michael Crosby <crosbymichael@gmail.com>
Date: Wed, 10 Dec 2014 18:08:40 -0800
Subject: [PATCH] Refactor put image function's redirect loop

Signed-off-by: Michael Crosby <crosbymichael@gmail.com>
---
 docs/session.go | 64 ++++++++++++++++++++++++++-----------------------
 1 file changed, 34 insertions(+), 30 deletions(-)

diff --git a/docs/session.go b/docs/session.go
index 2658ec1a8..4b2f55225 100644
--- a/docs/session.go
+++ b/docs/session.go
@@ -462,7 +462,6 @@ func (r *Session) PushRegistryTag(remote, revision, tag, registry string, token
 
 func (r *Session) PushImageJSONIndex(remote string, imgList []*ImgData, validate bool, regs []string) (*RepositoryData, error) {
 	cleanImgList := []*ImgData{}
-
 	if validate {
 		for _, elem := range imgList {
 			if elem.Checksum != "" {
@@ -484,44 +483,28 @@ func (r *Session) PushImageJSONIndex(remote string, imgList []*ImgData, validate
 	u := fmt.Sprintf("%srepositories/%s/%s", r.indexEndpoint.VersionString(1), remote, suffix)
 	log.Debugf("[registry] PUT %s", u)
 	log.Debugf("Image list pushed to index:\n%s", imgListJSON)
-	req, err := r.reqFactory.NewRequest("PUT", u, bytes.NewReader(imgListJSON))
-	if err != nil {
-		return nil, err
+	headers := map[string][]string{
+		"Content-type":   {"application/json"},
+		"X-Docker-Token": {"true"},
 	}
-	req.Header.Add("Content-type", "application/json")
-	req.SetBasicAuth(r.authConfig.Username, r.authConfig.Password)
-	req.ContentLength = int64(len(imgListJSON))
-	req.Header.Set("X-Docker-Token", "true")
 	if validate {
-		req.Header["X-Docker-Endpoints"] = regs
+		headers["X-Docker-Endpoints"] = regs
 	}
 
-	res, _, err := r.doRequest(req)
-	if err != nil {
-		return nil, err
-	}
-	defer res.Body.Close()
-
 	// Redirect if necessary
-	for res.StatusCode >= 300 && res.StatusCode < 400 {
-		log.Debugf("Redirected to %s", res.Header.Get("Location"))
-		req, err := r.reqFactory.NewRequest("PUT", res.Header.Get("Location"), bytes.NewReader(imgListJSON))
-		if err != nil {
+	var res *http.Response
+	for {
+		if res, err = r.putImageRequest(u, headers, imgListJSON); err != nil {
 			return nil, err
 		}
-		req.SetBasicAuth(r.authConfig.Username, r.authConfig.Password)
-		req.ContentLength = int64(len(imgListJSON))
-		req.Header.Set("X-Docker-Token", "true")
-		if validate {
-			req.Header["X-Docker-Endpoints"] = regs
+		if !shouldRedirect(res) {
+			break
 		}
-		redirect, _, err := r.doRequest(req)
-		if err != nil {
-			return nil, err
-		}
-		res = redirect
-		defer res.Body.Close()
+		res.Body.Close()
+		u = res.Header.Get("Location")
+		log.Debugf("Redirected to %s", u)
 	}
+	defer res.Body.Close()
 
 	var tokens, endpoints []string
 	if !validate {
@@ -564,6 +547,27 @@ func (r *Session) PushImageJSONIndex(remote string, imgList []*ImgData, validate
 	}, nil
 }
 
+func (r *Session) putImageRequest(u string, headers map[string][]string, body []byte) (*http.Response, error) {
+	req, err := r.reqFactory.NewRequest("PUT", u, bytes.NewReader(body))
+	if err != nil {
+		return nil, err
+	}
+	req.SetBasicAuth(r.authConfig.Username, r.authConfig.Password)
+	req.ContentLength = int64(len(body))
+	for k, v := range headers {
+		req.Header[k] = v
+	}
+	response, _, err := r.doRequest(req)
+	if err != nil {
+		return nil, err
+	}
+	return response, nil
+}
+
+func shouldRedirect(response *http.Response) bool {
+	return response.StatusCode >= 300 && response.StatusCode < 400
+}
+
 func (r *Session) SearchRepositories(term string) (*SearchResults, error) {
 	log.Debugf("Index server: %s", r.indexEndpoint)
 	u := r.indexEndpoint.VersionString(1) + "search?q=" + url.QueryEscape(term)