add header-name flag, set header-name to default auth if set, add tests for header-name and cookie configs

Author: whwalter

internal/cli/cli.go

@@ -35,6 +35,7 @@ func NewRootCommand() *cobra.Command {
 	cmd.PersistentFlags().String("grafana-proxy-url", "http://grafana.example.com", "Grafana url to proxy to")
 	cmd.PersistentFlags().String("grafana-user-header", "X-WEBAUTH-USER", "Header to containing the user to authenticate")
 	cmd.PersistentFlags().String("cookie-name", "auth_token", "Cookie name with jwt token. If set will take precedence over auth header")
+	cmd.PersistentFlags().String("header-name", "", "header name with jwt token. If set will take precedence over cookie-name")
 	cmd.PersistentFlags().String("admin-user", "admin", "Admin user")
 	cmd.PersistentFlags().String("admin-password", "", "Admin password")
 	cmd.PersistentFlags().String("jwt-claim-login", "email", "JWT claim to be used as user Login in Grafana. Valid values are 'email' or 'sub'")
@@ -84,6 +85,7 @@ func defaultServerOptions() []server.ServerFuncOpt {
 
 	opts := []server.ServerFuncOpt{
 		server.WithCookieName(viper.GetString("cookie-name")),
+		server.WithHeaderName(viper.GetString("header-name")),
 		server.WithGrafanaResponseHeaders(responseHeaders),
 	}
 

internal/server/options.go

@@ -14,6 +14,13 @@ func WithCookieName(cookie string) ServerFuncOpt {
 	}
 }
 
+func WithHeaderName(header string) ServerFuncOpt {
+	return func(s *Server) error {
+		s.headerName = header
+		return nil
+	}
+}
+
 func WithConfigGroups(groups config.Groups) ServerFuncOpt {
 	return func(s *Server) error {
 		s.groups = groups

internal/server/server.go

@@ -26,6 +26,7 @@ type GrafanaClaimsConfig struct {
 type Server struct {
 	router                 *http.ServeMux
 	cookieName             string
+	headerName             string
 	groups                 config.Groups
 	grafanaProxyUrl        *url.URL
 	grafanaClient          *grafana.Client
@@ -65,15 +66,27 @@ func getValidClaim(claims *jwt.Claims, input string) string {
 
 func (s *Server) handleRoot() http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
-		// Extract token from cookie
-		cookie, err := r.Cookie(s.cookieName)
-		if err != nil {
-			logAndError(w, http.StatusUnauthorized, err, "error reading cookie")
-			return
+		// Get token
+		var token string
+
+		if s.headerName != "" {
+			token = r.Header.Get(s.headerName)
+			// replicate the cookie look up behavior for a missing header
+			if token == "" {
+				logAndError(w, http.StatusUnauthorized, fmt.Errorf("No value for header %s", s.headerName), "error reading header")
+				return
+			}
+		} else {
+			cookie, err := r.Cookie(s.cookieName)
+			if err != nil {
+				logAndError(w, http.StatusUnauthorized, err, "error reading cookie")
+				return
+			}
+			token = cookie.Value
 		}
 
 		// Get claims from token
-		claims, err := jwt.TokenClaims(cookie.Value)
+		claims, err := jwt.TokenClaims(token)
 		if err != nil {
 			logAndError(w, http.StatusUnauthorized, err, "error reading claims from jwt token")
 			return

internal/server/server_test.go

@@ -37,33 +37,41 @@ func TestTokenValidations(t *testing.T) {
 	defer backendServer.Close()
 	backendURL, _ := url.Parse(backendServer.URL)
 
+	headerName := "X-Test-Header"
+
+	validToken := newTestJWTToken("jhon")
+	noSubToken := newTestJWTToken("")
+	invalidToken := "this-is-no-valid-jwt"
+	emptyToken := ""
+
 	client := grafana.NewMockClient(gapi.User{Login: "jhon", ID: 1}, map[int64]grafana.RoleType{})
 
 	tests := []struct {
 		name       string
 		cookie     *http.Cookie
+		header     *string
 		authorized bool
 	}{
 		{
 			name: "valid JWT token and cookie",
 			cookie: &http.Cookie{
 				Name:  "auth_token",
-				Value: newTestJWTToken("jhon"),
+				Value: validToken,
 			},
 			authorized: true,
 		},
 		{
 			name: "valid JWT token without sub",
 			cookie: &http.Cookie{
 				Name:  "auth_token",
-				Value: newTestJWTToken(""),
+				Value: noSubToken,
 			},
 		},
 		{
 			name: "invalid JWT token",
 			cookie: &http.Cookie{
 				Name:  "auth_token",
-				Value: "this-is-no-valid-jwt",
+				Value: invalidToken,
 			},
 		},
 		{
@@ -72,34 +80,74 @@ func TestTokenValidations(t *testing.T) {
 				Name: "",
 			},
 		},
+		{
+			name:       "valid JWT token and header",
+			header:     &validToken,
+			authorized: true,
+		},
+		{
+			name:   "valid JWT token without sub",
+			header: &noSubToken,
+		},
+		{
+			name:   "invalid JWT token",
+			header: &invalidToken,
+		},
+		{
+			name:   "Empty Header",
+			header: &emptyToken,
+		},
+		{
+			name:   "valid JWT token header and cookie",
+			header: &validToken,
+			cookie: &http.Cookie{
+				Name:  "auth_token",
+				Value: validToken,
+			},
+			authorized: true,
+		},
+		{
+			name:   "valid JWT token header invalid cookie",
+			header: &validToken,
+			cookie: &http.Cookie{
+				Name:  "auth_token",
+				Value: invalidToken,
+			},
+			authorized: true,
+		},
 	}
 
 	for _, test := range tests {
 		req := httptest.NewRequest("GET", "/", nil)
 		req.Host = "http://grafana.example.com"
 
-		if test.cookie.Name != "" {
-			req.AddCookie(test.cookie)
-		}
-
-		server, err := New(
+		opts := []ServerFuncOpt{
 			WithGrafanaProxyURL(backendURL),
-			WithCookieName(test.cookie.Name),
 			WithConfigGroups(config.Groups{}),
 			WithGrafanaClient(client),
 			WithGrafanaResponseHeaders(GrafanaResponseHeaders{
 				User: "X-WEBAUTH-USER",
 			}),
-		)
+		}
+
+		if test.cookie != nil && test.cookie.Name != "" {
+			req.AddCookie(test.cookie)
+			opts = append(opts, WithCookieName(test.cookie.Name))
+		}
+		if test.header != nil {
+			req.Header.Add(headerName, *test.header)
+			opts = append(opts, WithHeaderName(headerName))
+		}
+		server, err := New(opts...)
 		assert.NoError(t, err)
 
 		w := httptest.NewRecorder()
 		server.ServeHTTP(w, req)
 
 		if test.authorized {
-			assert.Equal(t, http.StatusOK, w.Code)
+			assert.Equal(t, http.StatusOK, w.Code, test.name)
 		} else {
-			assert.Equal(t, http.StatusUnauthorized, w.Code)
+			assert.Equal(t, http.StatusUnauthorized, w.Code, test.name)
 		}
 	}
 }