From 3c33198a1d2533ed0e6c90734ff3952b960147bd Mon Sep 17 00:00:00 2001 From: yatuk <92943471+yatuk@users.noreply.github.com> Date: Fri, 7 Aug 2026 20:36:09 +0300 Subject: [PATCH] feat: support Vault JWT auth method for gateway deployments Adds VAULT_AUTH_METHOD=jwt as an alternative to the current static VAULT_TOKEN / X-Vault-Token auth. In gateway-fronted deployments (Toolhive, any OIDC-terminating reverse proxy) every user currently reaches Vault as the same service-account identity, since only a static token is supported. That collapses per-user RBAC into a single shared identity and breaks audit attribution at the Vault layer. With JWT mode enabled, VaultContextMiddleware reads the incoming Authorization: Bearer header, exchanges it via Vault's POST /v1/auth//login (role + jwt), and injects the resulting short-lived, user-scoped Vault token into the request context in place of a static one. The exchanged token is cached in memory, keyed by a hash of the JWT, until close to its lease expiry (or a configured VAULT_AUTH_JWT_CACHE_TTL cap), so repeat requests from the same session don't re-hit /login every time. Static-token mode remains the default and is unchanged; the new code path only runs when VAULT_AUTH_METHOD=jwt is set. A failed exchange (missing JWT, Vault rejecting it, Vault unreachable) returns 401/403/503 with a reason and never falls back to a static token. New env vars: VAULT_AUTH_METHOD, VAULT_AUTH_JWT_PATH, VAULT_AUTH_JWT_ROLE, VAULT_AUTH_JWT_HEADER, VAULT_AUTH_JWT_CACHE_TTL. --- README.md | 26 +++ pkg/client/client.go | 20 ++- pkg/client/jwt_auth.go | 188 +++++++++++++++++++++ pkg/client/jwt_auth_test.go | 299 ++++++++++++++++++++++++++++++++++ pkg/client/middleware.go | 54 +++++- pkg/client/middleware_test.go | 130 +++++++++++++++ 6 files changed, 711 insertions(+), 6 deletions(-) create mode 100644 pkg/client/jwt_auth.go create mode 100644 pkg/client/jwt_auth_test.go diff --git a/README.md b/README.md index 111399e..23d8205 100644 --- a/README.md +++ b/README.md @@ -79,6 +79,11 @@ The server can be configured using environment variables: - `MCP_TLS_KEY_FILE`: Location of the TLS key file (e.g. `/path/to/key.pem`)(default: `""`) - `MCP_RATE_LIMIT_GLOBAL`: Global rate limit (format: `rps:burst`) (default: `10:20`) - `MCP_RATE_LIMIT_SESSION`: Per-session rate limit (format: `rps:burst`) (default: `5:10`) +- `VAULT_AUTH_METHOD`: `token` (default) or `jwt`. See [JWT Auth Mode](#jwt-auth-mode-gateway-deployments) below. +- `VAULT_AUTH_JWT_PATH`: Mount path of Vault's JWT auth method (default: `jwt`) +- `VAULT_AUTH_JWT_ROLE`: Vault role to authenticate against (required when `VAULT_AUTH_METHOD=jwt`) +- `VAULT_AUTH_JWT_HEADER`: Header the incoming JWT is read from (default: `Authorization`, expects a `Bearer ` value) +- `VAULT_AUTH_JWT_CACHE_TTL`: Optional cap, in seconds, on how long an exchanged Vault token is cached (default: the token's own lease duration) ## HTTP Mode Configuration @@ -96,6 +101,27 @@ The HTTP server includes a comprehensive middleware stack: - **Vault Context Middleware**: Extracts Vault configuration and adds to request context - **Logging Middleware**: Structured HTTP request logging +### JWT auth mode (gateway deployments) + +By default the server authenticates to Vault with a static token (`VAULT_TOKEN` or `X-Vault-Token`). Every request then reaches Vault as the same identity, which doesn't work well behind a gateway that already knows who the user is. + +Set `VAULT_AUTH_METHOD=jwt` to have the server exchange the caller's JWT for a short-lived, user-scoped Vault token on each request instead: + +1. A reverse proxy or gateway (e.g. [Toolhive](https://docs.stacklok.com/toolhive)) authenticates the user via OIDC and forwards their token as `Authorization: Bearer `. +2. The server reads that header and calls Vault's `POST /v1/auth//login` with the configured `VAULT_AUTH_JWT_ROLE` and the JWT. +3. Vault verifies the JWT itself (signature, issuer, audience, per the JWT auth method's configuration) and returns a token scoped to whatever policies that role maps to. +4. The resulting token is used for the rest of the request and cached in memory, keyed by the JWT, until it's close to expiry (or until `VAULT_AUTH_JWT_CACHE_TTL` elapses, if set). + +`VAULT_TOKEN` and `X-Vault-Token` are ignored while JWT mode is active. If the JWT is missing or Vault rejects it, the request fails with 401 or 403 rather than falling back to a static token. + +```bash +export VAULT_AUTH_METHOD=jwt +export VAULT_AUTH_JWT_ROLE=mcp-gateway +export VAULT_ADDR=https://vault.internal:8200 +``` + +This assumes Vault's JWT auth method is already enabled and mapped to your identity provider; see [Vault's JWT auth docs](https://developer.hashicorp.com/vault/docs/auth/jwt) for that side of the setup. + ## Integration with Visual Studio Code 1. In your project workspace root, create or open the `.vscode/mcp.json` configuration file. Alternatively, to add an MCP to your user configuration, run the `MCP: Open User Configuration` command, which opens the mcp.json file in your user profile. If the file does not exist, VS Code creates it for you. diff --git a/pkg/client/client.go b/pkg/client/client.go index 70fea37..0d12dec 100644 --- a/pkg/client/client.go +++ b/pkg/client/client.go @@ -56,9 +56,9 @@ func getEnv(key, fallback string) string { return fallback } -// NewVaultClient creates a new Vault client for the given session -func NewVaultClient(sessionId string, vaultAddress string, vaultSkipTLSVerify bool, vaultToken string, vaultNamespace string) (*api.Client, error) { - // Initialize Vault client +// buildVaultClient creates an *api.Client for the given address and TLS +// setting, without touching the session cache or setting a token. +func buildVaultClient(vaultAddress string, vaultSkipTLSVerify bool, vaultNamespace string) (*api.Client, error) { config := api.DefaultConfig() config.Address = vaultAddress @@ -72,12 +72,22 @@ func NewVaultClient(sessionId string, vaultAddress string, vaultSkipTLSVerify bo return nil, fmt.Errorf("api.NewClient failed to create Vault client: %v", err) } - client.SetToken(vaultToken) - if vaultNamespace != "" { client.SetNamespace(vaultNamespace) } + return client, nil +} + +// NewVaultClient creates a new Vault client for the given session +func NewVaultClient(sessionId string, vaultAddress string, vaultSkipTLSVerify bool, vaultToken string, vaultNamespace string) (*api.Client, error) { + client, err := buildVaultClient(vaultAddress, vaultSkipTLSVerify, vaultNamespace) + if err != nil { + return nil, err + } + + client.SetToken(vaultToken) + activeClients.Store(sessionId, &sessionEntry{client: client, tokenHash: hashToken(vaultToken)}) return client, nil diff --git a/pkg/client/jwt_auth.go b/pkg/client/jwt_auth.go new file mode 100644 index 0000000..8bbd5d1 --- /dev/null +++ b/pkg/client/jwt_auth.go @@ -0,0 +1,188 @@ +// Copyright IBM Corp. 2025, 2026 +// SPDX-License-Identifier: MPL-2.0 + +package client + +import ( + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "sync" + "time" + + "github.com/hashicorp/vault/api" + log "github.com/sirupsen/logrus" +) + +const ( + VaultAuthMethod = "VAULT_AUTH_METHOD" + VaultAuthJWTPath = "VAULT_AUTH_JWT_PATH" + VaultAuthJWTRole = "VAULT_AUTH_JWT_ROLE" + VaultAuthJWTHeader = "VAULT_AUTH_JWT_HEADER" + VaultAuthJWTCacheTTL = "VAULT_AUTH_JWT_CACHE_TTL" +) + +const ( + defaultJWTAuthPath = "jwt" + defaultJWTAuthHeader = "Authorization" + + // cacheSafetyMargin keeps a cached token from being handed out right at + // the edge of its lease expiry. + cacheSafetyMargin = 10 * time.Second +) + +// jwtAuthEnabled reports whether VAULT_AUTH_METHOD=jwt is configured. +func jwtAuthEnabled() bool { + return strings.EqualFold(getEnv(VaultAuthMethod, "token"), "jwt") +} + +// jwtAuthHeaderName returns the header the incoming JWT is read from. +func jwtAuthHeaderName() string { + return getEnv(VaultAuthJWTHeader, defaultJWTAuthHeader) +} + +// jwtLoginError is a JWT exchange failure carrying the HTTP status the +// caller should see. Acceptance criteria for this feature require that +// exchange failures are surfaced clearly rather than silently falling back +// to a static token, so this type is what the middleware inspects to pick +// the response status. +type jwtLoginError struct { + status int + message string +} + +func (e *jwtLoginError) Error() string { return e.message } +func (e *jwtLoginError) StatusCode() int { return e.status } + +// extractBearerToken pulls the JWT out of a header value such as +// "Bearer ". A bare token, with no "Bearer " prefix, is also +// accepted since some gateways forward the JWT without it. +func extractBearerToken(headerValue string) (string, error) { + headerValue = strings.TrimSpace(headerValue) + if headerValue == "" { + return "", &jwtLoginError{http.StatusUnauthorized, fmt.Sprintf("missing JWT: %s header not provided", jwtAuthHeaderName())} + } + if rest, ok := strings.CutPrefix(headerValue, "Bearer "); ok { + headerValue = strings.TrimSpace(rest) + } else if strings.EqualFold(headerValue, "Bearer") { + headerValue = "" + } + if headerValue == "" { + return "", &jwtLoginError{http.StatusUnauthorized, "missing JWT: empty bearer token"} + } + return headerValue, nil +} + +type cachedToken struct { + token string + expiresAt time.Time +} + +// jwtTokenCache maps sha256(jwt) to the Vault token obtained for it, so +// repeat requests within the token's lease don't re-hit /login. +var jwtTokenCache sync.Map + +func cacheKey(jwt string) string { + sum := sha256.Sum256([]byte(jwt)) + return hex.EncodeToString(sum[:]) +} + +func lookupCachedToken(jwt string) (string, bool) { + key := cacheKey(jwt) + value, ok := jwtTokenCache.Load(key) + if !ok { + return "", false + } + entry := value.(cachedToken) + if time.Now().After(entry.expiresAt) { + jwtTokenCache.Delete(key) + return "", false + } + return entry.token, true +} + +func storeCachedToken(jwt, vaultToken string, leaseDuration time.Duration) { + if leaseDuration <= cacheSafetyMargin { + // Lease too short to be worth caching; every request will just + // exchange again. + return + } + + if cap := getEnv(VaultAuthJWTCacheTTL, ""); cap != "" { + if capSeconds, err := strconv.Atoi(cap); err == nil { + if capDuration := time.Duration(capSeconds) * time.Second; capDuration < leaseDuration { + leaseDuration = capDuration + } + } + } + + jwtTokenCache.Store(cacheKey(jwt), cachedToken{ + token: vaultToken, + expiresAt: time.Now().Add(leaseDuration - cacheSafetyMargin), + }) +} + +// exchangeJWTForVaultToken calls Vault's JWT auth login endpoint and returns +// a short-lived, user-scoped Vault token plus its lease duration. +func exchangeJWTForVaultToken(vaultAddress, vaultNamespace string, vaultSkipTLSVerify bool, jwt string) (string, time.Duration, error) { + role := getEnv(VaultAuthJWTRole, "") + if role == "" { + return "", 0, &jwtLoginError{http.StatusUnauthorized, "VAULT_AUTH_JWT_ROLE is not configured"} + } + mount := getEnv(VaultAuthJWTPath, defaultJWTAuthPath) + + loginClient, err := buildVaultClient(vaultAddress, vaultSkipTLSVerify, vaultNamespace) + if err != nil { + return "", 0, fmt.Errorf("failed to build Vault client for JWT login: %w", err) + } + + secret, err := loginClient.Logical().Write(fmt.Sprintf("auth/%s/login", mount), map[string]interface{}{ + "role": role, + "jwt": jwt, + }) + if err != nil { + var respErr *api.ResponseError + if errors.As(err, &respErr) { + status := http.StatusUnauthorized + if respErr.StatusCode == http.StatusForbidden { + status = http.StatusForbidden + } + msg := "vault rejected the JWT" + if len(respErr.Errors) > 0 { + msg = strings.Join(respErr.Errors, "; ") + } + return "", 0, &jwtLoginError{status, msg} + } + // Not a Vault API error, e.g. the login request never reached + // Vault. That's a connectivity problem, not an auth rejection. + return "", 0, &jwtLoginError{http.StatusServiceUnavailable, fmt.Sprintf("could not reach Vault to exchange JWT: %v", err)} + } + if secret == nil || secret.Auth == nil || secret.Auth.ClientToken == "" { + return "", 0, &jwtLoginError{http.StatusUnauthorized, "vault JWT login returned no token"} + } + + return secret.Auth.ClientToken, time.Duration(secret.Auth.LeaseDuration) * time.Second, nil +} + +// resolveJWTVaultToken returns a Vault token for the given JWT, using the +// in-memory cache when possible and exchanging with Vault otherwise. +func resolveJWTVaultToken(vaultAddress, vaultNamespace string, vaultSkipTLSVerify bool, jwt string, logger *log.Logger) (string, error) { + if cached, ok := lookupCachedToken(jwt); ok { + return cached, nil + } + + vaultToken, leaseDuration, err := exchangeJWTForVaultToken(vaultAddress, vaultNamespace, vaultSkipTLSVerify, jwt) + if err != nil { + return "", err + } + + storeCachedToken(jwt, vaultToken, leaseDuration) + if logger != nil { + logger.Debug("Vault token obtained via JWT exchange") + } + return vaultToken, nil +} diff --git a/pkg/client/jwt_auth_test.go b/pkg/client/jwt_auth_test.go new file mode 100644 index 0000000..0897181 --- /dev/null +++ b/pkg/client/jwt_auth_test.go @@ -0,0 +1,299 @@ +// Copyright IBM Corp. 2025, 2026 +// SPDX-License-Identifier: MPL-2.0 + +package client + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "sync" + "testing" + "time" + + log "github.com/sirupsen/logrus" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func resetJWTCache() { + jwtTokenCache = sync.Map{} +} + +func testLogger() *log.Logger { + logger := log.New() + logger.SetOutput(os.Stdout) + logger.SetLevel(log.ErrorLevel) + return logger +} + +func TestExtractBearerToken(t *testing.T) { + tests := []struct { + name string + header string + wantToken string + wantErr bool + }{ + {"bearer prefix", "Bearer abc.def.ghi", "abc.def.ghi", false}, + {"bare token", "abc.def.ghi", "abc.def.ghi", false}, + {"empty header", "", "", true}, + {"bearer with no token", "Bearer ", "", true}, + {"whitespace only", " ", "", true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + token, err := extractBearerToken(tt.header) + if tt.wantErr { + require.Error(t, err) + var loginErr *jwtLoginError + require.ErrorAs(t, err, &loginErr) + assert.Equal(t, http.StatusUnauthorized, loginErr.StatusCode()) + return + } + require.NoError(t, err) + assert.Equal(t, tt.wantToken, token) + }) + } +} + +func TestExchangeJWTForVaultToken(t *testing.T) { + t.Run("missing role is rejected before any request is made", func(t *testing.T) { + os.Unsetenv(VaultAuthJWTRole) + _, _, err := exchangeJWTForVaultToken("http://127.0.0.1:1", "", false, "some.jwt.token") + require.Error(t, err) + var loginErr *jwtLoginError + require.ErrorAs(t, err, &loginErr) + assert.Equal(t, http.StatusUnauthorized, loginErr.StatusCode()) + assert.Contains(t, loginErr.Error(), "VAULT_AUTH_JWT_ROLE") + }) + + t.Run("successful login returns token and lease", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "/v1/auth/jwt/login", r.URL.Path) + + var body map[string]string + require.NoError(t, json.NewDecoder(r.Body).Decode(&body)) + assert.Equal(t, "my-role", body["role"]) + assert.Equal(t, "my.jwt.token", body["jwt"]) + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "auth": map[string]interface{}{ + "client_token": "s.generated-token", + "lease_duration": 3600, + }, + }) + })) + defer server.Close() + + os.Setenv(VaultAuthJWTRole, "my-role") + defer os.Unsetenv(VaultAuthJWTRole) + + token, lease, err := exchangeJWTForVaultToken(server.URL, "", false, "my.jwt.token") + require.NoError(t, err) + assert.Equal(t, "s.generated-token", token) + assert.Equal(t, 3600*time.Second, lease) + }) + + t.Run("custom mount path is used", func(t *testing.T) { + var gotPath string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "auth": map[string]interface{}{ + "client_token": "s.token", + "lease_duration": 60, + }, + }) + })) + defer server.Close() + + os.Setenv(VaultAuthJWTRole, "my-role") + os.Setenv(VaultAuthJWTPath, "custom-jwt-mount") + defer os.Unsetenv(VaultAuthJWTRole) + defer os.Unsetenv(VaultAuthJWTPath) + + _, _, err := exchangeJWTForVaultToken(server.URL, "", false, "jwt") + require.NoError(t, err) + assert.Equal(t, "/v1/auth/custom-jwt-mount/login", gotPath) + }) + + t.Run("403 from vault is surfaced as forbidden", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusForbidden) + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "errors": []string{"permission denied"}, + }) + })) + defer server.Close() + + os.Setenv(VaultAuthJWTRole, "my-role") + defer os.Unsetenv(VaultAuthJWTRole) + + _, _, err := exchangeJWTForVaultToken(server.URL, "", false, "bad.jwt.token") + require.Error(t, err) + var loginErr *jwtLoginError + require.ErrorAs(t, err, &loginErr) + assert.Equal(t, http.StatusForbidden, loginErr.StatusCode()) + assert.Contains(t, loginErr.Error(), "permission denied") + }) + + t.Run("400 from vault is surfaced as unauthorized", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "errors": []string{"missing jwt"}, + }) + })) + defer server.Close() + + os.Setenv(VaultAuthJWTRole, "my-role") + defer os.Unsetenv(VaultAuthJWTRole) + + _, _, err := exchangeJWTForVaultToken(server.URL, "", false, "") + require.Error(t, err) + var loginErr *jwtLoginError + require.ErrorAs(t, err, &loginErr) + assert.Equal(t, http.StatusUnauthorized, loginErr.StatusCode()) + }) + + t.Run("unreachable vault is surfaced as service unavailable", func(t *testing.T) { + os.Setenv(VaultAuthJWTRole, "my-role") + defer os.Unsetenv(VaultAuthJWTRole) + + _, _, err := exchangeJWTForVaultToken("http://127.0.0.1:1", "", false, "some.jwt.token") + require.Error(t, err) + var loginErr *jwtLoginError + require.ErrorAs(t, err, &loginErr) + assert.Equal(t, http.StatusServiceUnavailable, loginErr.StatusCode()) + }) +} + +func TestResolveJWTVaultToken(t *testing.T) { + logger := testLogger() + + t.Run("caches token across calls with the same jwt", func(t *testing.T) { + resetJWTCache() + var loginCalls int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + loginCalls++ + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "auth": map[string]interface{}{ + "client_token": "s.cached-token", + "lease_duration": 3600, + }, + }) + })) + defer server.Close() + + os.Setenv(VaultAuthJWTRole, "my-role") + defer os.Unsetenv(VaultAuthJWTRole) + + token1, err := resolveJWTVaultToken(server.URL, "", false, "same.jwt.token", logger) + require.NoError(t, err) + token2, err := resolveJWTVaultToken(server.URL, "", false, "same.jwt.token", logger) + require.NoError(t, err) + + assert.Equal(t, token1, token2) + assert.Equal(t, 1, loginCalls) + }) + + t.Run("does not cache a lease shorter than the safety margin", func(t *testing.T) { + resetJWTCache() + var loginCalls int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + loginCalls++ + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "auth": map[string]interface{}{ + "client_token": "s.short-lived-token", + "lease_duration": 5, + }, + }) + })) + defer server.Close() + + os.Setenv(VaultAuthJWTRole, "my-role") + defer os.Unsetenv(VaultAuthJWTRole) + + _, err := resolveJWTVaultToken(server.URL, "", false, "short.jwt.token", logger) + require.NoError(t, err) + _, err = resolveJWTVaultToken(server.URL, "", false, "short.jwt.token", logger) + require.NoError(t, err) + + assert.Equal(t, 2, loginCalls) + }) + + t.Run("different jwts get independent cache entries", func(t *testing.T) { + resetJWTCache() + var loginCalls int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + loginCalls++ + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "auth": map[string]interface{}{ + "client_token": "s.token-for-user", + "lease_duration": 3600, + }, + }) + })) + defer server.Close() + + os.Setenv(VaultAuthJWTRole, "my-role") + defer os.Unsetenv(VaultAuthJWTRole) + + _, err := resolveJWTVaultToken(server.URL, "", false, "user-a.jwt", logger) + require.NoError(t, err) + _, err = resolveJWTVaultToken(server.URL, "", false, "user-b.jwt", logger) + require.NoError(t, err) + + assert.Equal(t, 2, loginCalls) + }) + + t.Run("VAULT_AUTH_JWT_CACHE_TTL caps the cache lifetime", func(t *testing.T) { + resetJWTCache() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "auth": map[string]interface{}{ + "client_token": "s.capped-token", + "lease_duration": 3600, + }, + }) + })) + defer server.Close() + + os.Setenv(VaultAuthJWTRole, "my-role") + os.Setenv(VaultAuthJWTCacheTTL, "20") + defer os.Unsetenv(VaultAuthJWTRole) + defer os.Unsetenv(VaultAuthJWTCacheTTL) + + before := time.Now() + _, err := resolveJWTVaultToken(server.URL, "", false, "capped.jwt", logger) + require.NoError(t, err) + + value, ok := jwtTokenCache.Load(cacheKey("capped.jwt")) + require.True(t, ok) + entry := value.(cachedToken) + // Lease was 3600s but the cap is 20s, so the cache entry should + // expire in roughly cap-margin seconds, not close to an hour. + assert.WithinDuration(t, before.Add(20*time.Second-cacheSafetyMargin), entry.expiresAt, 2*time.Second) + }) + + t.Run("login failure is propagated and nothing is cached", func(t *testing.T) { + resetJWTCache() + os.Unsetenv(VaultAuthJWTRole) + + _, err := resolveJWTVaultToken("http://127.0.0.1:1", "", false, "unconfigured.jwt", logger) + require.Error(t, err) + + _, ok := jwtTokenCache.Load(cacheKey("unconfigured.jwt")) + assert.False(t, ok) + }) +} diff --git a/pkg/client/middleware.go b/pkg/client/middleware.go index 0e9ebc8..b1acb0e 100644 --- a/pkg/client/middleware.go +++ b/pkg/client/middleware.go @@ -5,10 +5,12 @@ package client import ( "context" + "errors" "fmt" "net/http" "net/textproto" "os" + "strconv" "strings" log "github.com/sirupsen/logrus" @@ -130,7 +132,12 @@ func (h *securityHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { func VaultContextMiddleware(logger *log.Logger) func(http.Handler) http.Handler { return func(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - requiredHeaders := []string{VaultAddress, VaultToken, VaultHeaderToken, VaultSkipTLSVerify} + requiredHeaders := []string{VaultAddress, VaultSkipTLSVerify} + if !jwtAuthEnabled() { + // In JWT mode the token comes from the JWT exchange below, + // not from a static VAULT_TOKEN / X-Vault-Token value. + requiredHeaders = append(requiredHeaders, VaultToken, VaultHeaderToken) + } ctx := r.Context() for _, header := range requiredHeaders { @@ -184,12 +191,57 @@ func VaultContextMiddleware(logger *log.Logger) func(http.Handler) http.Handler logger.Debug("Vault namespace configured via request context") } + if jwtAuthEnabled() { + rawHeader := r.Header.Get(textproto.CanonicalMIMEHeaderKey(jwtAuthHeaderName())) + jwt, err := extractBearerToken(rawHeader) + if err != nil { + writeJWTAuthError(w, logger, r.RemoteAddr, err) + return + } + + vaultAddress, _ := ctx.Value(contextKey(VaultAddress)).(string) + if vaultAddress == "" { + vaultAddress = getEnv(VaultAddress, DefaultVaultAddress) + } + vaultNamespace, _ := ctx.Value(contextKey(VaultNamespace)).(string) + var vaultSkipTLSVerify bool + if skipStr, ok := ctx.Value(contextKey(VaultSkipTLSVerify)).(string); ok { + vaultSkipTLSVerify, _ = strconv.ParseBool(skipStr) + } + + vaultToken, err := resolveJWTVaultToken(vaultAddress, vaultNamespace, vaultSkipTLSVerify, jwt, logger) + if err != nil { + writeJWTAuthError(w, logger, r.RemoteAddr, err) + return + } + + ctx = context.WithValue(ctx, contextKey(VaultToken), vaultToken) + } + // Call the next handler with the enriched context next.ServeHTTP(w, r.WithContext(ctx)) }) } } +// writeJWTAuthError writes the HTTP status and message for a failed JWT +// exchange. Errors that aren't a *jwtLoginError, such as failing to build +// the login client, fall back to 401 so a JWT auth failure never falls +// through to the static-token path. +func writeJWTAuthError(w http.ResponseWriter, logger *log.Logger, remoteAddr string, err error) { + var loginErr *jwtLoginError + status := http.StatusUnauthorized + message := err.Error() + if errors.As(err, &loginErr) { + status = loginErr.StatusCode() + message = loginErr.Error() + } + if logger != nil { + logger.WithField("remote_addr", remoteAddr).Warn("Vault JWT authentication failed: " + message) + } + http.Error(w, message, status) +} + // LoggingMiddleware logs HTTP requests with structured logging func LoggingMiddleware(logger *log.Logger) func(http.Handler) http.Handler { return func(next http.Handler) http.Handler { diff --git a/pkg/client/middleware_test.go b/pkg/client/middleware_test.go index 771f764..6fe376e 100644 --- a/pkg/client/middleware_test.go +++ b/pkg/client/middleware_test.go @@ -490,6 +490,136 @@ func TestVaultContextMiddleware_SecurityLogging(t *testing.T) { }) } +// TestVaultContextMiddleware_JWTMode tests the middleware's behavior when +// VAULT_AUTH_METHOD=jwt is configured: it should exchange the incoming JWT +// for a Vault token instead of reading VAULT_TOKEN / X-Vault-Token. +func TestVaultContextMiddleware_JWTMode(t *testing.T) { + logger := log.New() + logger.SetLevel(log.ErrorLevel) + + setJWTMode := func(t *testing.T, extra map[string]string) { + t.Helper() + os.Setenv(VaultAuthMethod, "jwt") + t.Cleanup(func() { os.Unsetenv(VaultAuthMethod) }) + for k, v := range extra { + os.Setenv(k, v) + t.Cleanup(func() { os.Unsetenv(k) }) + } + } + + t.Run("exchanges the bearer JWT for a vault token", func(t *testing.T) { + resetJWTCache() + vaultServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"auth":{"client_token":"s.exchanged-token","lease_duration":3600}}`)) + })) + defer vaultServer.Close() + + setJWTMode(t, map[string]string{ + VaultAddress: vaultServer.URL, + VaultAuthJWTRole: "my-role", + }) + + var gotToken string + testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + token, _ := r.Context().Value(contextKey(VaultToken)).(string) + gotToken = token + w.WriteHeader(http.StatusOK) + }) + + handler := VaultContextMiddleware(logger)(testHandler) + + req := httptest.NewRequest("GET", "/mcp", nil) + req.Header.Set("Authorization", "Bearer user.jwt.token") + + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + assert.Equal(t, http.StatusOK, rr.Code) + assert.Equal(t, "s.exchanged-token", gotToken) + }) + + t.Run("missing JWT header is rejected with 401", func(t *testing.T) { + resetJWTCache() + setJWTMode(t, map[string]string{VaultAuthJWTRole: "my-role"}) + + testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) + handler := VaultContextMiddleware(logger)(testHandler) + + req := httptest.NewRequest("GET", "/mcp", nil) + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + assert.Equal(t, http.StatusUnauthorized, rr.Code) + }) + + t.Run("vault login failure is surfaced, not silently ignored", func(t *testing.T) { + resetJWTCache() + vaultServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusForbidden) + w.Write([]byte(`{"errors":["permission denied"]}`)) + })) + defer vaultServer.Close() + + setJWTMode(t, map[string]string{ + VaultAddress: vaultServer.URL, + VaultAuthJWTRole: "my-role", + }) + + called := false + testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + w.WriteHeader(http.StatusOK) + }) + handler := VaultContextMiddleware(logger)(testHandler) + + req := httptest.NewRequest("GET", "/mcp", nil) + req.Header.Set("Authorization", "Bearer rejected.jwt.token") + + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + assert.Equal(t, http.StatusForbidden, rr.Code) + assert.False(t, called, "the next handler must not run when the JWT exchange fails") + assert.Contains(t, rr.Body.String(), "permission denied") + }) + + t.Run("static VAULT_TOKEN header is ignored in jwt mode", func(t *testing.T) { + resetJWTCache() + vaultServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"auth":{"client_token":"s.jwt-token","lease_duration":3600}}`)) + })) + defer vaultServer.Close() + + setJWTMode(t, map[string]string{ + VaultAddress: vaultServer.URL, + VaultAuthJWTRole: "my-role", + }) + + var gotToken string + testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + token, _ := r.Context().Value(contextKey(VaultToken)).(string) + gotToken = token + w.WriteHeader(http.StatusOK) + }) + handler := VaultContextMiddleware(logger)(testHandler) + + req := httptest.NewRequest("GET", "/mcp", nil) + req.Header.Set("Authorization", "Bearer user.jwt.token") + req.Header.Set(VaultToken, "should-be-ignored") + + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + assert.Equal(t, http.StatusOK, rr.Code) + assert.Equal(t, "s.jwt-token", gotToken) + }) +} + // TestVaultContextMiddleware_EdgeCases tests edge cases and error conditions func TestVaultContextMiddleware_EdgeCases(t *testing.T) { logger := log.New()