From d724bb1d24576d89ecde2916d7e1e25f486f1f16 Mon Sep 17 00:00:00 2001
From: yuzou <zouyu7@huawei.com>
Date: Thu, 15 Oct 2015 16:56:28 +0800
Subject: [PATCH] use the scheme and host from x-forward-proto and
 x-forward-host if they exits and correct the scheme for Location header
 during image upload

Signed-off-by: yuzou <zouyu7@huawei.com>
---
 registry/api/v2/urls.go      |  4 ++-
 registry/api/v2/urls_test.go | 49 +++++++++++++++++++++++++++++-------
 2 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/registry/api/v2/urls.go b/registry/api/v2/urls.go
index 429743940..6ba39cc9b 100644
--- a/registry/api/v2/urls.go
+++ b/registry/api/v2/urls.go
@@ -204,7 +204,9 @@ func (cr clonedRoute) URL(pairs ...string) (*url.URL, error) {
 		routeURL.Path = routeURL.Path[1:]
 	}
 
-	return cr.root.ResolveReference(routeURL), nil
+	url := cr.root.ResolveReference(routeURL)
+	url.Scheme = cr.root.Scheme
+	return url, nil
 }
 
 // appendValuesURL appends the parameters to the url.
diff --git a/registry/api/v2/urls_test.go b/registry/api/v2/urls_test.go
index 16e05695a..0ad33add8 100644
--- a/registry/api/v2/urls_test.go
+++ b/registry/api/v2/urls_test.go
@@ -166,6 +166,11 @@ func TestBuilderFromRequest(t *testing.T) {
 			request: &http.Request{URL: u, Host: u.Host},
 			base:    "http://example.com",
 		},
+
+		{
+			request: &http.Request{URL: u, Host: u.Host, Header: forwardedProtoHeader},
+			base:    "http://example.com",
+		},
 		{
 			request: &http.Request{URL: u, Host: u.Host, Header: forwardedProtoHeader},
 			base:    "https://example.com",
@@ -197,15 +202,26 @@ func TestBuilderFromRequest(t *testing.T) {
 		}
 
 		for _, testCase := range makeURLBuilderTestCases(builder) {
-			url, err := testCase.build()
+			buildURL, err := testCase.build()
 			if err != nil {
 				t.Fatalf("%s: error building url: %v", testCase.description, err)
 			}
 
-			expectedURL := tr.base + testCase.expectedPath
+			var expectedURL string
+			proto, ok := tr.request.Header["X-Forwarded-Proto"]
+			if !ok {
+				expectedURL = tr.base + testCase.expectedPath
+			} else {
+				urlBase, err := url.Parse(tr.base)
+				if err != nil {
+					t.Fatal(err)
+				}
+				urlBase.Scheme = proto[0]
+				expectedURL = urlBase.String() + testCase.expectedPath
+			}
 
-			if url != expectedURL {
-				t.Fatalf("%s: %q != %q", testCase.description, url, expectedURL)
+			if buildURL != expectedURL {
+				t.Fatalf("%s: %q != %q", testCase.description, buildURL, expectedURL)
 			}
 		}
 	}
@@ -229,6 +245,11 @@ func TestBuilderFromRequestWithPrefix(t *testing.T) {
 			request: &http.Request{URL: u, Host: u.Host},
 			base:    "http://example.com/prefix/",
 		},
+
+		{
+			request: &http.Request{URL: u, Host: u.Host, Header: forwardedProtoHeader},
+			base:    "http://example.com/prefix/",
+		},
 		{
 			request: &http.Request{URL: u, Host: u.Host, Header: forwardedProtoHeader},
 			base:    "https://example.com/prefix/",
@@ -253,15 +274,25 @@ func TestBuilderFromRequestWithPrefix(t *testing.T) {
 		}
 
 		for _, testCase := range makeURLBuilderTestCases(builder) {
-			url, err := testCase.build()
+			buildURL, err := testCase.build()
 			if err != nil {
 				t.Fatalf("%s: error building url: %v", testCase.description, err)
 			}
+			var expectedURL string
+			proto, ok := tr.request.Header["X-Forwarded-Proto"]
+			if !ok {
+				expectedURL = tr.base[0:len(tr.base)-1] + testCase.expectedPath
+			} else {
+				urlBase, err := url.Parse(tr.base)
+				if err != nil {
+					t.Fatal(err)
+				}
+				urlBase.Scheme = proto[0]
+				expectedURL = urlBase.String()[0:len(urlBase.String())-1] + testCase.expectedPath
+			}
 
-			expectedURL := tr.base[0:len(tr.base)-1] + testCase.expectedPath
-
-			if url != expectedURL {
-				t.Fatalf("%s: %q != %q", testCase.description, url, expectedURL)
+			if buildURL != expectedURL {
+				t.Fatalf("%s: %q != %q", testCase.description, buildURL, expectedURL)
 			}
 		}
 	}