Refactor acme tests and middleware for lazy domain resolution

4 files changed, 99 insertions, 90 deletions

diff --git a/internal/acme/acme.go b/internal/acme/acme.go
index 7b086eeb..c83fcc08 100644
--- a/internal/acme/acme.go
+++ b/internal/acme/acme.go
@@ -10,43 +10,29 @@ import (
 	"gitlab.com/gitlab-org/gitlab-pages/internal/logging"
 )
 
-// Middleware handles acme challenges by redirecting them to GitLab instance
-type Middleware struct {
-	GitlabURL string
-}
-
-// Domain interface represent D from domain package
-type Domain interface {
-	ServeFileHTTP(w http.ResponseWriter, r *http.Request) bool
-}
+// FallbackStrategy try to solve the acme challenge before redirecting to GitLab
+type FallbackStrategy func(http.ResponseWriter, *http.Request) bool
 
 // ServeAcmeChallenges identifies if request is acme-challenge and redirects to GitLab in that case
-func (m *Middleware) ServeAcmeChallenges(w http.ResponseWriter, r *http.Request, domain Domain) bool {
-	if m == nil {
-		return false
-	}
-
+func ServeAcmeChallenges(w http.ResponseWriter, r *http.Request, fallback FallbackStrategy, gitlabURL *url.URL) bool {
 	if !isAcmeChallenge(r.URL.Path) {
 		return false
 	}
 
-	if domain.ServeFileHTTP(w, r) {
+	if fallback(w, r) {
 		return true
 	}
 
-	return m.redirectToGitlab(w, r)
+	redirectToGitlab(w, r, gitlabURL)
+	return true
 }
 
 func isAcmeChallenge(path string) bool {
 	return strings.HasPrefix(filepath.Clean(path), "/.well-known/acme-challenge/")
 }
 
-func (m *Middleware) redirectToGitlab(w http.ResponseWriter, r *http.Request) bool {
-	redirectURL, err := url.Parse(m.GitlabURL)
-	if err != nil {
-		logging.LogRequest(r).WithError(err).Error("Can't parse GitLab URL for acme challenge redirect")
-		return false
-	}
+func redirectToGitlab(w http.ResponseWriter, r *http.Request, gitlabURL *url.URL) {
+	redirectURL := *gitlabURL
 
 	redirectURL.Path = "/-/acme-challenge"
 	query := redirectURL.Query()
@@ -57,5 +43,4 @@ func (m *Middleware) redirectToGitlab(w http.ResponseWriter, r *http.Request) bo
 	logging.LogRequest(r).WithField("redirect_url", redirectURL).Debug("Redirecting to GitLab for processing acme challenge")
 
 	http.Redirect(w, r, redirectURL.String(), http.StatusTemporaryRedirect)
-	return true
 }
diff --git a/internal/acme/acme_test.go b/internal/acme/acme_test.go
index 225d1dc4..148e953f 100644
--- a/internal/acme/acme_test.go
+++ b/internal/acme/acme_test.go
@@ -1,65 +1,64 @@
-package acme
+package acme_test
 
 import (
 	"net/http"
+	"net/url"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 
-	"gitlab.com/gitlab-org/gitlab-pages/internal/testhelpers"
+	"gitlab.com/gitlab-org/gitlab-pages/internal/acme"
 )
 
-type domainStub struct {
-	hasAcmeChallenge bool
-}
-
-func (d *domainStub) ServeFileHTTP(w http.ResponseWriter, r *http.Request) bool {
-	if r.URL.Path == "/.well-known/acme-challenge/token" {
-		return d.hasAcmeChallenge
-	}
-
-	return false
-}
-
-func serveAcmeOrNotFound(m *Middleware, domain Domain) http.HandlerFunc {
-	return func(w http.ResponseWriter, r *http.Request) {
-		if !m.ServeAcmeChallenges(w, r, domain) {
-			http.NotFound(w, r)
-		}
-	}
-}
-
 const (
 	baseURL      = "http://example.com"
 	indexURL     = baseURL + "/index.html"
 	challengeURL = baseURL + "/.well-known/acme-challenge/token"
 )
 
-var (
-	domainWithChallenge = &domainStub{hasAcmeChallenge: true}
-	domain              = &domainStub{hasAcmeChallenge: false}
-	middleware          = &Middleware{GitlabURL: "https://gitlab.example.com"}
-	middlewareMalformed = &Middleware{GitlabURL: ":foo"}
-)
-
-func TestServeAcmeChallengesNotConfigured(t *testing.T) {
-	require.HTTPStatusCode(t, serveAcmeOrNotFound(nil, domain), http.MethodGet, challengeURL, nil, http.StatusNotFound)
-}
-
-func TestServeAcmeChallengeMalformed(t *testing.T) {
-	require.HTTPStatusCode(t, serveAcmeOrNotFound(middlewareMalformed, domain), http.MethodGet, challengeURL, nil, http.StatusNotFound)
-}
+func TestAcmeMiddleware(t *testing.T) {
+	u, err := url.Parse("https://gitlab.example.com")
+	require.NoError(t, err)
 
-func TestServeAcmeChallengeWhenPresent(t *testing.T) {
-	require.HTTPStatusCode(t, serveAcmeOrNotFound(middleware, domainWithChallenge), http.MethodGet, challengeURL, nil, http.StatusOK)
-}
+	testCases := []struct {
+		name           string
+		f              acme.FallbackStrategy
+		path           string
+		expectedStatus int
+	}{
+		{
+			name:           "not an acme request",
+			path:           indexURL,
+			expectedStatus: http.StatusAccepted,
+		},
+		{
+			name: "acme challenge redirect to gitlab if missing",
+			f: func(w http.ResponseWriter, r *http.Request) bool {
+				return false
+			},
+			path:           challengeURL,
+			expectedStatus: http.StatusTemporaryRedirect,
+		},
+		{
+			name: "acme challenge served from disk if present",
+			f: func(w http.ResponseWriter, r *http.Request) bool {
+				w.WriteHeader(http.StatusOK)
+				return true
+			},
+			path:           challengeURL,
+			expectedStatus: http.StatusOK,
+		},
+	}
 
-func TestServeAcmeChallengeWhenMissing(t *testing.T) {
-	testhelpers.AssertRedirectTo(
-		t, serveAcmeOrNotFound(middleware, domain),
-		"GET", challengeURL, nil,
-		"https://gitlab.example.com/-/acme-challenge?domain=example.com&token=token",
-	)
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				if !acme.ServeAcmeChallenges(w, r, tc.f, u) {
+					w.WriteHeader(http.StatusAccepted)
+				}
+			})
 
-	require.HTTPStatusCode(t, serveAcmeOrNotFound(middleware, domain), http.MethodGet, indexURL, nil, http.StatusNotFound)
+			require.HTTPStatusCode(t, h.ServeHTTP, http.MethodGet, tc.path, nil, tc.expectedStatus)
+		})
+	}
 }
diff --git a/internal/acme/middleware.go b/internal/acme/middleware.go
deleted file mode 100644
index 6fe5cc82..00000000
--- a/internal/acme/middleware.go
+++ /dev/null
@@ -1,21 +0,0 @@
-package acme
-
-import (
-	"net/http"
-
-	// TODO: break this dependency too https://gitlab.com/gitlab-org/gitlab-pages/-/issues/650
-	domainCfg "gitlab.com/gitlab-org/gitlab-pages/internal/domain"
-)
-
-// AcmeMiddleware handles ACME challenges
-func (m *Middleware) AcmeMiddleware(handler http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		domain := domainCfg.FromRequest(r)
-
-		if m.ServeAcmeChallenges(w, r, domain) {
-			return
-		}
-
-		handler.ServeHTTP(w, r)
-	})
-}
diff --git a/internal/handlers/acme.go b/internal/handlers/acme.go
new file mode 100644
index 00000000..258aaebe
--- /dev/null
+++ b/internal/handlers/acme.go
@@ -0,0 +1,46 @@
+package handlers
+
+import (
+	"errors"
+	"net/http"
+	"net/url"
+
+	"gitlab.com/gitlab-org/gitlab-pages/internal/acme"
+	"gitlab.com/gitlab-org/gitlab-pages/internal/domain"
+	"gitlab.com/gitlab-org/gitlab-pages/internal/httperrors"
+	"gitlab.com/gitlab-org/gitlab-pages/internal/logging"
+	"gitlab.com/gitlab-org/gitlab-pages/internal/request"
+	"gitlab.com/gitlab-org/gitlab-pages/internal/source"
+)
+
+func AcmeMiddleware(handler http.Handler, s source.Source, gitlabURL string) http.Handler {
+	if gitlabURL == "" {
+		return handler
+	}
+
+	u, _ := url.Parse(gitlabURL)
+	fn := serveFromDomain(s)
+
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if acme.ServeAcmeChallenges(w, r, fn, u) {
+			return
+		}
+
+		handler.ServeHTTP(w, r)
+	})
+}
+
+func serveFromDomain(s source.Source) acme.FallbackStrategy {
+	return func(w http.ResponseWriter, r *http.Request) bool {
+		d, err := s.GetDomain(r.Context(), request.GetHostWithoutPort(r))
+
+		if err != nil && !errors.Is(err, domain.ErrDomainDoesNotExist) {
+			logging.LogRequest(r).WithError(err).Error("could not fetch domain information from a source")
+
+			httperrors.Serve502(w)
+			return true
+		}
+
+		return d.ServeFileHTTP(w, r)
+	}
+}