Merge branch 'security-102-fix-1-28' into '1-28-stable'

Sign OAuth code during auth flow

See merge request gitlab-org/security/gitlab-pages!8

6 files changed, 553 insertions, 145 deletions

diff --git a/acceptance_test.go b/acceptance_test.go
index 69ec8742..1ded9be0 100644
--- a/acceptance_test.go
+++ b/acceptance_test.go
@@ -929,7 +929,7 @@ func TestPrivateArtifactProxyRequest(t *testing.T) {
 			)
 			defer teardown()
 
-			resp, err := GetRedirectPage(t, httpListener, tt.host, tt.path)
+			resp, err := GetRedirectPage(t, httpsListener, tt.host, tt.path)
 			require.NoError(t, err)
 			defer resp.Body.Close()
 
@@ -1131,7 +1131,7 @@ func TestWhenLoginCallbackWithWrongStateShouldFail(t *testing.T) {
 	require.Equal(t, http.StatusUnauthorized, authrsp.StatusCode)
 }
 
-func TestWhenLoginCallbackWithCorrectStateWithoutEndpoint(t *testing.T) {
+func TestWhenLoginCallbackWithUnencryptedCode(t *testing.T) {
 	skipUnlessEnabled(t)
 	teardown := RunPagesProcessWithAuth(t, *pagesBinary, listeners, "")
 	defer teardown()
@@ -1153,8 +1153,92 @@ func TestWhenLoginCallbackWithCorrectStateWithoutEndpoint(t *testing.T) {
 	require.NoError(t, err)
 	defer authrsp.Body.Close()
 
-	// Will cause 503 because token endpoint is not available
-	require.Equal(t, http.StatusServiceUnavailable, authrsp.StatusCode)
+	// Will cause 500 because the code is not encrypted
+	require.Equal(t, http.StatusInternalServerError, authrsp.StatusCode)
+}
+
+// This proves the fix for https://gitlab.com/gitlab-org/gitlab-pages/-/issues/262
+// Read the issue description if any changes to internal/auth/ break this test.
+// Related to https://tools.ietf.org/html/rfc6749#section-10.6.
+func TestHijackedCode(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()
+
+	/****ATTACKER******/
+	// get valid cookie for a different private project
+	targetDomain := "private.domain.com"
+	attackersDomain := "group.auth.gitlab-example.com"
+	attackerCookie, attackerState := getValidCookieAndState(t, targetDomain)
+
+	/****TARGET******/
+	// fool target to click on modified URL with attacker's domain for redirect with a valid state
+	hackedURL := fmt.Sprintf("/auth?domain=http://%s&state=%s", attackersDomain, "irrelevant")
+	maliciousResp, err := GetProxyRedirectPageWithCookie(t, proxyListener, "projects.gitlab-example.com", hackedURL, "", true)
+	require.NoError(t, err)
+	defer maliciousResp.Body.Close()
+
+	pagesCookie := maliciousResp.Header.Get("Set-Cookie")
+
+	/*
+	   OAuth flow happens here...
+	*/
+	maliciousRespURL, err := url.Parse(maliciousResp.Header.Get("Location"))
+	require.NoError(t, err)
+	maliciousState := maliciousRespURL.Query().Get("state")
+
+	// Go to auth page with correct state and code "obtained" from GitLab
+	authrsp, err := GetProxyRedirectPageWithCookie(t, proxyListener,
+		"projects.gitlab-example.com", "/auth?code=1&state="+maliciousState,
+		pagesCookie, true)
+
+	require.NoError(t, err)
+	defer authrsp.Body.Close()
+
+	/****ATTACKER******/
+	// Target is redirected to attacker's domain and attacker receives the proper code
+	require.Equal(t, http.StatusFound, authrsp.StatusCode, "should redirect to attacker's domain")
+	authrspURL, err := url.Parse(authrsp.Header.Get("Location"))
+	require.NoError(t, err)
+	require.Contains(t, authrspURL.String(), attackersDomain)
+
+	// attacker's got the code
+	hijackedCode := authrspURL.Query().Get("code")
+	require.NotEmpty(t, hijackedCode)
+
+	// attacker tries to access private pages content
+	impersonatingRes, err := GetProxyRedirectPageWithCookie(t, proxyListener, targetDomain,
+		"/auth?code="+hijackedCode+"&state="+attackerState, attackerCookie, true)
+	require.NoError(t, err)
+	defer authrsp.Body.Close()
+
+	require.Equal(t, impersonatingRes.StatusCode, http.StatusInternalServerError, "should fail to decode code")
+}
+
+func getValidCookieAndState(t *testing.T, domain string) (string, string) {
+	t.Helper()
+
+	// follow flow to get a valid cookie
+	// visit https://<domain>/
+	rsp, err := GetProxyRedirectPageWithCookie(t, proxyListener, domain, "/", "", true)
+	require.NoError(t, err)
+	defer rsp.Body.Close()
+
+	cookie := rsp.Header.Get("Set-Cookie")
+	require.NotEmpty(t, cookie)
+
+	redirectURL, err := url.Parse(rsp.Header.Get("Location"))
+	require.NoError(t, err)
+
+	state := redirectURL.Query().Get("state")
+	require.NotEmpty(t, state)
+
+	return cookie, state
 }
 
 // makeGitLabPagesAccessStub provides a stub *httptest.Server to check pages_access API call.
@@ -1354,11 +1438,13 @@ func TestAccessControlUnderCustomDomain(t *testing.T) {
 
 	// Will redirect to custom domain
 	require.Equal(t, "private.domain.com", url.Host)
-	require.Equal(t, "1", url.Query().Get("code"))
+	// code must have changed since it's encrypted
+	code := url.Query().Get("code")
+	require.NotEqual(t, "1", code)
 	require.Equal(t, state, url.Query().Get("state"))
 
 	// Run auth callback in custom domain
-	authrsp, err = GetRedirectPageWithCookie(t, httpListener, "private.domain.com", "/auth?code=1&state="+
+	authrsp, err = GetRedirectPageWithCookie(t, httpListener, "private.domain.com", "/auth?code="+code+"&state="+
 		state, cookie)
 
 	require.NoError(t, err)
@@ -1454,11 +1540,13 @@ func TestCustomErrorPageWithAuth(t *testing.T) {
 
 			// Will redirect to custom domain
 			require.Equal(t, tt.domain, url.Host)
-			require.Equal(t, "1", url.Query().Get("code"))
+			// code must have changed since it's encrypted
+			code := url.Query().Get("code")
+			require.NotEqual(t, "1", code)
 			require.Equal(t, state, url.Query().Get("state"))
 
 			// Run auth callback in custom domain
-			authrsp, err = GetRedirectPageWithCookie(t, httpListener, tt.domain, "/auth?code=1&state="+
+			authrsp, err = GetRedirectPageWithCookie(t, httpListener, tt.domain, "/auth?code="+code+"&state="+
 				state, cookie)
 
 			require.NoError(t, err)
@@ -1527,12 +1615,14 @@ func TestAccessControlUnderCustomDomainWithHTTPSProxy(t *testing.T) {
 
 	// Will redirect to custom domain
 	require.Equal(t, "private.domain.com", url.Host)
-	require.Equal(t, "1", url.Query().Get("code"))
+	// code must have changed since it's encrypted
+	code := url.Query().Get("code")
+	require.NotEqual(t, "1", 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)
+		"/auth?code="+code+"&state="+state, cookie, true)
 
 	require.NoError(t, err)
 	defer authrsp.Body.Close()
diff --git a/app.go b/app.go
index 5a195396..baac7631 100644
--- a/app.go
+++ b/app.go
@@ -452,10 +452,7 @@ func runApp(config appConfig) {
 		a.Artifact = artifact.New(config.ArtifactsServer, config.ArtifactsServerTimeout, config.Domain)
 	}
 
-	if config.ClientID != "" {
-		a.Auth = auth.New(config.Domain, config.StoreSecret, config.ClientID, config.ClientSecret,
-			config.RedirectURI, config.GitLabServer)
-	}
+	a.setAuth(config)
 
 	a.Handlers = handlers.New(a.Auth, a.Artifact)
 
@@ -478,6 +475,19 @@ func runApp(config appConfig) {
 	a.Run()
 }
 
+func (a *theApp) setAuth(config appConfig) {
+	if config.ClientID == "" {
+		return
+	}
+
+	var err error
+	a.Auth, err = auth.New(config.Domain, config.StoreSecret, config.ClientID, config.ClientSecret,
+		config.RedirectURI, config.GitLabServer)
+	if err != nil {
+		log.WithError(err).Fatal("could not initialize auth package")
+	}
+}
+
 // fatal will log a fatal error and exit.
 func fatal(err error, message string) {
 	log.WithError(err).Fatal(message)
diff --git a/internal/auth/auth.go b/internal/auth/auth.go
index eaf3c25d..252954a6 100644
--- a/internal/auth/auth.go
+++ b/internal/auth/auth.go
@@ -16,14 +16,14 @@ import (
 	"github.com/gorilla/securecookie"
 	"github.com/gorilla/sessions"
 	log "github.com/sirupsen/logrus"
+	"golang.org/x/crypto/hkdf"
+
 	"gitlab.com/gitlab-org/labkit/errortracking"
 
 	"gitlab.com/gitlab-org/gitlab-pages/internal/httperrors"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/httptransport"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/request"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/source"
-
-	"golang.org/x/crypto/hkdf"
 )
 
 // nolint: gosec
@@ -47,17 +47,23 @@ var (
 	errFailAuth          = errors.New("Failed to authenticate request")
 	errAuthNotConfigured = errors.New("Authentication is not configured")
 	errQueryParameter    = errors.New("Failed to parse domain query parameter")
+
+	errGenerateKeys = errors.New("could not generate auth keys")
 )
 
 // Auth handles authenticating users with GitLab API
 type Auth struct {
-	pagesDomain  string
-	clientID     string
-	clientSecret string
-	redirectURI  string
-	gitLabServer string
-	apiClient    *http.Client
-	store        sessions.Store
+	pagesDomain   string
+	clientID      string
+	clientSecret  string
+	redirectURI   string
+	gitLabServer  string
+	authSecret    string
+	jwtSigningKey []byte
+	jwtExpiry     time.Duration
+	apiClient     *http.Client
+	store         sessions.Store
+	now           func() time.Time // allows to stub time.Now() easily in tests
 }
 
 type tokenResponse struct {
@@ -111,7 +117,7 @@ func (a *Auth) checkSession(w http.ResponseWriter, r *http.Request) (*sessions.S
 	return session, nil
 }
 
-// TryAuthenticate tries to authenticate user and fetch access token if request is a callback to auth
+// TryAuthenticate tries to authenticate user and fetch access token if request is a callback to /auth?
 func (a *Auth) TryAuthenticate(w http.ResponseWriter, r *http.Request, domains source.Source) bool {
 	if a == nil {
 		return false
@@ -166,11 +172,18 @@ func (a *Auth) checkAuthenticationResponse(session *sessions.Session, w http.Res
 		return
 	}
 
-	// Fetch access token with authorization code
-	token, err := a.fetchAccessToken(r.URL.Query().Get("code"))
+	decryptedCode, err := a.DecryptCode(r.URL.Query().Get("code"), getRequestDomain(r))
+	if err != nil {
+		logRequest(r).WithError(err).Error("failed to decrypt secure code")
+		errortracking.Capture(err, errortracking.WithRequest(r))
+		httperrors.Serve500(w)
+		return
+	}
 
-	// Fetching token not OK
+	// Fetch access token with authorization code
+	token, err := a.fetchAccessToken(decryptedCode)
 	if err != nil {
+		// Fetching token not OK
 		logRequest(r).WithError(err).WithField(
 			"redirect_uri", redirectURI,
 		).Error(errFetchAccessToken)
@@ -216,8 +229,8 @@ func (a *Auth) domainAllowed(name string, domains source.Source) bool {
 }
 
 func (a *Auth) handleProxyingAuth(session *sessions.Session, w http.ResponseWriter, r *http.Request, domains source.Source) bool {
-	// If request is for authenticating via custom domain
-	if shouldProxyAuth(r) {
+	// handle auth callback e.g. https://gitlab.io/auth?domain&domain&state=state
+	if shouldProxyAuthToGitlab(r) {
 		domain := r.URL.Query().Get("domain")
 		state := r.URL.Query().Get("state")
 
@@ -266,6 +279,7 @@ func (a *Auth) handleProxyingAuth(session *sessions.Session, w http.ResponseWrit
 	}
 
 	// If auth request callback should be proxied to custom domain
+	// redirect to originating domain set in the cookie as proxy_auth_domain
 	if shouldProxyCallbackToCustomDomain(r, session) {
 		// Get domain started auth process
 		proxyDomain := session.Values["proxy_auth_domain"].(string)
@@ -283,9 +297,30 @@ func (a *Auth) handleProxyingAuth(session *sessions.Session, w http.ResponseWrit
 			return true
 		}
 
-		// Redirect pages under custom domain
-		http.Redirect(w, r, proxyDomain+r.URL.Path+"?"+r.URL.RawQuery, 302)
+		query := r.URL.Query()
+
+		// prevent https://tools.ietf.org/html/rfc6749#section-10.6 and
+		// https://gitlab.com/gitlab-org/gitlab-pages/-/issues/262 by encrypting
+		// and signing the OAuth code
+		signedCode, err := a.EncryptAndSignCode(proxyDomain, query.Get("code"))
+		if err != nil {
+			logRequest(r).WithError(err).Error(errSaveSession)
+			errortracking.Capture(err, errortracking.WithRequest(r))
+
+			httperrors.Serve503(w)
+			return true
+		}
+
+		// prevent forwarding access token, more context on the security issue
+		// https://gitlab.com/gitlab-org/gitlab/-/issues/285244#note_451266051
+		query.Del("token")
+
+		// replace code with signed code
+		query.Set("code", signedCode)
 
+		// Redirect pages to originating domain with code and state to finish
+		// authentication process
+		http.Redirect(w, r, proxyDomain+r.URL.Path+"?"+query.Encode(), 302)
 		return true
 	}
 
@@ -306,7 +341,7 @@ func getRequestDomain(r *http.Request) string {
 	return "http://" + r.Host
 }
 
-func shouldProxyAuth(r *http.Request) bool {
+func shouldProxyAuthToGitlab(r *http.Request) bool {
 	return r.URL.Query().Get("domain") != "" && r.URL.Query().Get("state") != ""
 }
 
@@ -376,6 +411,7 @@ func (a *Auth) checkSessionIsValid(w http.ResponseWriter, r *http.Request) *sess
 		return nil
 	}
 
+	// redirect to /auth?domain=%s&state=%s
 	if a.checkTokenExists(session, w, r) {
 		return nil
 	}
@@ -586,28 +622,37 @@ func logRequest(r *http.Request) *log.Entry {
 	})
 }
 
-// generateKeyPair returns key pair for secure cookie: signing and encryption key
-func generateKeyPair(storeSecret string) ([]byte, []byte) {
-	hash := sha256.New
-	hkdf := hkdf.New(hash, []byte(storeSecret), []byte{}, []byte("PAGES_SIGNING_AND_ENCRYPTION_KEY"))
-	var keys [][]byte
-	for i := 0; i < 2; i++ {
+// generateKeys derives count hkdf keys from a secret, ensuring the key is
+// the same for the same secret used across multiple instances
+func generateKeys(secret string, count int) ([][]byte, error) {
+	keys := make([][]byte, count)
+	hkdfReader := hkdf.New(sha256.New, []byte(secret), []byte{}, []byte("PAGES_SIGNING_AND_ENCRYPTION_KEY"))
+
+	for i := 0; i < count; i++ {
 		key := make([]byte, 32)
-		if _, err := io.ReadFull(hkdf, key); err != nil {
-			log.WithError(err).Fatal("Can't generate key pair for secure cookies")
+		if _, err := io.ReadFull(hkdfReader, key); err != nil {
+			return nil, err
 		}
-		keys = append(keys, key)
+
+		keys[i] = key
+	}
+
+	if len(keys) < count {
+		return nil, errGenerateKeys
 	}
-	return keys[0], keys[1]
-}
 
-func createCookieStore(storeSecret string) sessions.Store {
-	return sessions.NewCookieStore(generateKeyPair(storeSecret))
+	return keys, nil
 }
 
 // New when authentication supported this will be used to create authentication handler
 func New(pagesDomain string, storeSecret string, clientID string, clientSecret string,
-	redirectURI string, gitLabServer string) *Auth {
+	redirectURI string, gitLabServer string) (*Auth, error) {
+	// generate 3 keys, 2 for the cookie store and 1 for JWT signing
+	keys, err := generateKeys(storeSecret, 3)
+	if err != nil {
+		return nil, err
+	}
+
 	return &Auth{
 		pagesDomain:  pagesDomain,
 		clientID:     clientID,
@@ -618,6 +663,10 @@ func New(pagesDomain string, storeSecret string, clientID string, clientSecret s
 			Timeout:   5 * time.Second,
 			Transport: httptransport.InternalTransport,
 		},
-		store: createCookieStore(storeSecret),
-	}
+		store:         sessions.NewCookieStore(keys[0], keys[1]),
+		authSecret:    storeSecret,
+		jwtSigningKey: keys[2],
+		jwtExpiry:     time.Minute,
+		now:           time.Now,
+	}, nil
 }
diff --git a/internal/auth/auth_code.go b/internal/auth/auth_code.go
new file mode 100644
index 00000000..d2fea5a9
--- /dev/null
+++ b/internal/auth/auth_code.go
@@ -0,0 +1,147 @@
+package auth
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/sha256"
+	"encoding/base64"
+	"encoding/hex"
+	"errors"
+	"fmt"
+	"io"
+
+	"github.com/dgrijalva/jwt-go"
+	"github.com/gorilla/securecookie"
+	"golang.org/x/crypto/hkdf"
+)
+
+var (
+	errInvalidToken      = errors.New("invalid token")
+	errEmptyDomainOrCode = errors.New("empty domain or code")
+	errInvalidNonce      = errors.New("invalid nonce")
+	errInvalidCode       = errors.New("invalid code")
+)
+
+// EncryptAndSignCode encrypts the OAuth code deriving the key from the domain.
+// It adds the code and domain as JWT token claims and signs it using signingKey derived from
+// the Auth secret.
+func (a *Auth) EncryptAndSignCode(domain, code string) (string, error) {
+	if domain == "" || code == "" {
+		return "", errEmptyDomainOrCode
+	}
+
+	nonce := base64.URLEncoding.EncodeToString(securecookie.GenerateRandomKey(16))
+
+	aesGcm, err := a.newAesGcmCipher(domain, nonce)
+	if err != nil {
+		return "", err
+	}
+
+	// encrypt code with a randomly generated nonce
+	encryptedCode := aesGcm.Seal(nil, []byte(nonce), []byte(code), nil)
+
+	// generate JWT token claims with encrypted code
+	claims := jwt.MapClaims{
+		// standard claims
+		"iss": "gitlab-pages",
+		"iat": a.now().Unix(),
+		"exp": a.now().Add(a.jwtExpiry).Unix(),
+		// custom claims
+		"domain": domain, // pass the domain so we can validate the signed domain matches the requested domain
+		"code":   hex.EncodeToString(encryptedCode),
+		"nonce":  nonce,
+	}
+
+	return jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString(a.jwtSigningKey)
+}
+
+// DecryptCode decodes the secureCode as a JWT token and validates its signature.
+// It then decrypts the code from the token claims and returns it.
+func (a *Auth) DecryptCode(jwt, domain string) (string, error) {
+	claims, err := a.parseJWTClaims(jwt)
+	if err != nil {
+		return "", err
+	}
+
+	// get nonce and encryptedCode from the JWT claims
+	nonce, ok := claims["nonce"].(string)
+	if !ok {
+		return "", errInvalidNonce
+	}
+
+	encryptedCode, ok := claims["code"].(string)
+	if !ok {
+		return "", errInvalidCode
+	}
+
+	cipherText, err := hex.DecodeString(encryptedCode)
+	if err != nil {
+		return "", err
+	}
+
+	aesGcm, err := a.newAesGcmCipher(domain, nonce)
+	if err != nil {
+		return "", err
+	}
+
+	decryptedCode, err := aesGcm.Open(nil, []byte(nonce), cipherText, nil)
+	if err != nil {
+		return "", err
+	}
+
+	return string(decryptedCode), nil
+}
+
+func (a *Auth) codeKey(domain string) ([]byte, error) {
+	hkdfReader := hkdf.New(sha256.New, []byte(a.authSecret), []byte(domain), []byte("PAGES_AUTH_CODE_ENCRYPTION_KEY"))
+
+	key := make([]byte, 32)
+	if _, err := io.ReadFull(hkdfReader, key); err != nil {
+		return nil, err
+	}
+
+	return key, nil
+}
+
+func (a *Auth) parseJWTClaims(secureCode string) (jwt.MapClaims, error) {
+	token, err := jwt.Parse(secureCode, a.getSigningKey)
+	if err != nil {
+		return nil, err
+	}
+
+	claims, ok := token.Claims.(jwt.MapClaims)
+	if !ok || !token.Valid {
+		return nil, errInvalidToken
+	}
+
+	return claims, nil
+}
+
+func (a *Auth) getSigningKey(token *jwt.Token) (interface{}, error) {
+	// Don't forget to validate the alg is what you expect:
+	if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
+		return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
+	}
+
+	return a.jwtSigningKey, nil
+}
+
+func (a *Auth) newAesGcmCipher(domain, nonce string) (cipher.AEAD, error) {
+	// get the same key for a domain
+	key, err := a.codeKey(domain)
+	if err != nil {
+		return nil, err
+	}
+
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+
+	aesGcm, err := cipher.NewGCMWithNonceSize(block, len(nonce))
+	if err != nil {
+		return nil, err
+	}
+
+	return aesGcm, nil
+}
diff --git a/internal/auth/auth_code_test.go b/internal/auth/auth_code_test.go
new file mode 100644
index 00000000..d54fcc7e
--- /dev/null
+++ b/internal/auth/auth_code_test.go
@@ -0,0 +1,99 @@
+package auth
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestEncryptAndDecryptSignedCode(t *testing.T) {
+	auth := createTestAuth(t, "")
+
+	tests := map[string]struct {
+		auth              *Auth
+		encDomain         string
+		code              string
+		expectedEncErrMsg string
+		decDomain         string
+		expectedDecErrMsg string
+	}{
+		"happy_path": {
+			auth:      auth,
+			encDomain: "domain",
+			decDomain: "domain",
+			code:      "code",
+		},
+		"empty_domain": {
+			auth:              auth,
+			encDomain:         "",
+			code:              "code",
+			expectedEncErrMsg: "empty domain or code",
+		},
+		"empty_code": {
+			auth:              auth,
+			encDomain:         "domain",
+			code:              "",
+			expectedEncErrMsg: "empty domain or code",
+		},
+		"different_dec_domain": {
+			auth:              auth,
+			encDomain:         "domain",
+			decDomain:         "another",
+			code:              "code",
+			expectedDecErrMsg: "cipher: message authentication failed",
+		},
+		"expired_token": {
+			auth: func() *Auth {
+				newAuth := *auth
+				newAuth.jwtExpiry = time.Nanosecond
+				newAuth.now = func() time.Time {
+					return time.Time{}
+				}
+
+				return &newAuth
+			}(),
+			encDomain:         "domain",
+			code:              "code",
+			decDomain:         "domain",
+			expectedDecErrMsg: "Token is expired",
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			encCode, err := test.auth.EncryptAndSignCode(test.encDomain, test.code)
+			if test.expectedEncErrMsg != "" {
+				require.EqualError(t, err, test.expectedEncErrMsg)
+				require.Empty(t, encCode)
+				return
+			}
+
+			require.NoError(t, err)
+			require.NotEmpty(t, encCode)
+
+			decCode, err := test.auth.DecryptCode(encCode, test.decDomain)
+			if test.expectedDecErrMsg != "" {
+				require.EqualError(t, err, test.expectedDecErrMsg)
+				require.Empty(t, decCode)
+				return
+			}
+
+			require.NoError(t, err)
+			require.Equal(t, test.code, decCode)
+		})
+	}
+}
+
+func TestDecryptCodeWithInvalidJWT(t *testing.T) {
+	auth1 := createTestAuth(t, "")
+	auth2 := createTestAuth(t, "")
+	auth2.jwtSigningKey = []byte("another signing key")
+
+	encCode, err := auth1.EncryptAndSignCode("domain", "code")
+	require.NoError(t, err)
+
+	decCode, err := auth2.DecryptCode(encCode, "domain")
+	require.EqualError(t, err, "signature is invalid")
+	require.Empty(t, decCode)
+}
diff --git a/internal/auth/auth_test.go b/internal/auth/auth_test.go
index 39a533b3..ce7d8320 100644
--- a/internal/auth/auth_test.go
+++ b/internal/auth/auth_test.go
@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/url"
+	"strings"
 	"testing"
 
 	"github.com/gorilla/sessions"
@@ -16,17 +17,19 @@ import (
 	"gitlab.com/gitlab-org/gitlab-pages/internal/source"
 )
 
-func createAuth(t *testing.T) *Auth {
-	return New("pages.gitlab-example.com",
+func createTestAuth(t *testing.T, url string) *Auth {
+	t.Helper()
+
+	a, err := New("pages.gitlab-example.com",
 		"something-very-secret",
 		"id",
 		"secret",
 		"http://pages.gitlab-example.com/auth",
-		"http://gitlab-example.com")
-}
+		url)
+
+	require.NoError(t, err)
 
-func defaultCookieStore() sessions.Store {
-	return createCookieStore("something-very-secret")
+	return a
 }
 
 type domainMock struct {
@@ -48,10 +51,13 @@ func (dm *domainMock) ServeNotFoundAuthFailed(w http.ResponseWriter, r *http.Req
 // Which leads to negative side effects: we can't test encryption, and cookie params
 // like max-age and secure are not being properly set
 // To avoid that we use fake request, and set only session cookie without copying context
-func setSessionValues(r *http.Request, values map[interface{}]interface{}) {
-	tmpRequest, _ := http.NewRequest("GET", "/", nil)
+func setSessionValues(t *testing.T, r *http.Request, store sessions.Store, values map[interface{}]interface{}) {
+	t.Helper()
+
+	tmpRequest, err := http.NewRequest("GET", "/", nil)
+	require.NoError(t, err)
+
 	result := httptest.NewRecorder()
-	store := defaultCookieStore()
 
 	session, _ := store.Get(tmpRequest, "gitlab-pages")
 	session.Values = values
@@ -63,7 +69,7 @@ func setSessionValues(r *http.Request, values map[interface{}]interface{}) {
 }
 
 func TestTryAuthenticate(t *testing.T) {
-	auth := createAuth(t)
+	auth := createTestAuth(t, "")
 
 	result := httptest.NewRecorder()
 	reqURL, err := url.Parse("/something/else")
@@ -75,11 +81,12 @@ func TestTryAuthenticate(t *testing.T) {
 }
 
 func TestTryAuthenticateWithError(t *testing.T) {
-	auth := createAuth(t)
+	auth := createTestAuth(t, "")
 
 	result := httptest.NewRecorder()
 	reqURL, err := url.Parse("/auth?error=access_denied")
 	require.NoError(t, err)
+
 	reqURL.Scheme = request.SchemeHTTPS
 	r := &http.Request{URL: reqURL}
 
@@ -88,8 +95,7 @@ func TestTryAuthenticateWithError(t *testing.T) {
 }
 
 func TestTryAuthenticateWithCodeButInvalidState(t *testing.T) {
-	store := sessions.NewCookieStore([]byte("something-very-secret"))
-	auth := createAuth(t)
+	auth := createTestAuth(t, "")
 
 	result := httptest.NewRecorder()
 	reqURL, err := url.Parse("/auth?code=1&state=invalid")
@@ -97,7 +103,9 @@ func TestTryAuthenticateWithCodeButInvalidState(t *testing.T) {
 	reqURL.Scheme = request.SchemeHTTPS
 	r := &http.Request{URL: reqURL}
 
-	session, _ := store.Get(r, "gitlab-pages")
+	session, err := auth.store.Get(r, "gitlab-pages")
+	require.NoError(t, err)
+
 	session.Values["state"] = "state"
 	session.Save(r, result)
 
@@ -105,7 +113,36 @@ func TestTryAuthenticateWithCodeButInvalidState(t *testing.T) {
 	require.Equal(t, 401, result.Code)
 }
 
+func TestTryAuthenticateRemoveTokenFromRedirect(t *testing.T) {
+	auth := createTestAuth(t, "")
+
+	result := httptest.NewRecorder()
+	reqURL, err := url.Parse("/auth?code=1&state=state&token=secret")
+	require.NoError(t, err)
+
+	require.Equal(t, reqURL.Query().Get("token"), "secret", "token is present before redirecting")
+	reqURL.Scheme = request.SchemeHTTPS
+	r := &http.Request{URL: reqURL}
+
+	session, err := auth.store.Get(r, "gitlab-pages")
+	require.NoError(t, err)
+
+	session.Values["state"] = "state"
+	session.Values["proxy_auth_domain"] = "https://domain.com"
+	session.Save(r, result)
+
+	require.Equal(t, true, auth.TryAuthenticate(result, r, source.NewMockSource()))
+	require.Equal(t, http.StatusFound, result.Code)
+
+	redirect, err := url.Parse(result.Header().Get("Location"))
+	require.NoError(t, err)
+
+	require.Empty(t, redirect.Query().Get("token"), "token is gone after redirecting")
+}
+
 func testTryAuthenticateWithCodeAndState(t *testing.T, https bool) {
+	t.Helper()
+
 	apiServer := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		switch r.URL.Path {
 		case "/oauth/token":
@@ -125,14 +162,17 @@ func testTryAuthenticateWithCodeAndState(t *testing.T, https bool) {
 	apiServer.Start()
 	defer apiServer.Close()
 
-	auth := New("pages.gitlab-example.com",
-		"something-very-secret",
-		"id",
-		"secret",
-		"http://pages.gitlab-example.com/auth",
-		apiServer.URL)
+	auth := createTestAuth(t, apiServer.URL)
+
+	domain := apiServer.URL
+	if https {
+		domain = strings.Replace(apiServer.URL, "http://", "https://", -1)
+	}
 
-	r, err := http.NewRequest("GET", "/auth?code=1&state=state", nil)
+	code, err := auth.EncryptAndSignCode(domain, "1")
+	require.NoError(t, err)
+
+	r, err := http.NewRequest("GET", "/auth?code="+code+"&state=state", nil)
 	require.NoError(t, err)
 	if https {
 		r.URL.Scheme = request.SchemeHTTPS
@@ -140,14 +180,16 @@ func testTryAuthenticateWithCodeAndState(t *testing.T, https bool) {
 		r.URL.Scheme = request.SchemeHTTP
 	}
 
-	setSessionValues(r, map[interface{}]interface{}{
+	r.Host = strings.TrimPrefix(apiServer.URL, "http://")
+
+	setSessionValues(t, r, auth.store, map[interface{}]interface{}{
 		"uri":   "https://pages.gitlab-example.com/project/",
 		"state": "state",
 	})
 
 	result := httptest.NewRecorder()
 	require.Equal(t, true, auth.TryAuthenticate(result, r, source.NewMockSource()))
-	require.Equal(t, 302, result.Code)
+	require.Equal(t, http.StatusFound, result.Code)
 	require.Equal(t, "https://pages.gitlab-example.com/project/", result.Header().Get("Location"))
 	require.Equal(t, 600, result.Result().Cookies()[0].MaxAge)
 	require.Equal(t, https, result.Result().Cookies()[0].Secure)
@@ -177,13 +219,7 @@ func TestCheckAuthenticationWhenAccess(t *testing.T) {
 	apiServer.Start()
 	defer apiServer.Close()
 
-	store := defaultCookieStore()
-	auth := New("pages.gitlab-example.com",
-		"something-very-secret",
-		"id",
-		"secret",
-		"http://pages.gitlab-example.com/auth",
-		apiServer.URL)
+	auth := createTestAuth(t, apiServer.URL)
 
 	result := httptest.NewRecorder()
 	reqURL, err := url.Parse("/auth?code=1&state=state")
@@ -191,7 +227,9 @@ func TestCheckAuthenticationWhenAccess(t *testing.T) {
 	reqURL.Scheme = request.SchemeHTTPS
 	r := &http.Request{URL: reqURL}
 
-	session, _ := store.Get(r, "gitlab-pages")
+	session, err := auth.store.Get(r, "gitlab-pages")
+	require.NoError(t, err)
+
 	session.Values["access_token"] = "abc"
 	session.Save(r, result)
 	contentServed := auth.CheckAuthentication(result, r, &domainMock{projectID: 1000})
@@ -217,13 +255,7 @@ func TestCheckAuthenticationWhenNoAccess(t *testing.T) {
 	apiServer.Start()
 	defer apiServer.Close()
 
-	store := defaultCookieStore()
-	auth := New("pages.gitlab-example.com",
-		"something-very-secret",
-		"id",
-		"secret",
-		"http://pages.gitlab-example.com/auth",
-		apiServer.URL)
+	auth := createTestAuth(t, apiServer.URL)
 
 	w := httptest.NewRecorder()
 
@@ -232,7 +264,9 @@ func TestCheckAuthenticationWhenNoAccess(t *testing.T) {
 	reqURL.Scheme = request.SchemeHTTPS
 	r := &http.Request{URL: reqURL}
 
-	session, _ := store.Get(r, "gitlab-pages")
+	session, err := auth.store.Get(r, "gitlab-pages")
+	require.NoError(t, err)
+
 	session.Values["access_token"] = "abc"
 	session.Save(r, w)
 
@@ -265,22 +299,19 @@ func TestCheckAuthenticationWhenInvalidToken(t *testing.T) {
 	apiServer.Start()
 	defer apiServer.Close()
 
-	store := defaultCookieStore()
-	auth := New("pages.gitlab-example.com",
-		"something-very-secret",
-		"id",
-		"secret",
-		"http://pages.gitlab-example.com/auth",
-		apiServer.URL)
+	auth := createTestAuth(t, apiServer.URL)
 
 	result := httptest.NewRecorder()
 	reqURL, err := url.Parse("/auth?code=1&state=state")
 	require.NoError(t, err)
 	r := &http.Request{URL: reqURL}
 
-	session, _ := store.Get(r, "gitlab-pages")
+	session, err := auth.store.Get(r, "gitlab-pages")
+	require.NoError(t, err)
+
 	session.Values["access_token"] = "abc"
-	session.Save(r, result)
+	err = session.Save(r, result)
+	require.NoError(t, err)
 
 	contentServed := auth.CheckAuthentication(result, r, &domainMock{projectID: 1000})
 	require.True(t, contentServed)
@@ -303,13 +334,7 @@ func TestCheckAuthenticationWithoutProject(t *testing.T) {
 	apiServer.Start()
 	defer apiServer.Close()
 
-	store := defaultCookieStore()
-	auth := New("pages.gitlab-example.com",
-		"something-very-secret",
-		"id",
-		"secret",
-		"http://pages.gitlab-example.com/auth",
-		apiServer.URL)
+	auth := createTestAuth(t, apiServer.URL)
 
 	result := httptest.NewRecorder()
 	reqURL, err := url.Parse("/auth?code=1&state=state")
@@ -317,7 +342,9 @@ func TestCheckAuthenticationWithoutProject(t *testing.T) {
 	reqURL.Scheme = request.SchemeHTTPS
 	r := &http.Request{URL: reqURL}
 
-	session, _ := store.Get(r, "gitlab-pages")
+	session, err := auth.store.Get(r, "gitlab-pages")
+	require.NoError(t, err)
+
 	session.Values["access_token"] = "abc"
 	session.Save(r, result)
 
@@ -343,19 +370,16 @@ func TestCheckAuthenticationWithoutProjectWhenInvalidToken(t *testing.T) {
 	apiServer.Start()
 	defer apiServer.Close()
 
-	store := defaultCookieStore()
-	auth := New("pages.gitlab-example.com",
-		"something-very-secret",
-		"id",
-		"secret",
-		"http://pages.gitlab-example.com/auth",
-		apiServer.URL)
+	auth := createTestAuth(t, apiServer.URL)
 
 	result := httptest.NewRecorder()
 	reqURL, err := url.Parse("/auth?code=1&state=state")
 	require.NoError(t, err)
 	r := &http.Request{URL: reqURL}
-	session, _ := store.Get(r, "gitlab-pages")
+
+	session, err := auth.store.Get(r, "gitlab-pages")
+	require.NoError(t, err)
+
 	session.Values["access_token"] = "abc"
 	session.Save(r, result)
 
@@ -364,28 +388,31 @@ func TestCheckAuthenticationWithoutProjectWhenInvalidToken(t *testing.T) {
 	require.Equal(t, 302, result.Code)
 }
 
-func TestGenerateKeyPair(t *testing.T) {
-	signingSecret, encryptionSecret := generateKeyPair("something-very-secret")
-	require.NotEqual(t, fmt.Sprint(signingSecret), fmt.Sprint(encryptionSecret))
-	require.Equal(t, len(signingSecret), 32)
-	require.Equal(t, len(encryptionSecret), 32)
+func TestGenerateKeys(t *testing.T) {
+	keys, err := generateKeys("something-very-secret", 3)
+	require.NoError(t, err)
+	require.Len(t, keys, 3)
+
+	require.NotEqual(t, fmt.Sprint(keys[0]), fmt.Sprint(keys[1]))
+	require.NotEqual(t, fmt.Sprint(keys[0]), fmt.Sprint(keys[2]))
+	require.NotEqual(t, fmt.Sprint(keys[1]), fmt.Sprint(keys[2]))
+
+	require.Equal(t, len(keys[0]), 32)
+	require.Equal(t, len(keys[1]), 32)
+	require.Equal(t, len(keys[2]), 32)
 }
 
 func TestGetTokenIfExistsWhenTokenExists(t *testing.T) {
-	store := sessions.NewCookieStore([]byte("something-very-secret"))
-	auth := New("pages.gitlab-example.com",
-		"something-very-secret",
-		"id",
-		"secret",
-		"http://pages.gitlab-example.com/auth",
-		"")
+	auth := createTestAuth(t, "")
 
 	result := httptest.NewRecorder()
 	reqURL, err := url.Parse("/")
 	require.NoError(t, err)
 	r := &http.Request{URL: reqURL}
 
-	session, _ := store.Get(r, "gitlab-pages")
+	session, err := auth.store.Get(r, "gitlab-pages")
+	require.NoError(t, err)
+
 	session.Values["access_token"] = "abc"
 	session.Save(r, result)
 
@@ -395,20 +422,16 @@ func TestGetTokenIfExistsWhenTokenExists(t *testing.T) {
 }
 
 func TestGetTokenIfExistsWhenTokenDoesNotExist(t *testing.T) {
-	store := sessions.NewCookieStore([]byte("something-very-secret"))
-	auth := New("pages.gitlab-example.com",
-		"something-very-secret",
-		"id",
-		"secret",
-		"http://pages.gitlab-example.com/auth",
-		"")
+	auth := createTestAuth(t, "")
 
 	result := httptest.NewRecorder()
 	reqURL, err := url.Parse("http://pages.gitlab-example.com/test")
 	require.NoError(t, err)
 	r := &http.Request{URL: reqURL, Host: "pages.gitlab-example.com", RequestURI: "/test"}
 
-	session, _ := store.Get(r, "gitlab-pages")
+	session, err := auth.store.Get(r, "gitlab-pages")
+	require.NoError(t, err)
+
 	session.Save(r, result)
 
 	token, err := auth.GetTokenIfExists(result, r)
@@ -417,12 +440,7 @@ func TestGetTokenIfExistsWhenTokenDoesNotExist(t *testing.T) {
 }
 
 func TestCheckResponseForInvalidTokenWhenInvalidToken(t *testing.T) {
-	auth := New("pages.gitlab-example.com",
-		"something-very-secret",
-		"id",
-		"secret",
-		"http://pages.gitlab-example.com/auth",
-		"")
+	auth := createTestAuth(t, "")
 
 	result := httptest.NewRecorder()
 	reqURL, err := url.Parse("http://pages.gitlab-example.com/test")
@@ -437,12 +455,7 @@ func TestCheckResponseForInvalidTokenWhenInvalidToken(t *testing.T) {
 }
 
 func TestCheckResponseForInvalidTokenWhenNotInvalidToken(t *testing.T) {
-	auth := New("pages.gitlab-example.com",
-		"something-very-secret",
-		"id",
-		"secret",
-		"http://pages.gitlab-example.com/auth",
-		"")
+	auth := createTestAuth(t, "")
 
 	result := httptest.NewRecorder()
 	reqURL, err := url.Parse("/something")