add support for arbitrary token validation funcs

Author: WhammyLeaf

config.go

@@ -5,6 +5,7 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/golang-jwt/jwt/v5"
 	"github.com/tuannvm/oauth-mcp-proxy/provider"
 )
 
@@ -35,10 +36,11 @@ type Config struct {
 	// integrate with your application's logging system (e.g., zap, logrus).
 	Logger Logger
 
-	// Validation skip configuration
-	SkipIssuerCheck   bool
-	SkipAudienceCheck bool
-	SkipExpiryCheck   bool
+	// Token validation configuration
+	SkipIssuerCheck      bool
+	SkipAudienceCheck    bool
+	SkipExpiryCheck      bool
+	TokenValidationFuncs []func(claims jwt.MapClaims) error
 }
 
 // Validate validates the configuration
@@ -127,14 +129,15 @@ func SetupOAuth(cfg *Config) (provider.TokenValidator, error) {
 func createValidator(cfg *Config, logger Logger) (provider.TokenValidator, error) {
 	// Convert root Config to provider.Config
 	providerCfg := &provider.Config{
-		Provider:          cfg.Provider,
-		Issuer:            cfg.Issuer,
-		Audience:          cfg.Audience,
-		JWTSecret:         cfg.JWTSecret,
-		Logger:            logger,
-		SkipIssuerCheck:   cfg.SkipIssuerCheck,
-		SkipAudienceCheck: cfg.SkipAudienceCheck,
-		SkipExpiryCheck:   cfg.SkipExpiryCheck,
+		Provider:            cfg.Provider,
+		Issuer:              cfg.Issuer,
+		Audience:            cfg.Audience,
+		JWTSecret:           cfg.JWTSecret,
+		Logger:              logger,
+		SkipIssuerCheck:     cfg.SkipIssuerCheck,
+		SkipAudienceCheck:   cfg.SkipAudienceCheck,
+		SkipExpiryCheck:     cfg.SkipExpiryCheck,
+		TokenValidatorFuncs: cfg.TokenValidationFuncs,
 	}
 
 	var validator provider.TokenValidator

provider/provider.go

@@ -30,14 +30,15 @@ type Logger interface {
 
 // Config holds OAuth configuration (subset needed by provider)
 type Config struct {
-	Provider          string
-	Issuer            string
-	Audience          string
-	JWTSecret         []byte
-	Logger            Logger
-	SkipIssuerCheck   bool
-	SkipAudienceCheck bool
-	SkipExpiryCheck   bool
+	Provider            string
+	Issuer              string
+	Audience            string
+	JWTSecret           []byte
+	Logger              Logger
+	SkipIssuerCheck     bool
+	SkipAudienceCheck   bool
+	SkipExpiryCheck     bool
+	TokenValidatorFuncs []func(claims jwt.MapClaims) error
 }
 
 // TokenValidator interface for OAuth token validation
@@ -55,10 +56,11 @@ type HMACValidator struct {
 
 // OIDCValidator validates JWT tokens using OIDC/JWKS (Okta, Google, Azure)
 type OIDCValidator struct {
-	verifier *oidc.IDTokenVerifier
-	provider *oidc.Provider
-	audience string
-	logger   Logger
+	verifier            *oidc.IDTokenVerifier
+	provider            *oidc.Provider
+	audience            string
+	TokenValidatorFuncs []func(claims jwt.MapClaims) error
+	logger              Logger
 }
 
 // Initialize sets up the HMAC validator with JWT secret and audience
@@ -215,6 +217,7 @@ func (v *OIDCValidator) Initialize(cfg *Config) error {
 
 	v.provider = provider
 	v.verifier = verifier
+	v.TokenValidatorFuncs = cfg.TokenValidatorFuncs
 	return nil
 }
 
@@ -263,6 +266,14 @@ func (v *OIDCValidator) ValidateToken(ctx context.Context, tokenString string) (
 		return nil, fmt.Errorf("audience validation failed: %w", err)
 	}
 
+	// Run extra validation functions
+	for _, fn := range v.TokenValidatorFuncs {
+		err := fn(rawClaims)
+		if err != nil {
+			return nil, fmt.Errorf("validation function failed with error: %w", err)
+		}
+	}
+
 	return &User{
 		Subject:  claims.Subject,
 		Username: claims.PreferredUsername,