Encrypt and sign OAuth code

Add AES GCM encryption/decryption to auth

Add signing key to Auth

Abstract key generation and Auth init to their own funcs.
Cleanup and DRY unit tests.

Use same code parameter in auth redirect

Cleanup auth and add tests for enc/dec oauth code

Add acceptance test for fix

Apply suggestion from review

Add missing test and apply feedback

Fix unit test

Simplify acceptance test

4 files changed, 439 insertions, 131 deletions

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")