feat: Implement OAuth 2.1 Resource Server for JWT authentication, including JWKS caching, discovery endpoints, and configuration updates.

Author: om252345

.gitignore

@@ -17,3 +17,4 @@ vendor/
 # User Specific configuration
 agentgate.yaml
 .env
+chinook.db

README.md

@@ -49,19 +49,22 @@ Relying on "system prompts" for security is a guaranteed way to get your databas
 
 ## 🔥 Core Features
 
-### 1. Semantic RBAC & Parameter Sandbox 🔒
+### 1. Centralized OAuth 2.1 Resource Server 🛡️
+Stop writing custom OAuth logic for every single MCP tool you build! AgentGate acts as a spec-compliant OAuth 2.1 Resource Server. It validates JWTs, fetches JWKS keys (with background rotation), and bounces unauthenticated AI clients with `WWW-Authenticate` headers — completely decoupling auth from your business logic.
+
+### 2. Semantic RBAC & Parameter Sandbox 🔒
 Whitelist exactly which tools an agent can use. Go deeper with **regex rules on the parameters themselves** — e.g., the agent can only read files ending in `.log`, or can only `SELECT` but never `DELETE`.
 
-### 2. Human-in-the-Loop (HITL) ⏸️
+### 3. Human-in-the-Loop (HITL) ⏸️
 Automatically pause high-risk tool execution. AgentGate intercepts the request, pings your **Slack** or a CLI webhook, and physically holds the HTTP connection open until a human clicks **Approve** or **Deny**.
 
-### 3. Runaway Loop Breaker (Rate Limiting) ⏱️
+### 4. Runaway Loop Breaker (Rate Limiting) ⏱️
 Defeat hallucination loops. Cap tool executions per minute globally or **per MCP server**. If an agent spams a function, it instantly receives `HTTP 429`.
 
-### 4. The IPC Panic Button 🛑
+### 5. The IPC Panic Button 🛑
 If an agent goes completely rogue, type `agentgate service pause` in your terminal. This uses an isolated Unix Domain Socket to **instantly sever all autonomous tool execution** with a `503`, without exposing an admin endpoint to the network.
 
-### 5. stdio → HTTP Bridge 🌉
+### 6. stdio → HTTP Bridge 🌉
 MCP natively uses local `stdio`. AgentGate translates this to standard **HTTP/SSE**, letting you run tools in an isolated container or VPC while the LLM client stays local.
 
 ---

auth/context.go

@@ -0,0 +1,44 @@
+package auth
+
+import "context"
+
+// ctxKey is an unexported type for context keys in this package.
+// Using a package-local type prevents collisions with other packages
+// that also use context.WithValue.
+type ctxKey int
+
+const (
+	ctxKeySub    ctxKey = iota // JWT "sub" claim (user/client identity)
+	ctxKeyScopes               // JWT "scope" claim (space-delimited, split into []string)
+)
+
+// ContextWithClaims returns a new context enriched with the JWT subject and scopes.
+func ContextWithClaims(ctx context.Context, sub string, scopes []string) context.Context {
+	ctx = context.WithValue(ctx, ctxKeySub, sub)
+	ctx = context.WithValue(ctx, ctxKeyScopes, scopes)
+	return ctx
+}
+
+// SubFromContext extracts the JWT subject (sub) from a request context.
+// Returns an empty string if not present (i.e., OAuth2 is disabled).
+func SubFromContext(ctx context.Context) string {
+	v, _ := ctx.Value(ctxKeySub).(string)
+	return v
+}
+
+// ScopesFromContext extracts the JWT scopes from a request context.
+// Returns nil if not present.
+func ScopesFromContext(ctx context.Context) []string {
+	v, _ := ctx.Value(ctxKeyScopes).([]string)
+	return v
+}
+
+// HasScope returns true if the given scope is present in the request context.
+func HasScope(ctx context.Context, required string) bool {
+	for _, s := range ScopesFromContext(ctx) {
+		if s == required {
+			return true
+		}
+	}
+	return false
+}

auth/jwks.go

@@ -0,0 +1,195 @@
+package auth
+
+import (
+	"crypto"
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rsa"
+	"encoding/base64"
+	"encoding/json"
+	"fmt"
+	"log"
+	"math/big"
+	"net/http"
+	"sync"
+	"time"
+)
+
+// jwksResponse is the raw JSON structure from the /.well-known/jwks.json endpoint.
+type jwksResponse struct {
+	Keys []jwk `json:"keys"`
+}
+
+// jwk represents a single JSON Web Key (RSA or EC public key fields).
+type jwk struct {
+	Kid string `json:"kid"` // Key ID
+	Kty string `json:"kty"` // Key type ("RSA" or "EC")
+	Alg string `json:"alg"` // Algorithm (e.g. "RS256", "ES256")
+	
+	// RSA fields:
+	N   string `json:"n"`   // RSA modulus (base64url encoded)
+	E   string `json:"e"`   // RSA exponent (base64url encoded)
+	
+	// EC fields:
+	Crv string `json:"crv"` // Curve (e.g. "P-256")
+	X   string `json:"x"`   // EC X coordinate
+	Y   string `json:"y"`   // EC Y coordinate
+}
+
+// JWKSCache fetches and caches public keys from a JWKS endpoint.
+// It refreshes the cache periodically to handle key rotation.
+type JWKSCache struct {
+	jwksURL string
+	mu      sync.RWMutex
+	keys    map[string]crypto.PublicKey // kid → public key
+}
+
+// NewJWKSCache creates a JWKSCache and performs an initial key fetch.
+// It also spawns a background goroutine that refreshes the cache every
+// refreshInterval to transparently handle key rotation at the IdP.
+func NewJWKSCache(jwksURL string, refreshInterval time.Duration) (*JWKSCache, error) {
+	c := &JWKSCache{
+		jwksURL: jwksURL,
+		keys:    make(map[string]crypto.PublicKey),
+	}
+
+	// Perform the initial fetch synchronously so we fail fast on misconfiguration.
+	if err := c.refresh(); err != nil {
+		return nil, fmt.Errorf("initial JWKS fetch from %s failed: %w", jwksURL, err)
+	}
+
+	// Background refresh goroutine — handles key rotation.
+	go func() {
+		ticker := time.NewTicker(refreshInterval)
+		defer ticker.Stop()
+		for range ticker.C {
+			if err := c.refresh(); err != nil {
+				log.Printf("[JWKS] Background refresh failed (will retry): %v", err)
+			} else {
+				log.Printf("[JWKS] Key cache refreshed from %s", jwksURL)
+			}
+		}
+	}()
+
+	return c, nil
+}
+
+// GetKey returns the public key for the given key ID (kid).
+// Returns an error if the kid is not found in the cache.
+func (c *JWKSCache) GetKey(kid string) (crypto.PublicKey, error) {
+	c.mu.RLock()
+	key, ok := c.keys[kid]
+	c.mu.RUnlock()
+
+	if !ok {
+		// Kid not found — try one eager refresh in case the IdP just rotated keys.
+		log.Printf("[JWKS] kid %q not in cache — attempting eager refresh", kid)
+		if err := c.refresh(); err != nil {
+			return nil, fmt.Errorf("kid %q not found and refresh failed: %w", kid, err)
+		}
+		c.mu.RLock()
+		key, ok = c.keys[kid]
+		c.mu.RUnlock()
+		if !ok {
+			return nil, fmt.Errorf("kid %q not found in JWKS after refresh", kid)
+		}
+	}
+	return key, nil
+}
+
+// refresh fetches the JWKS endpoint and atomically replaces the key cache.
+func (c *JWKSCache) refresh() error {
+	resp, err := http.Get(c.jwksURL) //nolint:noctx
+	if err != nil {
+		return fmt.Errorf("GET %s: %w", c.jwksURL, err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("GET %s: unexpected HTTP %d", c.jwksURL, resp.StatusCode)
+	}
+
+	var jwks jwksResponse
+	if err := json.NewDecoder(resp.Body).Decode(&jwks); err != nil {
+		return fmt.Errorf("decode JWKS: %w", err)
+	}
+
+	newKeys := make(map[string]crypto.PublicKey, len(jwks.Keys))
+	for _, k := range jwks.Keys {
+		if k.Kty == "RSA" {
+			pub, err := parseRSAPublicKey(k)
+			if err != nil {
+				log.Printf("[JWKS] Skipping RSA key kid=%q: %v", k.Kid, err)
+				continue
+			}
+			newKeys[k.Kid] = pub
+		} else if k.Kty == "EC" && k.Crv == "P-256" {
+			pub, err := parseECPublicKey(k)
+			if err != nil {
+				log.Printf("[JWKS] Skipping EC key kid=%q: %v", k.Kid, err)
+				continue
+			}
+			newKeys[k.Kid] = pub
+		} else {
+			continue // Skip unsupported key types/curves
+		}
+	}
+
+	if len(newKeys) == 0 {
+		return fmt.Errorf("JWKS response contained no supported RSA or EC keys")
+	}
+
+	c.mu.Lock()
+	c.keys = newKeys
+	c.mu.Unlock()
+
+	log.Printf("[JWKS] Loaded %d key(s) from %s", len(newKeys), c.jwksURL)
+	return nil
+}
+
+// parseECPublicKey converts a JSON Web Key into a Go *ecdsa.PublicKey
+// using only the standard library.
+func parseECPublicKey(k jwk) (*ecdsa.PublicKey, error) {
+	xBytes, err := base64.RawURLEncoding.DecodeString(k.X)
+	if err != nil {
+		return nil, fmt.Errorf("decode X: %w", err)
+	}
+
+	yBytes, err := base64.RawURLEncoding.DecodeString(k.Y)
+	if err != nil {
+		return nil, fmt.Errorf("decode Y: %w", err)
+	}
+
+	return &ecdsa.PublicKey{
+		Curve: elliptic.P256(),
+		X:     new(big.Int).SetBytes(xBytes),
+		Y:     new(big.Int).SetBytes(yBytes),
+	}, nil
+}
+
+// parseRSAPublicKey converts a JSON Web Key into a Go *rsa.PublicKey
+// using only the standard library (no external JWT packages).
+func parseRSAPublicKey(k jwk) (*rsa.PublicKey, error) {
+	// Decode the base64url-encoded modulus (N)
+	nBytes, err := base64.RawURLEncoding.DecodeString(k.N)
+	if err != nil {
+		return nil, fmt.Errorf("decode N: %w", err)
+	}
+
+	// Decode the base64url-encoded exponent (E)
+	eBytes, err := base64.RawURLEncoding.DecodeString(k.E)
+	if err != nil {
+		return nil, fmt.Errorf("decode E: %w", err)
+	}
+
+	// Convert E bytes to int
+	var eInt int
+	for _, b := range eBytes {
+		eInt = eInt<<8 | int(b)
+	}
+
+	return &rsa.PublicKey{
+		N: new(big.Int).SetBytes(nBytes),
+		E: eInt,
+	}, nil
+}

auth/middleware.go

@@ -0,0 +1,100 @@
+package auth
+
+import (
+	"fmt"
+	"log"
+	"net/http"
+	"strings"
+
+	"github.com/agentgate/agentgate/config"
+)
+
+// JWTAuthMiddleware is the OAuth 2.1 Resource Server interceptor.
+//
+// Behavior:
+//   - If oauth2.enabled = false: passes through immediately (backward compatible).
+//   - No/invalid Authorization header: HTTP 401 + WWW-Authenticate challenge.
+//   - Valid JWT: enriches context with sub+scopes, strips Authorization header,
+//     optionally injects X-AgentGate-User header for the upstream MCP server.
+//   - Passes to next handler.
+func JWTAuthMiddleware(cfg *config.Config, cache *JWKSCache, next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Fast path: OAuth2 not enabled — skip this entire layer.
+		if !cfg.OAuth2.Enabled {
+			next.ServeHTTP(w, r)
+			return
+		}
+
+		// ── Extract Bearer token ──────────────────────────────────────────────
+		authHeader := r.Header.Get("Authorization")
+		if authHeader == "" || !strings.HasPrefix(authHeader, "Bearer ") {
+			writeWWWAuthenticate(w, cfg.OAuth2.ResourceMetadata, "missing_token", "Authorization: Bearer <token> is required")
+			return
+		}
+		tokenString := strings.TrimPrefix(authHeader, "Bearer ")
+
+		// ── Validate JWT ──────────────────────────────────────────────────────
+		sub, scopes, err := ValidateJWT(tokenString, cache, cfg.OAuth2.Issuer, cfg.OAuth2.Audience)
+		if err != nil {
+			log.Printf("[JWTAuth] [ERROR] Token validation failed from %s: %v", r.RemoteAddr, err)
+			errCode := classifyError(err)
+			writeWWWAuthenticate(w, cfg.OAuth2.ResourceMetadata, errCode, err.Error())
+			return
+		}
+
+		log.Printf("[JWTAuth] Token valid — sub=%q scopes=%v from %s", sub, scopes, r.RemoteAddr)
+
+		// ── Enrich context with claims ────────────────────────────────────────
+		ctx := ContextWithClaims(r.Context(), sub, scopes)
+		r = r.WithContext(ctx)
+
+		// ── Strip upstream Authorization header (don't leak the JWT) ─────────
+		// We replace the original header with a custom user identity header
+		// so the upstream MCP server knows who made the request without having
+		// access to the raw credential.
+		r.Header.Del("Authorization")
+		if cfg.OAuth2.InjectUserHeader {
+			r.Header.Set("X-AgentGate-User", sub)
+			if len(scopes) > 0 {
+				r.Header.Set("X-AgentGate-Scopes", strings.Join(scopes, " "))
+			}
+		}
+
+		next.ServeHTTP(w, r)
+	})
+}
+
+// writeWWWAuthenticate sends an RFC 6750-compliant 401 response.
+//
+// Format per MCP spec:
+//
+//	WWW-Authenticate: Bearer realm="mcp", resource_metadata="<url>",
+//	                          error="<code>", error_description="<msg>"
+func writeWWWAuthenticate(w http.ResponseWriter, resourceMetadata, errCode, errDesc string) {
+	challenge := fmt.Sprintf(
+		`Bearer realm="mcp", resource_metadata="%s", error="%s", error_description="%s"`,
+		resourceMetadata, errCode, errDesc,
+	)
+	w.Header().Set("WWW-Authenticate", challenge)
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(http.StatusUnauthorized)
+	w.Write([]byte(`{"error":"unauthorized","error_description":"` + errDesc + `"}`))
+}
+
+// classifyError maps a validation error to an RFC 6750 error code string.
+func classifyError(err error) string {
+	msg := err.Error()
+	if strings.Contains(msg, "expired") {
+		return "invalid_token"
+	}
+	if strings.Contains(msg, "signature") {
+		return "invalid_token"
+	}
+	if strings.Contains(msg, "issuer") || strings.Contains(msg, "audience") {
+		return "invalid_token"
+	}
+	if strings.Contains(msg, "malformed") || strings.Contains(msg, "algorithm") {
+		return "invalid_token"
+	}
+	return "invalid_token"
+}