feat: derive token cache expiry from JWT claim

Currently the code will cache a JWT for 5 minutes. This creates a timing
bug and forces unnecessary caching.

Since tokens are cached for 5 minutes, it is possible that a token with
1 second of expiration left is cached. The cache code pulls the token
out of the cache without any validation, the expired token will be used
for the next 4 minutes and 59 seconds.

This work sets the cache expiration time based on the token expiration
time. This results in each token being cached once and evicted from the
cache at expiration.

In cases where the JWT needs to be valid beyond this check (e.g., for
downstream requests), users can optionally choose to expire the cache
before the JWT expiration time with `TokenExpiryBuffer` option. This
option also ensures JWTs that are within the expiration window are
treated as expired.

Signed-off-by: Tyler Drombosky <bowsky@gmail.com>

Author: drombosky

api_test.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"reflect"
 	"strings"
 	"testing"
 	"time"
@@ -708,3 +709,156 @@ func TestWrapMCPEndpointWithValidToken(t *testing.T) {
 		t.Errorf("status = %d, want 200", rec.Code)
 	}
 }
+
+// TestValidateTokenCached tests that only non-expired tokens are cached.
+func TestValidateTokenCached(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name              string
+		expirationFromNow time.Duration
+		tokenExpiryBuffer time.Duration
+		want              *User
+		wantErr           string
+	}{
+		{
+			name:              "success",
+			expirationFromNow: time.Minute,
+			want:              &User{Subject: "testuser"},
+		},
+		{
+			name:              "expired token",
+			expirationFromNow: -time.Minute,
+			wantErr:           "authentication failed: failed to parse and validate token: token has invalid claims: token is expired",
+		},
+		{
+			name:              "token expires in buffer",
+			tokenExpiryBuffer: 5 * time.Minute,
+			expirationFromNow: time.Minute,
+			wantErr:           "authentication failed: token expired or expiring too soon",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			cfg := &Config{
+				Mode:              "native",
+				Provider:          "hmac",
+				Audience:          "api://test",
+				JWTSecret:         []byte("test-secret-key-must-be-32-bytes-long!"),
+				ServerURL:         "https://test-server.com",
+				Issuer:            "https://test.example.com",
+				TokenExpiryBuffer: tt.tokenExpiryBuffer,
+			}
+			srv, err := NewServer(cfg)
+			if err != nil {
+				t.Fatalf("NewServer: %v", err)
+			}
+
+			// Create a valid HMAC token
+			token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
+				"sub": "testuser",
+				"aud": "api://test",
+				"iss": "https://test.example.com",
+				"exp": time.Now().Add(tt.expirationFromNow).Unix(),
+			})
+			tokenString, err := token.SignedString(cfg.JWTSecret)
+			if err != nil {
+				t.Fatalf("SignedString: %v", err)
+			}
+
+			user, err := srv.ValidateTokenCached(context.Background(), tokenString)
+			if tt.wantErr != "" {
+				if err == nil {
+					t.Fatalf("expected error %q, got nil", tt.wantErr)
+				}
+				if err.Error() != tt.wantErr {
+					t.Errorf("expected error %q, got %q", tt.wantErr, err.Error())
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			if !reflect.DeepEqual(user, tt.want) {
+				t.Errorf("expected user %v got user %v", tt.want, user)
+			}
+		})
+	}
+}
+
+// TestValidateTokenCached tests that the cache expires correctly.
+func TestValidateTokenCached_Expires(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name              string
+		expirationFromNow time.Duration
+		tokenExpiryBuffer time.Duration
+		wantErr           string
+	}{
+		{
+			name:              "default expiry buffer",
+			expirationFromNow: time.Second,
+			wantErr:           "authentication failed: failed to parse and validate token: token has invalid claims: token is expired",
+		},
+		{
+			name:              "custom expiry buffer",
+			expirationFromNow: 5 * time.Second,
+			tokenExpiryBuffer: 4 * time.Second,
+			wantErr:           "authentication failed: token expired or expiring too soon",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			cfg := &Config{
+				Mode:              "native",
+				Provider:          "hmac",
+				Audience:          "api://test",
+				JWTSecret:         []byte("test-secret-key-must-be-32-bytes-long!"),
+				ServerURL:         "https://test-server.com",
+				Issuer:            "https://test.example.com",
+				TokenExpiryBuffer: tt.tokenExpiryBuffer,
+			}
+			srv, err := NewServer(cfg)
+			if err != nil {
+				t.Fatalf("NewServer: %v", err)
+			}
+
+			// Create a valid HMAC token
+			token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
+				"sub": "testuser",
+				"aud": "api://test",
+				"iss": "https://test.example.com",
+				"exp": time.Now().Add(tt.expirationFromNow).Unix(),
+			})
+			tokenString, err := token.SignedString(cfg.JWTSecret)
+			if err != nil {
+				t.Fatalf("SignedString: %v", err)
+			}
+
+			// Token is successfully verified and cached
+			_, err = srv.ValidateTokenCached(context.Background(), tokenString)
+			if err != nil {
+				t.Fatalf("unexpected error %s", err)
+			}
+
+			// Wait twice as long as the token should take to expire.
+			time.Sleep(2 * (tt.expirationFromNow - tt.tokenExpiryBuffer))
+
+			_, err = srv.ValidateTokenCached(context.Background(), tokenString)
+			if err == nil {
+				t.Fatal("expected error, got nil")
+			}
+
+			if err.Error() != tt.wantErr {
+				t.Errorf("expected error %q, got %q", tt.wantErr, err.Error())
+			}
+		})
+	}
+}

config.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/tuannvm/oauth-mcp-proxy/provider"
 )
@@ -41,6 +42,11 @@ type Config struct {
 	// The issuer URL to use for issuer validation.
 	// This should only be set if the issuer in the token differs from the standard issuer URL.
 	ValidatorIssuer string
+
+	// TokenExpiryBuffer is subtracted from the JWT's exp claim to determine
+	// effective cache expiry. Tokens with less than this duration remaining
+	// are treated as expired and rejected. Defaults to 0 (no buffer).
+	TokenExpiryBuffer time.Duration
 }
 
 // Validate validates the configuration
@@ -95,6 +101,10 @@ func (c *Config) Validate() error {
 			return fmt.Errorf("proxy mode requires RedirectURIs or FixedRedirectURI")
 		}
 	}
+	// Validate TokenExpiryBuffer is positive.
+	if c.TokenExpiryBuffer < 0 {
+		return fmt.Errorf("TokenExpiryBuffer must be >= 0")
+	}
 
 	return nil
 }
@@ -265,6 +275,13 @@ func (b *ConfigBuilder) WithValidatorIssuer(validatorIssuer string) *ConfigBuild
 	return b
 }
 
+// WithTokenExpiryBuffer sets the buffer subtracted from JWT exp for cache expiry.
+// Tokens with less than d remaining until expiry are treated as expired and rejected.
+func (b *ConfigBuilder) WithTokenExpiryBuffer(d time.Duration) *ConfigBuilder {
+	b.config.TokenExpiryBuffer = d
+	return b
+}
+
 // WithServerURL sets the full server URL directly
 func (b *ConfigBuilder) WithServerURL(url string) *ConfigBuilder {
 	b.config.ServerURL = url

context_propagation_test.go

@@ -41,7 +41,7 @@ func TestContextPropagation(t *testing.T) {
 
 		// Test 1: Normal context works
 		ctx := context.Background()
-		user, err := server.validator.ValidateToken(ctx, tokenString)
+		user, _, err := server.validator.ValidateToken(ctx, tokenString)
 		if err != nil {
 			t.Fatalf("ValidateToken with normal context failed: %v", err)
 		}
@@ -83,7 +83,7 @@ func TestContextPropagation(t *testing.T) {
 
 		// For HMAC validator (local-only), this still succeeds
 		// because HMAC doesn't do I/O and doesn't check context cancellation
-		user, err := server.validator.ValidateToken(ctx, tokenString)
+		user, _, err := server.validator.ValidateToken(ctx, tokenString)
 
 		// HMAC validation is local-only, so it succeeds even with cancelled context
 		if err != nil {
@@ -126,7 +126,7 @@ func TestContextPropagation(t *testing.T) {
 		defer cancel()
 
 		// Validate with timeout context
-		user, err := server.validator.ValidateToken(ctx, tokenString)
+		user, _, err := server.validator.ValidateToken(ctx, tokenString)
 		if err != nil {
 			t.Fatalf("ValidateToken with timeout context failed: %v", err)
 		}

docs/SECURITY.md

@@ -200,7 +200,9 @@ oauth.WithOAuth(mux, &oauth.Config{
 
 ### Cache Behavior
 
-- **Cache TTL:** 5 minutes (hardcoded in v0.1.0)
+- **Cache TTL:** Until the JWT `exp` claim, minus the configured expiry buffer.
+  Tokens inside the buffer are rejected, and tokens without `exp` fall back
+  to a 5-minute cache lifetime.
 - **Cache scope:** Per Server instance
 - **Cache key:** SHA-256 hash of token
 

integration_test.go

@@ -47,7 +47,7 @@ func TestIntegration(t *testing.T) {
 		tokenString, _ := token.SignedString(cfg.JWTSecret)
 
 		// Validate token using provider package directly
-		user, err := validator.ValidateToken(context.Background(), tokenString)
+		user, _, err := validator.ValidateToken(context.Background(), tokenString)
 		if err != nil {
 			t.Fatalf("ValidateToken failed: %v", err)
 		}
@@ -90,7 +90,7 @@ func TestIntegration(t *testing.T) {
 
 		tokenString, _ := token.SignedString(rootCfg.JWTSecret)
 
-		user, err := validator.ValidateToken(context.Background(), tokenString)
+		user, _, err := validator.ValidateToken(context.Background(), tokenString)
 		if err != nil {
 			t.Fatalf("ValidateToken after conversion failed: %v", err)
 		}
@@ -276,7 +276,7 @@ func TestValidatorIntegration(t *testing.T) {
 
 		tokenString, _ := token.SignedString(cfg.JWTSecret)
 
-		user, err := v.ValidateToken(context.Background(), tokenString)
+		user, _, err := v.ValidateToken(context.Background(), tokenString)
 		if err != nil {
 			t.Fatalf("ValidateToken failed: %v", err)
 		}

middleware.go

@@ -2,12 +2,10 @@ package oauth
 
 import (
 	"context"
-	"crypto/sha256"
 	"fmt"
 	"log"
 	"net/http"
 	"strings"
-	"time"
 
 	"github.com/mark3labs/mcp-go/mcp"
 	"github.com/mark3labs/mcp-go/server"
@@ -38,36 +36,13 @@ func (s *Server) Middleware() func(server.ToolHandlerFunc) server.ToolHandlerFun
 				return nil, fmt.Errorf("authentication required: missing OAuth token")
 			}
 
-			// Create token hash for caching
-			tokenHash := fmt.Sprintf("%x", sha256.Sum256([]byte(tokenString)))
-
-			// Check cache first
-			if cached, exists := s.cache.getCachedToken(tokenHash); exists {
-				s.logger.Info("Using cached authentication for tool: %s (user: %s)", req.Params.Name, cached.User.Username)
-				ctx = context.WithValue(ctx, userContextKey, cached.User)
-				return next(ctx, req)
-			}
-
-			// Log token hash for debugging (prevents sensitive data exposure)
-			tokenHashFull := fmt.Sprintf("%x", sha256.Sum256([]byte(tokenString)))
-			tokenHashPreview := tokenHashFull[:16] + "..."
-			s.logger.Info("Validating token for tool %s (hash: %s)", req.Params.Name, tokenHashPreview)
-
-			// Validate token using configured provider (with request context for timeout/cancellation)
-			user, err := s.validator.ValidateToken(ctx, tokenString)
+			user, err := s.ValidateTokenCached(ctx, tokenString)
 			if err != nil {
 				s.logger.Error("Token validation failed for tool %s: %v", req.Params.Name, err)
-				return nil, fmt.Errorf("authentication failed: %w", err)
+				return nil, err
 			}
-
-			// Cache the validation result (expire in 5 minutes)
-			expiresAt := time.Now().Add(5 * time.Minute)
-			s.cache.setCachedToken(tokenHash, user, expiresAt)
-
-			// Add user to context for downstream handlers
 			ctx = context.WithValue(ctx, userContextKey, user)
-			s.logger.Info("Authenticated user %s for tool: %s (cached for 5 minutes)", user.Username, req.Params.Name)
-
+			s.logger.Info("Authenticated user %s for tool: %s", user.Username, req.Params.Name)
 			return next(ctx, req)
 		}
 	}

oauth.go

@@ -104,9 +104,9 @@ func (s *Server) RegisterHandlers(mux *http.ServeMux) {
 // This is the core validation method that SDK adapters can use.
 //
 // The method:
-//  1. Checks token cache (5-minute TTL)
+//  1. Checks token cache (keyed by SHA-256 hash)
 //  2. Validates token using configured provider if not cached
-//  3. Caches validation result for future requests
+//  3. Caches validation result until JWT exp (minus TokenExpiryBuffer)
 //  4. Returns authenticated User or error
 //
 // This method is used internally by both WrapHandler and adapter middleware.
@@ -120,16 +120,23 @@ func (s *Server) ValidateTokenCached(ctx context.Context, token string) (*User,
 
 	s.logger.Info("Validating token (hash: %s...)", tokenHash[:16])
 
-	user, err := s.validator.ValidateToken(ctx, token)
+	user, expiry, err := s.validator.ValidateToken(ctx, token)
 	if err != nil {
 		s.logger.Error("Token validation failed: %v", err)
 		return nil, fmt.Errorf("authentication failed: %w", err)
 	}
 
-	expiresAt := time.Now().Add(5 * time.Minute)
-	s.cache.setCachedToken(tokenHash, user, expiresAt)
+	if s.config != nil {
+		expiry = expiry.Add(-s.config.TokenExpiryBuffer)
+	}
+	if time.Now().After(expiry) {
+		s.logger.Info("Token rejected: expires within buffer (hash: %s...)", tokenHash[:16])
+		return nil, fmt.Errorf("authentication failed: token expired or expiring too soon")
+	}
+
+	s.cache.setCachedToken(tokenHash, user, expiry)
 
-	s.logger.Info("Authenticated user %s (cached for 5 minutes)", user.Username)
+	s.logger.Info("Authenticated user %s (cache expires at %s)", user.Username, expiry.Format(time.RFC3339))
 	return user, nil
 }