Fix https downgrade for pages behind proxy

We can't rely on r.TLS when pages are served behind proxy
So we save https flag to a context for later usage
Right now I'm trying to keep changes to a minimum since
I'm planning to backport this to older versions
That's why https flag is not refactored throughout the codebase

The alternative way would be to use gorilla's proxy headers
I'm planning to refactor to that version later

7 files changed, 143 insertions, 5 deletions

diff --git a/acceptance_test.go b/acceptance_test.go
index 8023c4ff..587fc27e 100644
--- a/acceptance_test.go
+++ b/acceptance_test.go
@@ -37,6 +37,7 @@ var listeners = []ListenSpec{
 var (
 	httpListener  = listeners[0]
 	httpsListener = listeners[2]
+	proxyListener = listeners[4]
 )
 
 func skipUnlessEnabled(t *testing.T, conditions ...string) {
@@ -944,6 +945,71 @@ func TestAccessControlUnderCustomDomain(t *testing.T) {
 	assert.Equal(t, http.StatusOK, authrsp.StatusCode)
 }
 
+func TestAccessControlUnderCustomDomainWithHTTPSProxy(t *testing.T) {
+	skipUnlessEnabled(t, "not-inplace-chroot")
+
+	testServer := makeGitLabPagesAccessStub(t)
+	testServer.Start()
+	defer testServer.Close()
+
+	teardown := RunPagesProcessWithAuthServer(t, *pagesBinary, listeners, "", testServer.URL)
+	defer teardown()
+
+	rsp, err := GetProxyRedirectPageWithCookie(t, proxyListener, "private.domain.com", "/", "", true)
+	require.NoError(t, err)
+	defer rsp.Body.Close()
+
+	cookie := rsp.Header.Get("Set-Cookie")
+
+	url, err := url.Parse(rsp.Header.Get("Location"))
+	require.NoError(t, err)
+
+	state := url.Query().Get("state")
+	require.Equal(t, url.Query().Get("domain"), "https://private.domain.com")
+	pagesrsp, err := GetProxyRedirectPageWithCookie(t, proxyListener, url.Host, url.Path+"?"+url.RawQuery, "", true)
+	require.NoError(t, err)
+	defer pagesrsp.Body.Close()
+
+	pagescookie := pagesrsp.Header.Get("Set-Cookie")
+
+	// Go to auth page with correct state will cause fetching the token
+	authrsp, err := GetProxyRedirectPageWithCookie(t, proxyListener,
+		"projects.gitlab-example.com", "/auth?code=1&state="+state,
+		pagescookie, true)
+
+	require.NoError(t, err)
+	defer authrsp.Body.Close()
+
+	url, err = url.Parse(authrsp.Header.Get("Location"))
+	require.NoError(t, err)
+
+	// Will redirect to custom domain
+	require.Equal(t, "private.domain.com", url.Host)
+	require.Equal(t, "1", url.Query().Get("code"))
+	require.Equal(t, state, url.Query().Get("state"))
+
+	// Run auth callback in custom domain
+	authrsp, err = GetProxyRedirectPageWithCookie(t, proxyListener, "private.domain.com",
+		"/auth?code=1&state="+state, cookie, true)
+
+	require.NoError(t, err)
+	defer authrsp.Body.Close()
+
+	// Will redirect to the page
+	cookie = authrsp.Header.Get("Set-Cookie")
+	require.Equal(t, http.StatusFound, authrsp.StatusCode)
+
+	url, err = url.Parse(authrsp.Header.Get("Location"))
+	require.NoError(t, err)
+
+	// Will redirect to custom domain
+	require.Equal(t, "https://private.domain.com/", url.String())
+	// Fetch page in custom domain
+	authrsp, err = GetProxyRedirectPageWithCookie(t, proxyListener, "private.domain.com", "/",
+		cookie, true)
+	require.Equal(t, http.StatusOK, authrsp.StatusCode)
+}
+
 func TestAccessControlGroupDomain404RedirectsAuth(t *testing.T) {
 	skipUnlessEnabled(t)
 	teardown := RunPagesProcessWithAuth(t, *pagesBinary, listeners, "")
diff --git a/app.go b/app.go
index da1a5315..3efc22c2 100644
--- a/app.go
+++ b/app.go
@@ -26,6 +26,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-pages/internal/domain"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/httperrors"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/netutil"
+	"gitlab.com/gitlab-org/gitlab-pages/internal/request"
 	"gitlab.com/gitlab-org/gitlab-pages/metrics"
 )
 
@@ -230,6 +231,8 @@ func (a *theApp) serveFileOrNotFound(domain *domain.D) http.HandlerFunc {
 
 func (a *theApp) ServeHTTP(ww http.ResponseWriter, r *http.Request) {
 	https := r.TLS != nil
+	r = request.WithHTTPSFlag(r, https)
+
 	headerConfig.AddCustomHeaders(ww, a.CustomHeaders)
 
 	a.serveContent(ww, r, https)
@@ -238,6 +241,7 @@ func (a *theApp) ServeHTTP(ww http.ResponseWriter, r *http.Request) {
 func (a *theApp) ServeProxy(ww http.ResponseWriter, r *http.Request) {
 	forwardedProto := r.Header.Get(xForwardedProto)
 	https := forwardedProto == xForwardedProtoHTTPS
+	r = request.WithHTTPSFlag(r, https)
 
 	if forwardedHost := r.Header.Get(xForwardedHost); forwardedHost != "" {
 		r.Host = forwardedHost
diff --git a/helpers_test.go b/helpers_test.go
index 61fa5279..ad7c65f1 100644
--- a/helpers_test.go
+++ b/helpers_test.go
@@ -314,15 +314,31 @@ func GetRedirectPage(t *testing.T, spec ListenSpec, host, urlsuffix string) (*ht
 	return GetRedirectPageWithCookie(t, spec, host, urlsuffix, "")
 }
 
+func GetProxyRedirectPageWithCookie(t *testing.T, spec ListenSpec, host string, urlsuffix string, cookie string, https bool) (*http.Response, error) {
+	schema := "http"
+	if https {
+		schema = "https"
+	}
+	header := http.Header{
+		"X-Forwarded-Proto": []string{schema},
+		"X-Forwarded-Host":  []string{host},
+		"cookie":            []string{cookie},
+	}
+
+	return GetRedirectPageWithHeaders(t, spec, host, urlsuffix, header)
+}
+
 func GetRedirectPageWithCookie(t *testing.T, spec ListenSpec, host, urlsuffix string, cookie string) (*http.Response, error) {
+	return GetRedirectPageWithHeaders(t, spec, host, urlsuffix, http.Header{"cookie": []string{cookie}})
+}
+
+func GetRedirectPageWithHeaders(t *testing.T, spec ListenSpec, host, urlsuffix string, header http.Header) (*http.Response, error) {
 	url := spec.URL(urlsuffix)
 	req, err := http.NewRequest("GET", url, nil)
 	if err != nil {
 		return nil, err
 	}
-	if cookie != "" {
-		req.Header.Set("Cookie", cookie)
-	}
+	req.Header = header
 
 	req.Host = host
 
diff --git a/internal/auth/auth.go b/internal/auth/auth.go
index d6cbdff1..920c5d12 100644
--- a/internal/auth/auth.go
+++ b/internal/auth/auth.go
@@ -22,6 +22,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-pages/internal/domain"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/httperrors"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/httptransport"
+	"gitlab.com/gitlab-org/gitlab-pages/internal/request"
 
 	"golang.org/x/crypto/hkdf"
 )
@@ -285,14 +286,14 @@ func (a *Auth) handleProxyingAuth(session *sessions.Session, w http.ResponseWrit
 }
 
 func getRequestAddress(r *http.Request) string {
-	if r.TLS != nil {
+	if request.IsHTTPS(r) {
 		return "https://" + r.Host + r.RequestURI
 	}
 	return "http://" + r.Host + r.RequestURI
 }
 
 func getRequestDomain(r *http.Request) string {
-	if r.TLS != nil {
+	if request.IsHTTPS(r) {
 		return "https://" + r.Host
 	}
 	return "http://" + r.Host
diff --git a/internal/auth/auth_test.go b/internal/auth/auth_test.go
index 2fbbb938..00cdbd5b 100644
--- a/internal/auth/auth_test.go
+++ b/internal/auth/auth_test.go
@@ -13,6 +13,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"gitlab.com/gitlab-org/gitlab-pages/internal/domain"
+	"gitlab.com/gitlab-org/gitlab-pages/internal/request"
 )
 
 func createAuth(t *testing.T) *Auth {
@@ -214,6 +215,7 @@ func TestCheckAuthenticationWhenInvalidToken(t *testing.T) {
 	reqURL, err := url.Parse("/auth?code=1&state=state")
 	require.NoError(t, err)
 	r := &http.Request{URL: reqURL}
+	r = request.WithHTTPSFlag(r, false)
 
 	session, _ := store.Get(r, "gitlab-pages")
 	session.Values["access_token"] = "abc"
@@ -289,6 +291,7 @@ func TestCheckAuthenticationWithoutProjectWhenInvalidToken(t *testing.T) {
 	reqURL, err := url.Parse("/auth?code=1&state=state")
 	require.NoError(t, err)
 	r := &http.Request{URL: reqURL}
+	r = request.WithHTTPSFlag(r, false)
 
 	session, _ := store.Get(r, "gitlab-pages")
 	session.Values["access_token"] = "abc"
diff --git a/internal/request/request.go b/internal/request/request.go
new file mode 100644
index 00000000..dad6af3d
--- /dev/null
+++ b/internal/request/request.go
@@ -0,0 +1,24 @@
+package request
+
+import (
+	"context"
+	"net/http"
+)
+
+type ctxKey string
+
+const (
+	ctxHTTPSKey ctxKey = "https"
+)
+
+// WithHTTPSFlag saves https flag in request's context
+func WithHTTPSFlag(r *http.Request, https bool) *http.Request {
+	ctx := context.WithValue(r.Context(), ctxHTTPSKey, https)
+
+	return r.WithContext(ctx)
+}
+
+// IsHTTPS restores https flag from request's context
+func IsHTTPS(r *http.Request) bool {
+	return r.Context().Value(ctxHTTPSKey).(bool)
+}
diff --git a/internal/request/request_test.go b/internal/request/request_test.go
new file mode 100644
index 00000000..1f47ee2e
--- /dev/null
+++ b/internal/request/request_test.go
@@ -0,0 +1,24 @@
+package request
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWithHTTPSFlag(t *testing.T) {
+	r, err := http.NewRequest("GET", "/", nil)
+	require.NoError(t, err)
+
+	assert.Panics(t, func() {
+		IsHTTPS(r)
+	})
+
+	httpsRequest := WithHTTPSFlag(r, true)
+	assert.True(t, IsHTTPS(httpsRequest))
+
+	httpRequest := WithHTTPSFlag(r, false)
+	assert.False(t, IsHTTPS(httpRequest))
+}