Fix problem with the public suffix listed pages domain

2 files changed, 40 insertions, 27 deletions

diff --git a/acceptance_test.go b/acceptance_test.go
index 8a842290..516b90d2 100644
--- a/acceptance_test.go
+++ b/acceptance_test.go
@@ -600,9 +600,15 @@ func TestWhenAuthIsEnabledPrivateWillRedirectToAuthorize(t *testing.T) {
 
 	assert.Equal(t, http.StatusFound, rsp.StatusCode)
 	assert.Equal(t, 1, len(rsp.Header["Location"]))
-
 	url, err := url.Parse(rsp.Header.Get("Location"))
 	require.NoError(t, err)
+	rsp, err = GetRedirectPage(t, httpsListener, url.Host, url.Path+"?"+url.RawQuery)
+
+	assert.Equal(t, http.StatusFound, rsp.StatusCode)
+	assert.Equal(t, 1, len(rsp.Header["Location"]))
+
+	url, err = url.Parse(rsp.Header.Get("Location"))
+	require.NoError(t, err)
 
 	assert.Equal(t, "https", url.Scheme)
 	assert.Equal(t, "gitlab-auth.com", url.Host)
@@ -849,19 +855,39 @@ func TestAccessControl(t *testing.T) {
 			defer rsp.Body.Close()
 
 			assert.Equal(t, http.StatusFound, rsp.StatusCode)
-
 			cookie := rsp.Header.Get("Set-Cookie")
 
+			// Redirects to the gitlab pages root domain for authentication flow
 			url, err := url.Parse(rsp.Header.Get("Location"))
 			require.NoError(t, err)
+			assert.Equal(t, "gitlab-example.com", url.Host)
+			assert.Equal(t, "/auth", url.Path)
+			state := url.Query().Get("state")
+
+			rsp, err = GetRedirectPage(t, httpsListener, url.Host, url.Path+"?"+url.RawQuery)
+
+			require.NoError(t, err)
+			defer rsp.Body.Close()
+
+			assert.Equal(t, http.StatusFound, rsp.StatusCode)
+			pagesDomainCookie := rsp.Header.Get("Set-Cookie")
 
 			// Go to auth page with correct state will cause fetching the token
 			authrsp, err := GetRedirectPageWithCookie(t, httpsListener, "gitlab-example.com", "/auth?code=1&state="+
-				url.Query().Get("state"), cookie)
+				state, pagesDomainCookie)
 
 			require.NoError(t, err)
 			defer authrsp.Body.Close()
 
+			// Will redirect auth callback to correct host
+			url, err = url.Parse(authrsp.Header.Get("Location"))
+			require.NoError(t, err)
+			assert.Equal(t, c.Host, url.Host)
+			assert.Equal(t, "/auth", url.Path)
+
+			// Request auth callback in project domain
+			authrsp, err = GetRedirectPageWithCookie(t, httpsListener, url.Host, url.Path+"?"+url.RawQuery, cookie)
+
 			// server returns the ticket, user will be redirected to the project page
 			assert.Equal(t, http.StatusFound, authrsp.StatusCode)
 			cookie = authrsp.Header.Get("Set-Cookie")
diff --git a/internal/auth/auth.go b/internal/auth/auth.go
index ea185ea2..e88cf7a2 100644
--- a/internal/auth/auth.go
+++ b/internal/auth/auth.go
@@ -51,18 +51,10 @@ type errorResponse struct {
 func (a *Auth) getSessionFromStore(r *http.Request) (*sessions.Session, error) {
 	store := sessions.NewCookieStore([]byte(a.storeSecret))
 
-	if strings.HasSuffix(r.Host, a.pagesDomain) {
-		// GitLab pages wide cookie
-		store.Options = &sessions.Options{
-			Path:   "/",
-			Domain: a.pagesDomain,
-		}
-	} else {
-		// Cookie just for this domain
-		store.Options = &sessions.Options{
-			Path:   "/",
-			Domain: r.Host,
-		}
+	// Cookie just for this domain
+	store.Options = &sessions.Options{
+		Path:   "/",
+		Domain: r.Host,
 	}
 
 	return store.Get(r, "gitlab-pages")
@@ -160,14 +152,14 @@ func (a *Auth) TryAuthenticate(w http.ResponseWriter, r *http.Request) bool {
 func (a *Auth) handleProxyingAuth(session *sessions.Session, w http.ResponseWriter, r *http.Request) bool {
 	// If request is for authenticating via custom domain
 	if shouldProxyAuth(r) {
-		customDomain := r.URL.Query().Get("domain")
+		domain := r.URL.Query().Get("domain")
 		state := r.URL.Query().Get("state")
-		log.WithField("domain", customDomain).Debug("User is authenticating via custom domain")
+		log.WithField("domain", domain).Debug("User is authenticating via domain")
 
 		if r.TLS != nil {
-			session.Values["proxy_auth_domain"] = "https://" + customDomain
+			session.Values["proxy_auth_domain"] = "https://" + domain
 		} else {
-			session.Values["proxy_auth_domain"] = "http://" + customDomain
+			session.Values["proxy_auth_domain"] = "http://" + domain
 		}
 		session.Save(r, w)
 
@@ -289,14 +281,9 @@ func (a *Auth) checkTokenExists(session *sessions.Session, w http.ResponseWriter
 
 		session.Save(r, w)
 
-		// If we are in custom domain, redirect to pages domain to trigger authorization flow
-		if !strings.HasSuffix(r.Host, a.pagesDomain) {
-			http.Redirect(w, r, a.getProxyAddress(r, state), 302)
-		} else {
-			// Otherwise just redirect to OAuth login
-			url := fmt.Sprintf(authorizeURLTemplate, a.gitLabServer, a.clientID, a.redirectURI, state)
-			http.Redirect(w, r, url, 302)
-		}
+		// Because the pages domain might be in public suffix list, we have to
+		// redirect to pages domain to trigger authorization flow
+		http.Redirect(w, r, a.getProxyAddress(r, state), 302)
 
 		return true
 	}