update getClaimsFromHeader

Author: duwenxin99

internal/auth/generic/generic.go

@@ -121,16 +121,14 @@ func (a AuthService) GetName() string {
 
 // Verifies generic JWT access token inside the Authorization header
 func (a AuthService) GetClaimsFromHeader(ctx context.Context, h http.Header) (map[string]any, error) {
-	authHeader := h.Get("Authorization")
-	if authHeader == "" {
-		return nil, nil // Return nil, nil if no authorization header is found
+	if a.McpEnabled {
+		return nil, nil
 	}
 
-	parts := strings.Split(authHeader, " ")
-	if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" {
-		return nil, fmt.Errorf("authorization header format must be Bearer {token}")
+	tokenString := h.Get(a.Name + "_token")
+	if tokenString == "" {
+		return nil, nil
 	}
-	tokenString := parts[1]
 
 	// Parse and verify the token signature
 	token, err := jwt.Parse(tokenString, a.kf.Keyfunc)
@@ -161,45 +159,10 @@ func (a AuthService) GetClaimsFromHeader(ctx context.Context, h http.Header) (ma
 		}
 	}
 
-	// Some IDPs use 'client_id' instead of 'aud' or put it as a single string, checking that if aud not found or not matched
-	if !isAudValid {
-		if clientIDClaim, ok := claims["client_id"].(string); ok && clientIDClaim == a.Audience {
-			isAudValid = true
-		} else if audStr, ok := claims["aud"].(string); ok && audStr == a.Audience {
-			isAudValid = true
-		}
-	}
-
 	if !isAudValid {
 		return nil, fmt.Errorf("audience validation failed: expected %s, got %v", a.Audience, aud)
 	}
 
-	// Validate 'scope' claim against ScopesRequired
-	if len(a.ScopesRequired) > 0 {
-		var tokenScopes []string
-
-		switch s := claims["scope"].(type) {
-		case string:
-			tokenScopes = strings.Split(s, " ") // space-separated string is common
-		case []interface{}:
-			for _, v := range s {
-				if str, ok := v.(string); ok {
-					tokenScopes = append(tokenScopes, str)
-				}
-			}
-		}
-
-		scopeMap := make(map[string]bool)
-		for _, s := range tokenScopes {
-			scopeMap[s] = true
-		}
-
-		for _, requiredScope := range a.ScopesRequired {
-			if !scopeMap[requiredScope] {
-				return nil, fmt.Errorf("missing required scope: %s", requiredScope)
-			}
-		}
-	}
 
 	// Return claims dynamically
 	claimsMap := make(map[string]any)

internal/auth/generic/generic_test.go

@@ -95,7 +95,7 @@ func TestGetClaimsFromHeader(t *testing.T) {
 		Name:                   "test-generic-auth",
 		Type:                   "generic",
 		Audience:               "my-audience",
-		McpEnabled:             true,
+		McpEnabled:             false,
 		AuthorizationServerURL: server.URL,
 		ScopesRequired:         []string{"read:files"},
 	}
@@ -129,7 +129,7 @@ func TestGetClaimsFromHeader(t *testing.T) {
 					"exp":   time.Now().Add(time.Hour).Unix(),
 				})
 				header := http.Header{}
-				header.Set("Authorization", "Bearer "+token)
+				header.Set("test-generic-auth_token", token)
 				return header
 			},
 			wantError: false,
@@ -151,16 +151,7 @@ func TestGetClaimsFromHeader(t *testing.T) {
 				}
 			},
 		},
-		{
-			name: "invalid formatting",
-			setupHeader: func() http.Header {
-				header := http.Header{}
-				header.Set("Authorization", "Token something")
-				return header
-			},
-			wantError:   true,
-			errContains: "authorization header format must be Bearer {token}",
-		},
+
 		{
 			name: "wrong audience",
 			setupHeader: func() http.Header {
@@ -170,41 +161,13 @@ func TestGetClaimsFromHeader(t *testing.T) {
 					"exp":   time.Now().Add(time.Hour).Unix(),
 				})
 				header := http.Header{}
-				header.Set("Authorization", "Bearer "+token)
+				header.Set("test-generic-auth_token", token)
 				return header
 			},
 			wantError:   true,
 			errContains: "audience validation failed",
 		},
-		{
-			name: "missing required scope",
-			setupHeader: func() http.Header {
-				token := generateValidToken(t, privateKey, keyID, jwt.MapClaims{
-					"aud":   "my-audience",
-					"scope": "some:other_scope",
-					"exp":   time.Now().Add(time.Hour).Unix(),
-				})
-				header := http.Header{}
-				header.Set("Authorization", "Bearer "+token)
-				return header
-			},
-			wantError:   true,
-			errContains: "missing required scope: read:files",
-		},
-		{
-			name: "client_id used instead of aud (valid)",
-			setupHeader: func() http.Header {
-				token := generateValidToken(t, privateKey, keyID, jwt.MapClaims{
-					"client_id": "my-audience",
-					"scope":     []interface{}{"read:files"}, // Testing slice type scopes
-					"exp":       time.Now().Add(time.Hour).Unix(),
-				})
-				header := http.Header{}
-				header.Set("Authorization", "Bearer "+token)
-				return header
-			},
-			wantError: false,
-		},
+
 		{
 			name: "expired token",
 			setupHeader: func() http.Header {
@@ -214,7 +177,7 @@ func TestGetClaimsFromHeader(t *testing.T) {
 					"exp":   time.Now().Add(-1 * time.Hour).Unix(),
 				})
 				header := http.Header{}
-				header.Set("Authorization", "Bearer "+token)
+				header.Set("test-generic-auth_token", token)
 				return header
 			},
 			wantError:   true,
@@ -231,8 +194,6 @@ func TestGetClaimsFromHeader(t *testing.T) {
 				if err == nil {
 					t.Fatalf("expected error, got nil")
 				}
-				// We don't check for exact prefix because jwt library errors can be complex,
-				// check for substring or simple failure instead.
 				if tc.errContains != "" && !strings.Contains(err.Error(), tc.errContains) {
 					t.Errorf("expected error containing %q, got: %v", tc.errContains, err)
 				}

tests/http/http_integration_test.go

@@ -17,6 +17,8 @@ package http
 import (
 	"bytes"
 	"context"
+	"crypto/rand"
+	"crypto/rsa"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -28,6 +30,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/MicahParks/jwkset"
+	"github.com/golang-jwt/jwt/v5"
 	"github.com/googleapis/genai-toolbox/internal/testutils"
 	"github.com/googleapis/genai-toolbox/internal/util/parameters"
 	"github.com/googleapis/genai-toolbox/tests"
@@ -307,9 +311,40 @@ func TestHttpToolEndpoints(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
 	defer cancel()
 
+	// Set up generic auth mock server
+	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	if err != nil {
+		t.Fatalf("failed to create RSA private key: %v", err)
+	}
+	jwksServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/.well-known/openid-configuration" {
+			w.Header().Set("Content-Type", "application/json")
+			_ = json.NewEncoder(w).Encode(map[string]interface{}{
+				"issuer":   "https://example.com",
+				"jwks_uri": "http://" + r.Host + "/jwks",
+			})
+			return
+		}
+		if r.URL.Path == "/jwks" {
+			options := jwkset.JWKOptions{
+				Metadata: jwkset.JWKMetadataOptions{
+					KID: "test-key-id",
+				},
+			}
+			jwk, _ := jwkset.NewJWKFromKey(privateKey.Public(), options)
+			w.Header().Set("Content-Type", "application/json")
+			_ = json.NewEncoder(w).Encode(map[string]interface{}{
+				"keys": []jwkset.JWKMarshal{jwk.Marshal()},
+			})
+			return
+		}
+		http.NotFound(w, r)
+	}))
+	defer jwksServer.Close()
+
 	var args []string
 
-	toolsFile := getHTTPToolsConfig(sourceConfig, HttpToolType)
+	toolsFile := getHTTPToolsConfig(sourceConfig, HttpToolType, jwksServer.URL)
 	cmd, cleanup, err := tests.StartCmd(ctx, toolsFile, args...)
 	if err != nil {
 		t.Fatalf("command initialization returned an error: %s", err)
@@ -329,6 +364,69 @@ func TestHttpToolEndpoints(t *testing.T) {
 	tests.RunToolInvokeTest(t, `"hello world"`, tests.DisableArrayTest())
 	runAdvancedHTTPInvokeTest(t)
 	runQueryParamInvokeTest(t)
+	runGenericAuthInvokeTest(t, privateKey)
+}
+
+func runGenericAuthInvokeTest(t *testing.T, privateKey *rsa.PrivateKey) {
+	// Generate valid token
+	token := jwt.NewWithClaims(jwt.SigningMethodRS256, jwt.MapClaims{
+		"aud":   "test-audience",
+		"scope": "read:files",
+		"sub":   "test-user",
+		"exp":   time.Now().Add(time.Hour).Unix(),
+	})
+	token.Header["kid"] = "test-key-id"
+	signedString, err := token.SignedString(privateKey)
+	if err != nil {
+		t.Fatalf("failed to sign token: %v", err)
+	}
+
+	api := "http://127.0.0.1:5000/api/tool/my-auth-required-generic-tool/invoke"
+
+	// Test without auth header (should fail)
+	t.Run("invoke generic auth tool without token", func(t *testing.T) {
+		req, _ := http.NewRequest(http.MethodPost, api, bytes.NewBuffer([]byte(`{}`)))
+		req.Header.Add("Content-type", "application/json")
+		resp, err := http.DefaultClient.Do(req)
+		if err != nil {
+			t.Fatalf("unable to send request: %s", err)
+		}
+		defer resp.Body.Close()
+
+		var body map[string]interface{}
+		json.NewDecoder(resp.Body).Decode(&body)
+		resultStr, _ := body["result"].(string)
+		if !strings.Contains(resultStr, "unauthorized") && !strings.Contains(resultStr, "missing") {
+			bodyBytes, _ := json.Marshal(body)
+			t.Fatalf("expected unauthorized error, got: %s", string(bodyBytes))
+		}
+	})
+
+	// Test with valid token
+	t.Run("invoke generic auth tool with valid token", func(t *testing.T) {
+		req, _ := http.NewRequest(http.MethodPost, api, bytes.NewBuffer([]byte(`{}`)))
+		req.Header.Add("Content-type", "application/json")
+		req.Header.Add("my-generic-auth_token", signedString)
+
+		resp, err := http.DefaultClient.Do(req)
+		if err != nil {
+			t.Fatalf("unable to send request: %s", err)
+		}
+		defer resp.Body.Close()
+
+		if resp.StatusCode != http.StatusOK {
+			bodyBytes, _ := io.ReadAll(resp.Body)
+			t.Fatalf("expected status 200, got %d: %s", resp.StatusCode, string(bodyBytes))
+		}
+
+		var body map[string]interface{}
+		json.NewDecoder(resp.Body).Decode(&body)
+		got, ok := body["result"].(string)
+		if !ok || got != `"hello world"` {
+			bodyBytes, _ := json.Marshal(body)
+			t.Fatalf("unexpected result: %s", string(bodyBytes))
+		}
+	})
 }
 
 // runQueryParamInvokeTest runs the tool invoke endpoint for the query param test tool
@@ -500,7 +598,7 @@ func runAdvancedHTTPInvokeTest(t *testing.T) {
 }
 
 // getHTTPToolsConfig returns a mock HTTP tool's config file
-func getHTTPToolsConfig(sourceConfig map[string]any, toolType string) map[string]any {
+func getHTTPToolsConfig(sourceConfig map[string]any, toolType string, jwksURL string) map[string]any {
 	// Write config into a file and pass it to command
 	otherSourceConfig := make(map[string]any)
 	for k, v := range sourceConfig {
@@ -519,6 +617,12 @@ func getHTTPToolsConfig(sourceConfig map[string]any, toolType string) map[string
 				"type":     "google",
 				"clientId": tests.ClientId,
 			},
+			"my-generic-auth": map[string]any{
+				"type":                   "generic",
+				"audience":               "test-audience",
+				"authorizationServerUrl": jwksURL,
+				"scopesRequired":         []string{"read:files"},
+			},
 		},
 		"tools": map[string]any{
 			"my-simple-tool": map[string]any{
@@ -598,6 +702,15 @@ func getHTTPToolsConfig(sourceConfig map[string]any, toolType string) map[string
 				"requestBody":  "{}",
 				"authRequired": []string{"my-google-auth"},
 			},
+			"my-auth-required-generic-tool": map[string]any{
+				"type":         toolType,
+				"source":       "my-instance",
+				"method":       "POST",
+				"path":         "/tool0",
+				"description":  "some description",
+				"requestBody":  "{}",
+				"authRequired": []string{"my-generic-auth"},
+			},
 			"my-advanced-tool": map[string]any{
 				"type":        toolType,
 				"source":      "other-instance",