Add support for previewing artifacts that are not public

Remove some duplicate logic on Auth module
Separate handling artifact to own handlers package
Unit test handlers by mocking auth and artifact modules
Add generate-mock step to Makefile
Use additional handler func to simplify TryMakeRequest return type
Always try with token if exists
Do not log RequestURI, log path only
Remove not used logRequest func

2 files changed, 83 insertions, 29 deletions

diff --git a/internal/artifact/artifact.go b/internal/artifact/artifact.go
index 5050b426..28a595bb 100644
--- a/internal/artifact/artifact.go
+++ b/internal/artifact/artifact.go
@@ -1,6 +1,7 @@
 package artifact
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -11,8 +12,11 @@ import (
 	"strings"
 	"time"
 
+	"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/logging"
 )
 
 const (
@@ -26,7 +30,10 @@ const (
 
 var (
 	// Captures subgroup + project, job ID and artifacts path
-	pathExtractor = regexp.MustCompile(`(?i)\A/-/(.*)/-/jobs/(\d+)/artifacts(/[^?]*)\z`)
+	pathExtractor            = regexp.MustCompile(`(?i)\A/-/(.*)/-/jobs/(\d+)/artifacts(/[^?]*)\z`)
+	errCreateArtifactRequest = errors.New("Failed to create the artifact request")
+	errArtifactRequest       = errors.New("Failed to request the artifact")
+	errArtifactResponse      = errors.New("Artifact request response was not successful")
 )
 
 // Artifact proxies requests for artifact files to the GitLab artifacts API
@@ -51,8 +58,9 @@ func New(server string, timeoutSeconds int, pagesDomain string) *Artifact {
 
 // TryMakeRequest will attempt to proxy a request and write it to the argument
 // http.ResponseWriter, ultimately returning a bool that indicates if the
-// http.ResponseWriter has been written to in any capacity.
-func (a *Artifact) TryMakeRequest(host string, w http.ResponseWriter, r *http.Request) bool {
+// http.ResponseWriter has been written to in any capacity. Additional handler func
+// may be given which should return true if it did handle the response.
+func (a *Artifact) TryMakeRequest(host string, w http.ResponseWriter, r *http.Request, token string, additionalHandler func(*http.Response) bool) bool {
 	if a == nil || a.server == "" || host == "" {
 		return false
 	}
@@ -62,30 +70,51 @@ func (a *Artifact) TryMakeRequest(host string, w http.ResponseWriter, r *http.Re
 		return false
 	}
 
-	a.makeRequest(w, reqURL)
+	a.makeRequest(w, r, reqURL, token, additionalHandler)
+
 	return true
 }
 
-func (a *Artifact) makeRequest(w http.ResponseWriter, reqURL *url.URL) {
-	resp, err := a.client.Get(reqURL.String())
+func (a *Artifact) makeRequest(w http.ResponseWriter, r *http.Request, reqURL *url.URL, token string, additionalHandler func(*http.Response) bool) {
+	req, err := http.NewRequest("GET", reqURL.String(), nil)
+	if err != nil {
+		logging.LogRequest(r).WithError(err).Error(errCreateArtifactRequest)
+		errortracking.Capture(err, errortracking.WithRequest(r))
+		httperrors.Serve500(w)
+		return
+	}
+
+	if token != "" {
+		req.Header.Add("Authorization", "Bearer "+token)
+	}
+	resp, err := a.client.Do(req)
+
 	if err != nil {
+		logging.LogRequest(r).WithError(err).Error(errArtifactRequest)
+		errortracking.Capture(err, errortracking.WithRequest(r))
 		httperrors.Serve502(w)
 		return
 	}
 
+	if additionalHandler(resp) {
+		return
+	}
+
 	if resp.StatusCode == http.StatusNotFound {
 		httperrors.Serve404(w)
 		return
 	}
 
 	if resp.StatusCode == http.StatusInternalServerError {
+		logging.LogRequest(r).Error(errArtifactResponse)
+		errortracking.Capture(errArtifactResponse, errortracking.WithRequest(r))
 		httperrors.Serve500(w)
 		return
 	}
 
-	// we only cache responses within the 2xx series response codes
-	if (resp.StatusCode >= minStatusCode) && (resp.StatusCode <= maxStatusCode) {
-		w.Header().Set("Cache-Control", "max-age=3600")
+	// we only cache responses within the 2xx series response codes and that were not private
+	if token == "" {
+		addCacheHeader(w, resp)
 	}
 
 	w.Header().Set("Content-Type", resp.Header.Get("Content-Type"))
@@ -95,6 +124,12 @@ func (a *Artifact) makeRequest(w http.ResponseWriter, reqURL *url.URL) {
 	return
 }
 
+func addCacheHeader(w http.ResponseWriter, resp *http.Response) {
+	if (resp.StatusCode >= minStatusCode) && (resp.StatusCode <= maxStatusCode) {
+		w.Header().Set("Cache-Control", "max-age=3600")
+	}
+}
+
 // BuildURL returns a pointer to a url.URL for where the request should be
 // proxied to. The returned bool will indicate if there is some sort of issue
 // with the url while it is being generated.
diff --git a/internal/artifact/artifact_test.go b/internal/artifact/artifact_test.go
index ef37384e..689effb5 100644
--- a/internal/artifact/artifact_test.go
+++ b/internal/artifact/artifact_test.go
@@ -15,29 +15,12 @@ import (
 func TestTryMakeRequest(t *testing.T) {
 	content := "<!DOCTYPE html><html><head><title>Title of the document</title></head><body></body></html>"
 	contentType := "text/html; charset=utf-8"
-	testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Content-Type", contentType)
-		switch r.URL.RawPath {
-		case "/projects/group%2Fsubgroup%2Fproject/jobs/1/artifacts/200.html":
-			w.WriteHeader(http.StatusOK)
-		case "/projects/group%2Fsubgroup%2Fproject/jobs/1/artifacts/max-caching.html":
-			w.WriteHeader(http.StatusIMUsed)
-		case "/projects/group%2Fsubgroup%2Fproject/jobs/1/artifacts/non-caching.html":
-			w.WriteHeader(http.StatusTeapot)
-		case "/projects/group%2Fsubgroup%2Fproject/jobs/1/artifacts/500.html":
-			w.WriteHeader(http.StatusInternalServerError)
-		case "/projects/group%2Fsubgroup%2Fgroup%2Fproject/jobs/1/artifacts/404.html":
-			w.WriteHeader(http.StatusNotFound)
-		default:
-			t.Log("Surprising r.URL.RawPath", r.URL.RawPath)
-			w.WriteHeader(999)
-		}
-		fmt.Fprint(w, content)
-	}))
+	testServer := makeArtifactServerStub(t, content, contentType)
 	defer testServer.Close()
 
 	cases := []struct {
 		Path         string
+		Token        string
 		Status       int
 		Content      string
 		Length       string
@@ -47,6 +30,7 @@ func TestTryMakeRequest(t *testing.T) {
 	}{
 		{
 			"/200.html",
+			"",
 			http.StatusOK,
 			content,
 			"90",
@@ -55,7 +39,18 @@ func TestTryMakeRequest(t *testing.T) {
 			"basic successful request",
 		},
 		{
+			"/200.html",
+			"token",
+			http.StatusOK,
+			content,
+			"90",
+			"",
+			"text/html; charset=utf-8",
+			"basic successful request",
+		},
+		{
 			"/max-caching.html",
+			"",
 			http.StatusIMUsed,
 			content,
 			"90",
@@ -65,6 +60,7 @@ func TestTryMakeRequest(t *testing.T) {
 		},
 		{
 			"/non-caching.html",
+			"",
 			http.StatusTeapot,
 			content,
 			"90",
@@ -82,7 +78,7 @@ func TestTryMakeRequest(t *testing.T) {
 			r := &http.Request{URL: reqURL}
 			art := artifact.New(testServer.URL, 1, "gitlab-example.io")
 
-			require.True(t, art.TryMakeRequest("group.gitlab-example.io", result, r))
+			require.True(t, art.TryMakeRequest("group.gitlab-example.io", result, r, c.Token, func(resp *http.Response) bool { return false }))
 			require.Equal(t, c.Status, result.Code)
 			require.Equal(t, c.ContentType, result.Header().Get("Content-Type"))
 			require.Equal(t, c.Length, result.Header().Get("Content-Length"))
@@ -93,6 +89,29 @@ func TestTryMakeRequest(t *testing.T) {
 	}
 }
 
+// provide stub for testing different artifact responses
+func makeArtifactServerStub(t *testing.T, content string, contentType string) *httptest.Server {
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", contentType)
+		switch r.URL.RawPath {
+		case "/projects/group%2Fsubgroup%2Fproject/jobs/1/artifacts/200.html":
+			w.WriteHeader(http.StatusOK)
+		case "/projects/group%2Fsubgroup%2Fproject/jobs/1/artifacts/max-caching.html":
+			w.WriteHeader(http.StatusIMUsed)
+		case "/projects/group%2Fsubgroup%2Fproject/jobs/1/artifacts/non-caching.html":
+			w.WriteHeader(http.StatusTeapot)
+		case "/projects/group%2Fsubgroup%2Fproject/jobs/1/artifacts/500.html":
+			w.WriteHeader(http.StatusInternalServerError)
+		case "/projects/group%2Fsubgroup%2Fgroup%2Fproject/jobs/1/artifacts/404.html":
+			w.WriteHeader(http.StatusNotFound)
+		default:
+			t.Log("Surprising r.URL.RawPath", r.URL.RawPath)
+			w.WriteHeader(999)
+		}
+		fmt.Fprint(w, content)
+	}))
+}
+
 func TestBuildURL(t *testing.T) {
 	cases := []struct {
 		RawServer   string