1 files changed, 131 insertions, 0 deletions
diff --git a/workhorse/authorization_test.go b/workhorse/authorization_test.go
new file mode 100644
index 00000000000..8be1fdcf73a
--- /dev/null
+++ b/workhorse/authorization_test.go
@@ -0,0 +1,131 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"regexp"
+	"testing"
+
+	"gitlab.com/gitlab-org/labkit/correlation"
+
+	"github.com/dgrijalva/jwt-go"
+	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/api"
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/helper"
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/secret"
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/testhelper"
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/upstream/roundtripper"
+)
+
+func okHandler(w http.ResponseWriter, _ *http.Request, _ *api.Response) {
+	w.WriteHeader(201)
+	fmt.Fprint(w, "{\"status\":\"ok\"}")
+}
+
+func runPreAuthorizeHandler(t *testing.T, ts *httptest.Server, suffix string, url *regexp.Regexp, apiResponse interface{}, returnCode, expectedCode int) *httptest.ResponseRecorder {
+	if ts == nil {
+		ts = testAuthServer(t, url, nil, returnCode, apiResponse)
+		defer ts.Close()
+	}
+
+	// Create http request
+	ctx := correlation.ContextWithCorrelation(context.Background(), "12345678")
+	httpRequest, err := http.NewRequestWithContext(ctx, "GET", "/address", nil)
+	require.NoError(t, err)
+	parsedURL := helper.URLMustParse(ts.URL)
+	testhelper.ConfigureSecret()
+	a := api.NewAPI(parsedURL, "123", roundtripper.NewTestBackendRoundTripper(parsedURL))
+
+	response := httptest.NewRecorder()
+	a.PreAuthorizeHandler(okHandler, suffix).ServeHTTP(response, httpRequest)
+	require.Equal(t, expectedCode, response.Code)
+	return response
+}
+
+func TestPreAuthorizeHappyPath(t *testing.T) {
+	runPreAuthorizeHandler(
+		t, nil, "/authorize",
+		regexp.MustCompile(`/authorize\z`),
+		&api.Response{},
+		200, 201)
+}
+
+func TestPreAuthorizeSuffix(t *testing.T) {
+	runPreAuthorizeHandler(
+		t, nil, "/different-authorize",
+		regexp.MustCompile(`/authorize\z`),
+		&api.Response{},
+		200, 404)
+}
+
+func TestPreAuthorizeJsonFailure(t *testing.T) {
+	runPreAuthorizeHandler(
+		t, nil, "/authorize",
+		regexp.MustCompile(`/authorize\z`),
+		"not-json",
+		200, 500)
+}
+
+func TestPreAuthorizeContentTypeFailure(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, err := w.Write([]byte(`{"hello":"world"}`))
+		require.NoError(t, err, "write auth response")
+	}))
+	defer ts.Close()
+
+	runPreAuthorizeHandler(
+		t, ts, "/authorize",
+		regexp.MustCompile(`/authorize\z`),
+		"",
+		200, 200)
+}
+
+func TestPreAuthorizeRedirect(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		http.Redirect(w, r, "/", http.StatusMovedPermanently)
+	}))
+	defer ts.Close()
+
+	runPreAuthorizeHandler(t, ts, "/willredirect",
+		regexp.MustCompile(`/willredirect\z`),
+		"",
+		301, 301)
+}
+
+func TestPreAuthorizeJWT(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		token, err := jwt.Parse(r.Header.Get(secret.RequestHeader), func(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"])
+			}
+			testhelper.ConfigureSecret()
+			secretBytes, err := secret.Bytes()
+			if err != nil {
+				return nil, fmt.Errorf("read secret from file: %v", err)
+			}
+
+			return secretBytes, nil
+		})
+		require.NoError(t, err, "decode token")
+
+		claims, ok := token.Claims.(jwt.MapClaims)
+		require.True(t, ok, "claims cast")
+		require.True(t, token.Valid, "JWT token valid")
+		require.Equal(t, "gitlab-workhorse", claims["iss"], "JWT token issuer")
+
+		w.Header().Set("Content-Type", api.ResponseContentType)
+		_, err = w.Write([]byte(`{"hello":"world"}`))
+		require.NoError(t, err, "write auth response")
+	}))
+	defer ts.Close()
+
+	runPreAuthorizeHandler(
+		t, ts, "/authorize",
+		regexp.MustCompile(`/authorize\z`),
+		"",
+		200, 201)
+}