From a8549a2c408f4c72354b24aaa8f9bcafdd35982d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan-Otto=20Kr=C3=B6pke?= <mail@jkroepke.de>
Date: Sun, 8 Jul 2018 15:42:04 +0200
Subject: [PATCH] when uploading layers reuse token so reader isn't closed
 premature. (#103)

---
 registry/basictransport.go |  2 +-
 registry/layer.go          | 13 ++++++++-----
 registry/tokentransport.go |  6 +++++-
 3 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/registry/basictransport.go b/registry/basictransport.go
index be169953..4ba6ad79 100644
--- a/registry/basictransport.go
+++ b/registry/basictransport.go
@@ -15,7 +15,7 @@ type BasicTransport struct {
 
 // RoundTrip defines the round tripper for basic auth transport.
 func (t *BasicTransport) RoundTrip(req *http.Request) (*http.Response, error) {
-	if strings.HasPrefix(req.URL.String(), t.URL) {
+	if strings.HasPrefix(req.URL.String(), t.URL) && req.Header.Get("Authorization") == "" {
 		if t.Username != "" || t.Password != "" {
 			req.SetBasicAuth(t.Username, t.Password)
 		}
diff --git a/registry/layer.go b/registry/layer.go
index cd5f1f1e..e382337d 100644
--- a/registry/layer.go
+++ b/registry/layer.go
@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"net/url"
 
+	"fmt"
 	"github.com/docker/distribution/reference"
 	"github.com/opencontainers/go-digest"
 )
@@ -24,7 +25,7 @@ func (r *Registry) DownloadLayer(repository string, digest digest.Digest) (io.Re
 
 // UploadLayer uploads a specific layer by digest for a repository.
 func (r *Registry) UploadLayer(repository string, digest reference.Reference, content io.Reader) error {
-	uploadURL, err := r.initiateUpload(repository)
+	uploadURL, token, err := r.initiateUpload(repository)
 	if err != nil {
 		return err
 	}
@@ -39,6 +40,7 @@ func (r *Registry) UploadLayer(repository string, digest reference.Reference, co
 		return err
 	}
 	upload.Header.Set("Content-Type", "application/octet-stream")
+	upload.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
 
 	_, err = r.Client.Do(upload)
 	return err
@@ -70,20 +72,21 @@ func (r *Registry) HasLayer(repository string, digest digest.Digest) (bool, erro
 	return false, err
 }
 
-func (r *Registry) initiateUpload(repository string) (*url.URL, error) {
+func (r *Registry) initiateUpload(repository string) (*url.URL, string, error) {
 	initiateURL := r.url("/v2/%s/blobs/uploads/", repository)
 	r.Logf("registry.layer.initiate-upload url=%s repository=%s", initiateURL, repository)
 
 	resp, err := r.Client.Post(initiateURL, "application/octet-stream", nil)
 	if err != nil {
-		return nil, err
+		return nil, "", err
 	}
+	token := resp.Header.Get("Request-Token")
 	defer resp.Body.Close()
 
 	location := resp.Header.Get("Location")
 	locationURL, err := url.Parse(location)
 	if err != nil {
-		return nil, err
+		return nil, token, err
 	}
-	return locationURL, nil
+	return locationURL, token, nil
 }
diff --git a/registry/tokentransport.go b/registry/tokentransport.go
index a21ccd75..d1cb8712 100644
--- a/registry/tokentransport.go
+++ b/registry/tokentransport.go
@@ -47,7 +47,11 @@ func (t *TokenTransport) authAndRetry(authService *authService, req *http.Reques
 		return authResp, err
 	}
 
-	return t.retry(req, token)
+	response, err := t.retry(req, token)
+	if response != nil {
+		response.Header.Set("request-token", token)
+	}
+	return response, err
 }
 
 func (t *TokenTransport) auth(authService *authService) (string, *http.Response, error) {