Get namespace domain if auth fails for a private domain

Add acceptance test and some more domains for testing

Move namespace domain serving logic

Restore go.sum

Remove redundant return

Fix linter

12 files changed, 252 insertions, 49 deletions

diff --git a/acceptance_test.go b/acceptance_test.go
index 0ba5d18f..ada77682 100644
--- a/acceptance_test.go
+++ b/acceptance_test.go
@@ -1206,6 +1206,101 @@ func TestAccessControlUnderCustomDomain(t *testing.T) {
 	require.Equal(t, http.StatusOK, authrsp.StatusCode)
 }
 
+func TestCustomErrorPageWithAuth(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()
+
+	tests := []struct {
+		name              string
+		domain            string
+		path              string
+		expectedErrorPage string
+	}{
+		{
+			name:              "private_project_authorized",
+			domain:            "group.404.gitlab-example.com",
+			path:              "/private_project/unknown",
+			expectedErrorPage: "Private custom 404 error page",
+		},
+		{
+			name:              "public_namespace_with_private_unauthorized_project",
+			domain:            "group.404.gitlab-example.com",
+			path:              "/private_unauthorized/unknown",
+			expectedErrorPage: "Custom 404 group page",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			domain := tt.domain
+			path := tt.path
+			rsp, err := GetRedirectPage(t, httpListener, domain, path)
+			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, "http://"+domain, url.Query().Get("domain"))
+
+			pagesrsp, err := GetRedirectPage(t, httpListener, url.Host, url.Path+"?"+url.RawQuery)
+			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 := GetRedirectPageWithCookie(t, httpListener, "projects.gitlab-example.com", "/auth?code=1&state="+
+				state, pagescookie)
+
+			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, domain, 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 = GetRedirectPageWithCookie(t, httpListener, domain, "/auth?code=1&state="+
+				state, cookie)
+
+			require.NoError(t, err)
+			defer authrsp.Body.Close()
+
+			// Will redirect to the page
+			groupCookie := 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 error page
+			require.Equal(t, "http://"+domain+path, url.String())
+
+			// Fetch page in custom domain
+			anotherResp, err := GetRedirectPageWithCookie(t, httpListener, domain, path, groupCookie)
+			require.NoError(t, err)
+
+			require.Equal(t, http.StatusNotFound, anotherResp.StatusCode)
+
+			page, err := ioutil.ReadAll(anotherResp.Body)
+			require.NoError(t, err)
+			require.Contains(t, string(page), tt.expectedErrorPage)
+		})
+	}
+}
+
 func TestAccessControlUnderCustomDomainWithHTTPSProxy(t *testing.T) {
 	skipUnlessEnabled(t, "not-inplace-chroot")
 
diff --git a/app.go b/app.go
index 4ead5003..a6e6006d 100644
--- a/app.go
+++ b/app.go
@@ -95,22 +95,23 @@ func (a *theApp) domain(host string) (*domain.Domain, error) {
 }
 
 func (a *theApp) checkAuthenticationIfNotExists(domain *domain.Domain, w http.ResponseWriter, r *http.Request) bool {
-	if domain == nil || !domain.HasLookupPath(r) {
+	if !domain.HasLookupPath(r) {
 		// Only if auth is supported
 		if a.Auth.IsAuthSupported() {
 			// To avoid user knowing if pages exist, we will force user to login and authorize pages
-			if a.Auth.CheckAuthenticationWithoutProject(w, r) {
+			if contentServed, authFailed := a.Auth.CheckAuthenticationWithoutProject(w, r); contentServed {
+				return true
+			} else if authFailed && domain != nil {
+				// try to serve custom namespace not found if exists and is public
+				domain.ServeNamespaceNotFound(w, r)
 				return true
 			}
 
-			// User is authenticated, show the 404
 			if domain != nil {
+				// User is authenticated, show the 404
 				domain.ServeNotFoundHTTP(w, r)
-			} else {
-				httperrors.Serve404(w)
+				return true
 			}
-
-			return true
 		}
 	}
 
@@ -250,7 +251,10 @@ func (a *theApp) accessControlMiddleware(handler http.Handler) http.Handler {
 		// Only for projects that have access control enabled
 		if domain.IsAccessControlEnabled(r) {
 			// accessControlMiddleware
-			if a.Auth.CheckAuthentication(w, r, domain) {
+			if contentServed, authFailed := a.Auth.CheckAuthentication(w, r, domain.GetProjectID(r)); contentServed {
+				return
+			} else if authFailed && domain != nil {
+				domain.ServeNamespaceNotFound(w, r)
 				return
 			}
 		}
@@ -274,12 +278,12 @@ func (a *theApp) serveFileOrNotFoundHandler() http.Handler {
 			// because the projects override the paths of the namespace project and they might be private even though
 			// namespace project is public.
 			if domain.IsNamespaceProject(r) {
-				if a.Auth.CheckAuthenticationWithoutProject(w, r) {
+				if contentServed, authFailed := a.Auth.CheckAuthenticationWithoutProject(w, r); contentServed {
+					return
+				} else if authFailed {
+					httperrors.Serve404(w)
 					return
 				}
-
-				domain.ServeNotFoundHTTP(w, r)
-				return
 			}
 
 			domain.ServeNotFoundHTTP(w, r)
diff --git a/go.sum b/go.sum
index ceb73a62..288a508a 100644
--- a/go.sum
+++ b/go.sum
@@ -327,6 +327,7 @@ github.com/yalp/jsonpath v0.0.0-20180802001716-5cc68e5049a0/go.mod h1:/LWChgwKmv
 github.com/yudai/gojsondiff v1.0.0/go.mod h1:AY32+k2cwILAkW1fbgxQ5mUmMiZFgLIV+FBNExI05xg=
 github.com/yudai/golcs v0.0.0-20170316035057-ecda9a501e82/go.mod h1:lgjkn3NuSvDfVJdfcVVdX+jpBxNmX4rDAzaS45IcYoM=
 github.com/yudai/pp v2.0.1+incompatible/go.mod h1:PuxR/8QJ7cyCkFp/aUDS+JY727OFEZkTdatxwunjIkc=
+github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 gitlab.com/gitlab-org/labkit v0.0.0-20200526151107-6dcf1319fcd0 h1:+TiSK1umKrr6PCYCR8rJmUMk39hNynOgjZIbJmYXHm0=
 gitlab.com/gitlab-org/labkit v0.0.0-20200526151107-6dcf1319fcd0/go.mod h1:SNfxkfUwVNECgtmluVayv0GWFgEjjBs5AzgsowPQuo0=
 gitlab.com/lupine/go-mimedb v0.0.0-20180307000149-e8af1d659877 h1:k5N2m0IPaMuwWmFTO9fyTK4IEnSm35GC/p1S7VRgUyM=
diff --git a/internal/auth/auth.go b/internal/auth/auth.go
index c12207ca..768e7f75 100644
--- a/internal/auth/auth.go
+++ b/internal/auth/auth.go
@@ -18,7 +18,6 @@ import (
 	log "github.com/sirupsen/logrus"
 	"gitlab.com/gitlab-org/labkit/errortracking"
 
-	"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"
@@ -437,10 +436,10 @@ func (a *Auth) IsAuthSupported() bool {
 	return a != nil
 }
 
-func (a *Auth) checkAuthentication(w http.ResponseWriter, r *http.Request, projectID uint64) bool {
+func (a *Auth) checkAuthentication(w http.ResponseWriter, r *http.Request, projectID uint64) (contentServed, authFailed bool) {
 	session := a.checkSessionIsValid(w, r)
 	if session == nil {
-		return true
+		return true, false
 	}
 
 	// Access token exists, authorize request
@@ -457,14 +456,14 @@ func (a *Auth) checkAuthentication(w http.ResponseWriter, r *http.Request, proje
 		errortracking.Capture(err, errortracking.WithRequest(req))
 
 		httperrors.Serve500(w)
-		return true
+		return true, false
 	}
 
 	req.Header.Add("Authorization", "Bearer "+session.Values["access_token"].(string))
 	resp, err := a.apiClient.Do(req)
 
 	if err == nil && checkResponseForInvalidToken(resp, session, w, r) {
-		return true
+		return true, false
 	}
 
 	if err != nil || resp.StatusCode != 200 {
@@ -472,17 +471,17 @@ func (a *Auth) checkAuthentication(w http.ResponseWriter, r *http.Request, proje
 			logRequest(r).WithError(err).Error("Failed to retrieve info with token")
 		}
 
-		return false
+		return false, true
 	}
 
-	return false
+	return false, false
 }
 
 // CheckAuthenticationWithoutProject checks if user is authenticated and has a valid token
-func (a *Auth) CheckAuthenticationWithoutProject(w http.ResponseWriter, r *http.Request) bool {
+func (a *Auth) CheckAuthenticationWithoutProject(w http.ResponseWriter, r *http.Request) (contentServed, authFailed bool) {
 	if a == nil {
 		// No auth supported
-		return false
+		return false, false
 	}
 
 	return a.checkAuthentication(w, r, 0)
@@ -512,7 +511,8 @@ func (a *Auth) RequireAuth(w http.ResponseWriter, r *http.Request) bool {
 }
 
 // CheckAuthentication checks if user is authenticated and has access to the project
-func (a *Auth) CheckAuthentication(w http.ResponseWriter, r *http.Request, domain *domain.Domain) bool {
+// will return contentServed = false when authFailed = true
+func (a *Auth) CheckAuthentication(w http.ResponseWriter, r *http.Request, projectID uint64) (contentServed, authFailed bool) {
 	logRequest(r).Debug("Authenticate request")
 
 	if a == nil {
@@ -520,29 +520,10 @@ func (a *Auth) CheckAuthentication(w http.ResponseWriter, r *http.Request, domai
 		errortracking.Capture(errAuthNotConfigured, errortracking.WithRequest(r))
 
 		httperrors.Serve500(w)
-		return true
-	}
-
-	if a.checkAuthentication(w, r, domain.GetProjectID(r)) {
-		// if auth fails, try to resolve parent namespace domain
-		r.URL.Path = "/"
-		parent, err := domain.Resolver.Resolve(r)
-		if err != nil {
-			httperrors.Serve404(w)
-			return true
-		}
-
-		// for namespace domains that have no access control enabled
-		if parent.LookupPath.IsNamespaceProject && !parent.LookupPath.HasAccessControl {
-			parent.ServeNotFoundHTTP(w, r)
-			return true
-		}
-
-		httperrors.Serve404(w)
-		return true
+		return true, false
 	}
 
-	return false
+	return a.checkAuthentication(w, r, projectID)
 }
 
 // CheckResponseForInvalidToken checks response for invalid token and destroys session if it was invalid
diff --git a/internal/auth/auth_test.go b/internal/auth/auth_test.go
index fc8ddb44..711dea78 100644
--- a/internal/auth/auth_test.go
+++ b/internal/auth/auth_test.go
@@ -181,7 +181,10 @@ func TestCheckAuthenticationWhenAccess(t *testing.T) {
 	session.Values["access_token"] = "abc"
 	session.Save(r, result)
 
-	require.Equal(t, false, auth.CheckAuthentication(result, r, 1000))
+	contentServed, authFailed := auth.CheckAuthentication(result, r, 1000)
+	require.False(t, contentServed)
+	require.False(t, authFailed)
+	// content wasn't served so the default response from CheckAuthentication should be 200
 	require.Equal(t, 200, result.Code)
 }
 
@@ -219,8 +222,11 @@ func TestCheckAuthenticationWhenNoAccess(t *testing.T) {
 	session.Values["access_token"] = "abc"
 	session.Save(r, result)
 
-	require.Equal(t, true, auth.CheckAuthentication(result, r, 1000))
-	require.Equal(t, 404, result.Code)
+	contentServed, authFailed := auth.CheckAuthentication(result, r, 1000)
+	require.False(t, contentServed)
+	require.True(t, authFailed)
+	// content wasn't served so the default response from CheckAuthentication should be 200
+	require.Equal(t, 200, result.Code)
 }
 
 func TestCheckAuthenticationWhenInvalidToken(t *testing.T) {
@@ -257,7 +263,9 @@ func TestCheckAuthenticationWhenInvalidToken(t *testing.T) {
 	session.Values["access_token"] = "abc"
 	session.Save(r, result)
 
-	require.Equal(t, true, auth.CheckAuthentication(result, r, 1000))
+	contentServed, authFailed := auth.CheckAuthentication(result, r, 1000)
+	require.True(t, contentServed)
+	require.False(t, authFailed)
 	require.Equal(t, 302, result.Code)
 }
 
@@ -295,7 +303,9 @@ func TestCheckAuthenticationWithoutProject(t *testing.T) {
 	session.Values["access_token"] = "abc"
 	session.Save(r, result)
 
-	require.Equal(t, false, auth.CheckAuthenticationWithoutProject(result, r))
+	contentServed, authFailed := auth.CheckAuthenticationWithoutProject(result, r)
+	require.False(t, contentServed)
+	require.False(t, authFailed)
 	require.Equal(t, 200, result.Code)
 }
 
@@ -332,7 +342,9 @@ func TestCheckAuthenticationWithoutProjectWhenInvalidToken(t *testing.T) {
 	session.Values["access_token"] = "abc"
 	session.Save(r, result)
 
-	require.Equal(t, true, auth.CheckAuthenticationWithoutProject(result, r))
+	contentServed, authFailed := auth.CheckAuthenticationWithoutProject(result, r)
+	require.True(t, contentServed)
+	require.False(t, authFailed)
 	require.Equal(t, 302, result.Code)
 }
 
diff --git a/internal/domain/domain.go b/internal/domain/domain.go
index 17a0e1d3..2a301121 100644
--- a/internal/domain/domain.go
+++ b/internal/domain/domain.go
@@ -168,3 +168,24 @@ func (d *Domain) ServeNotFoundHTTP(w http.ResponseWriter, r *http.Request) {
 
 	request.ServeNotFoundHTTP(w, r)
 }
+
+// ServeNamespaceNotFound will try to find a parent namespace domain for a request
+// that failed authentication so that we serve the custom namespace error page for
+// public namespace domains
+func (d *Domain) ServeNamespaceNotFound(w http.ResponseWriter, r *http.Request) {
+	// override the path nd try to resolve the domain name
+	r.URL.Path = "/"
+	namespaceDomain, err := d.Resolver.Resolve(r)
+	if err != nil {
+		httperrors.Serve404(w)
+		return
+	}
+
+	// for namespace domains that have no access control enabled
+	if namespaceDomain.LookupPath.IsNamespaceProject && !namespaceDomain.LookupPath.HasAccessControl {
+		namespaceDomain.ServeNotFoundHTTP(w, r)
+		return
+	}
+
+	httperrors.Serve404(w)
+}
diff --git a/internal/domain/domain_test.go b/internal/domain/domain_test.go
index 49d46cb3..e152b031 100644
--- a/internal/domain/domain_test.go
+++ b/internal/domain/domain_test.go
@@ -1,7 +1,10 @@
 package domain
 
 import (
+	"fmt"
+	"io/ioutil"
 	"net/http"
+	"net/http/httptest"
 	"os"
 	"testing"
 
@@ -153,3 +156,76 @@ func chdirInPath(t require.TestingT, path string) func() {
 		chdirSet = false
 	}
 }
+
+func TestDomain_ServeNamespaceNotFound(t *testing.T) {
+	// defaultNotFound := "The page you're looking for could not be found."
+	// customNotFound := "Custom error page"
+
+	tests := []struct {
+		name             string
+		domain           string
+		path             string
+		resolver         *stubbedResolver
+		expectedResponse string
+	}{
+		{
+			name:   "public_namespace_domain",
+			domain: "group.404.gitlab-example.com",
+			path:   "/unknown",
+			resolver: &stubbedResolver{
+				project: &serving.LookupPath{
+					Path:               "../../shared/pages/group.404/group.404.gitlab-example.com/public",
+					IsNamespaceProject: true,
+				},
+				subpath: "/unknown",
+			},
+			expectedResponse: "Custom 404 group page",
+		},
+		{
+			name:   "private_project_under_public_namespace_domain",
+			domain: "group.404.gitlab-example.com",
+			path:   "/private_project/unknown",
+			resolver: &stubbedResolver{
+				project: &serving.LookupPath{
+					Path:               "../../shared/pages/group.404/group.404.gitlab-example.com/public",
+					IsNamespaceProject: true,
+					HasAccessControl:   false,
+				},
+				subpath: "/",
+			},
+			expectedResponse: "Custom 404 group page",
+		},
+		{
+			name:   "private_namespace_domain",
+			domain: "group.404.gitlab-example.com",
+			path:   "/unknown",
+			resolver: &stubbedResolver{
+				project: &serving.LookupPath{
+					Path:               "../../shared/pages/group.404/group.404.gitlab-example.com/public",
+					IsNamespaceProject: true,
+					HasAccessControl:   true,
+				},
+				subpath: "/",
+			},
+			expectedResponse: "The page you're looking for could not be found.",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			d := &Domain{
+				Name:     tt.domain,
+				Resolver: tt.resolver,
+			}
+			w := httptest.NewRecorder()
+			r := httptest.NewRequest("GET", fmt.Sprintf("http://%s%s", tt.domain, tt.path), nil)
+			d.ServeNamespaceNotFound(w, r)
+
+			resp := w.Result()
+			require.Equal(t, http.StatusNotFound, resp.StatusCode)
+			body, err := ioutil.ReadAll(resp.Body)
+			require.NoError(t, err)
+
+			require.Contains(t, string(body), tt.expectedResponse)
+		})
+	}
+}
diff --git a/internal/source/disk/map_test.go b/internal/source/disk/map_test.go
index d904e6aa..d2162883 100644
--- a/internal/source/disk/map_test.go
+++ b/internal/source/disk/map_test.go
@@ -49,6 +49,7 @@ func TestReadProjects(t *testing.T) {
 		"group.acme.test.io",
 		"withacmechallenge.domain.com",
 		"capitalgroup.test.io",
+		"group.404.gitlab-example.com",
 	}
 
 	for _, expected := range domains {
diff --git a/shared/pages/group.404/private_project/config.json b/shared/pages/group.404/private_project/config.json
new file mode 100644
index 00000000..5c0ebb50
--- /dev/null
+++ b/shared/pages/group.404/private_project/config.json
@@ -0,0 +1,5 @@
+{ "domains": [
+  {
+    "Domain": "group.404.gitlab-example.com"
+  }
+], "id": 1000, "access_control": true }
diff --git a/shared/pages/group.404/private_project/public/404.html b/shared/pages/group.404/private_project/public/404.html
new file mode 100644
index 00000000..6993ae1a
--- /dev/null
+++ b/shared/pages/group.404/private_project/public/404.html
@@ -0,0 +1 @@
+Private custom 404 error page
diff --git a/shared/pages/group.404/private_unauthorized/config.json b/shared/pages/group.404/private_unauthorized/config.json
new file mode 100644
index 00000000..79349565
--- /dev/null
+++ b/shared/pages/group.404/private_unauthorized/config.json
@@ -0,0 +1,5 @@
+{ "domains": [
+  {
+    "Domain": "group.404.gitlab-example.com"
+  }
+], "id": 2000, "access_control": true }
diff --git a/shared/pages/group.404/private_unauthorized/public/404.html b/shared/pages/group.404/private_unauthorized/public/404.html
new file mode 100644
index 00000000..6993ae1a
--- /dev/null
+++ b/shared/pages/group.404/private_unauthorized/public/404.html
@@ -0,0 +1 @@
+Private custom 404 error page