Feat: OAuth 2.1 Authorization Framework

.golangci.yml

@@ -1,15 +1,11 @@
 run:
   timeout: 5m
-  skip-files:
-    - ".*\\.pb\\.go$"
   allow-parallel-runners: true
   go: '1.24'
 
 linters-settings:
   errcheck:
     check-type-assertions: true
-  govet:
-    check-shadowing: true
   gofmt:
     simplify: true
   gocyclo:
@@ -42,6 +38,8 @@ linters:
     - unused
 
 issues:
+  exclude-files:
+    - ".*\\.pb\\.go$"
   exclude-rules:
     - path: _test\.go
       linters:

docs/mcp-alignment-tasks.md

@@ -0,0 +1,107 @@
+# MCP 2025-03-26 Alignment Tasks
+
+This document outlines the tasks needed to align Cortex with the latest Model Context Protocol (MCP) specification (2025-03-26). The recent MCP update introduces several significant changes that require implementation in our codebase.
+
+## Major Changes in MCP 2025-03-26
+
+The latest MCP specification includes the following key updates:
+
+1. **Authorization Framework**: A comprehensive authorization framework based on OAuth 2.1
+2. **Streamable HTTP Transport**: Replacement of HTTP+SSE transport with a more flexible Streamable HTTP transport
+3. **JSON-RPC Batching**: Support for batched requests via JSON-RPC
+4. **Tool Annotations**: Comprehensive tool annotations for better describing tool behavior
+
+## Implementation Tasks
+
+### 1. OAuth 2.1 Authorization Framework
+
+- [x] **Task 1.1**: Define OAuth 2.1 authentication interfaces in `pkg/server/auth.go`
+- [x] **Task 1.2**: Implement OAuth 2.1 token validation and verification
+- [x] **Task 1.3**: Create middleware for OAuth token validation
+- [x] **Task 1.4**: Add scope-based permission system for tool access
+- [x] **Task 1.5**: Update PocketBase integration to support OAuth 2.1
+- [x] **Task 1.6**: Create documentation for authorization setup and configuration
+- [x] **Task 1.7**: Implement tests for OAuth authentication flows
+
+### 2. Streamable HTTP Transport
+
+- [ ] **Task 2.1**: Define interface for streamable HTTP transport in `pkg/server/transport.go`
+- [ ] **Task 2.2**: Implement streamable response writer
+- [ ] **Task 2.3**: Update existing SSE implementation to use new transport
+- [ ] **Task 2.4**: Add support for streaming binary data (for audio support)
+- [ ] **Task 2.5**: Create adapters for different streaming protocols
+- [ ] **Task 2.6**: Update client notification system to use new transport
+- [ ] **Task 2.7**: Implement tests for streamable transport
+- [ ] **Task 2.8**: Update documentation with new transport details
+
+### 3. JSON-RPC Batching
+
+- [ ] **Task 3.1**: Add batch request handler in `pkg/server/jsonrpc.go`
+- [ ] **Task 3.2**: Implement concurrent execution of batched requests
+- [ ] **Task 3.3**: Add result collation for batch responses
+- [ ] **Task 3.4**: Create error handling for partial batch failures
+- [ ] **Task 3.5**: Add batch size limits and validation
+- [ ] **Task 3.6**: Update HTTP handlers to support batch endpoints
+- [ ] **Task 3.7**: Implement tests for batch processing
+- [ ] **Task 3.8**: Update documentation with batching examples
+
+### 4. Tool Annotations
+
+- [ ] **Task 4.1**: Extend tool definition schema in `pkg/tools` to include annotations
+- [ ] **Task 4.2**: Add read-only/destructive operation flags
+- [ ] **Task 4.3**: Implement permission level requirements based on annotations
+- [ ] **Task 4.4**: Update tool registration to include annotation metadata
+- [ ] **Task 4.5**: Add validation for tool annotations
+- [ ] **Task 4.6**: Update integration examples to demonstrate annotations
+- [ ] **Task 4.7**: Create tests for tool annotation handling
+- [ ] **Task 4.8**: Update documentation with annotation guidelines
+
+### 5. Other Schema Updates
+
+- [ ] **Task 5.1**: Add `message` field to `ProgressNotification` in notification system
+- [ ] **Task 5.2**: Implement audio data support in content types
+- [ ] **Task 5.3**: Add `completions` capability flag for argument autocompletion
+- [ ] **Task 5.4**: Update schema validation to match latest MCP specification
+- [ ] **Task 5.5**: Update all example tools to use the new schema features
+- [ ] **Task 5.6**: Create tests for new schema elements
+- [ ] **Task 5.7**: Update documentation with new schema examples
+
+## Migration Strategy
+
+To ensure a smooth transition to the new MCP specification:
+
+1. **Backward Compatibility**: Maintain support for the previous specification (2024-11-05) during transition
+2. **Phased Implementation**: Implement changes in the following order:
+   - Tool Annotations (lowest impact)
+   - Schema Updates
+   - JSON-RPC Batching
+   - Streamable HTTP Transport
+   - OAuth 2.1 Authorization Framework (highest impact)
+3. **Version Flagging**: Add version headers to allow clients to request specific protocol versions
+4. **Documentation Updates**: Keep documentation in sync with implementation progress
+
+## Testing Approach
+
+For each implemented task:
+
+1. Write unit tests before implementation (TDD approach)
+2. Create integration tests that verify compatibility with the specification
+3. Implement example clients that exercise the new functionality
+4. Verify backward compatibility with existing clients
+
+## Timeline
+
+The estimated completion timeline for aligning with MCP 2025-03-26:
+
+- Phase 1 (Tool Annotations & Schema Updates): 2 weeks
+- Phase 2 (JSON-RPC Batching): 1 week
+- Phase 3 (Streamable HTTP Transport): 2 weeks
+- Phase 4 (OAuth 2.1 Framework): 3 weeks
+
+Total estimated time: 8 weeks
+
+## Resources
+
+- [MCP 2025-03-26 Specification](https://modelcontextprotocol.io/specification/2025-03-26/)
+- [MCP Changelog](https://modelcontextprotocol.io/specification/2025-03-26/changelog)
+- [OAuth 2.1 Specification](https://oauth.net/2.1/) 

docs/oauth-authorization.md

@@ -0,0 +1,264 @@
+# OAuth 2.1 Authorization Setup and Configuration
+
+This document explains how to set up and configure OAuth 2.1 authorization in Cortex. The authorization framework provides secure access control to Cortex API endpoints and tools, following OAuth 2.1 standards.
+
+## Overview
+
+Cortex implements a comprehensive OAuth 2.1 authorization framework with the following features:
+
+- JWT token validation
+- Scope-based access control
+- Tool-specific permissions
+- Multiple token extraction methods (header, query parameter, cookie)
+- Integration with external identity providers
+
+## Basic Setup
+
+### Step 1: Configure OAuth Settings
+
+Create an `OAuthConfig` with your authorization settings:
+
+```go
+config := &server.OAuthConfig{
+    Issuer:            "https://auth.example.com",     // Your OAuth issuer URL
+    Audience:          []string{"cortex-api"},         // Expected audience values
+    JWKSUrl:           "https://auth.example.com/.well-known/jwks.json", // JWKS endpoint
+    RequiredScopes:    []string{"cortex:api"},         // Global required scopes
+    TokenLookupScheme: "header,query",                 // Where to look for tokens
+    TokenHeaderName:   "Authorization",                // Header name (default)
+    TokenQueryParam:   "access_token",                 // Query parameter name
+}
+```
+
+### Step 2: Create a Token Validator
+
+Choose one of these validator implementations based on your needs:
+
+#### JWT Token Validator (recommended)
+
+```go
+// Create a key provider that fetches keys from your JWKS endpoint
+keyProvider := server.NewJWKSKeyProvider(config.JWKSUrl)
+
+// Create a JWT validator
+validator := server.NewJWTTokenValidator(config, keyProvider)
+
+// Create OAuth middleware
+middleware := server.NewOAuthMiddlewareWithConfig(validator, config)
+```
+
+#### OAuth 2.0 Introspection Validator
+
+```go
+// Create an introspector for RFC 7662 token introspection
+introspector := server.NewStandardIntrospector(
+    "https://auth.example.com/oauth/introspect", 
+    "client_id", 
+    "client_secret"
+)
+
+// Create an introspection validator
+validator := server.NewIntrospectionTokenValidator(config, introspector)
+
+// Create OAuth middleware
+middleware := server.NewOAuthMiddlewareWithConfig(validator, config)
+```
+
+### Step 3: Apply OAuth Middleware
+
+Apply the OAuth middleware to your HTTP handlers:
+
+```go
+// Create your handler
+handler := http.HandlerFunc(yourHandlerFunc)
+
+// Wrap with OAuth middleware
+protectedHandler := middleware.Middleware(handler)
+
+// Use in your HTTP server
+http.Handle("/api/protected", protectedHandler)
+```
+
+## Scope-Based Access Control
+
+### Defining Scopes
+
+Scopes are strings that represent permissions. In Cortex, we use a hierarchical naming convention:
+
+- `cortex:api` - General API access
+- `cortex:tool:read` - Read access to all tools
+- `cortex:tool:execute:{toolName}` - Execute permission for a specific tool
+
+### Requiring Scopes for Endpoints
+
+You can protect endpoints with scope requirements:
+
+```go
+// Require a single scope
+handler := middleware.RequireScope("cortex:api", yourHandler)
+
+// Require any one of multiple scopes
+handler := middleware.RequireAnyScope([]string{"cortex:admin", "cortex:tool:read"}, yourHandler)
+
+// Require all specified scopes
+handler := middleware.RequireAllScopes([]string{"cortex:api", "cortex:tool:read"}, yourHandler)
+```
+
+## Tool Permissions
+
+Cortex provides a dedicated permission system for tools with three permission types:
+
+- `ToolPermissionRead`: Read tool metadata
+- `ToolPermissionWrite`: Modify tool configuration
+- `ToolPermissionExecute`: Execute the tool
+
+### Setting Up Tool Permissions
+
+```go
+// Create tool permissions with the OAuth middleware
+toolPermissions := server.NewToolPermissions(middleware)
+
+// Protect a tool endpoint
+handler := toolPermissions.RequireToolPermission(
+    "calculator",                // Tool name
+    server.ToolPermissionExecute, // Permission type
+    yourToolHandler              // Handler to protect
+)
+```
+
+### Tool Permission Scopes
+
+Tool permissions use the following scope format:
+
+- `cortex:tool:{permission}:{toolName}`
+
+Examples:
+- `cortex:tool:read` - Global read access to all tools
+- `cortex:tool:execute:calculator` - Execute permission for the calculator tool
+- `cortex:tool:write:weather` - Write permission for the weather tool
+
+## PocketBase Integration
+
+If you're using PocketBase, you can set up OAuth with the Cortex plugin:
+
+```go
+// Create the plugin
+plugin := pocketbase.NewCortexPlugin()
+
+// Set up OAuth
+validator := CreateYourValidator() // See validator setup above
+middleware := server.NewOAuthMiddlewareWithConfig(validator, config)
+
+// Add OAuth to the plugin
+plugin.WithOAuth(middleware).WithOAuthConfig(config)
+
+// Register with PocketBase
+plugin.RegisterWithPocketBase(app)
+```
+
+## Accessing Token Claims
+
+In your HTTP handlers, you can access token claims from the context:
+
+```go
+func yourHandler(w http.ResponseWriter, r *http.Request) {
+    // Get token claims from context
+    claims, ok := server.GetTokenClaimsFromContext(r.Context())
+    if !ok {
+        http.Error(w, "Unauthorized", http.StatusUnauthorized)
+        return
+    }
+
+    // Use claims information
+    userID := claims.Subject
+    scopes := claims.Scopes
+
+    // Check permissions manually if needed
+    if !hasPermission(claims) {
+        http.Error(w, "Forbidden", http.StatusForbidden)
+        return
+    }
+
+    // Continue with authorized operation...
+}
+```
+
+## Custom Scope Checking
+
+If you need custom scope checking logic, you can implement the `ScopeChecker` interface:
+
+```go
+type CustomScopeChecker struct {
+    // Any fields you need
+}
+
+func (c *CustomScopeChecker) HasScope(claims *server.TokenClaims, requiredScope string) bool {
+    // Your custom logic here
+    return customScopeCheckLogic(claims, requiredScope)
+}
+
+func (c *CustomScopeChecker) HasAnyScope(claims *server.TokenClaims, requiredScopes ...string) bool {
+    // Your custom logic here
+    return customAnyScopeCheckLogic(claims, requiredScopes)
+}
+
+func (c *CustomScopeChecker) HasAllScopes(claims *server.TokenClaims, requiredScopes ...string) bool {
+    // Your custom logic here
+    return customAllScopesCheckLogic(claims, requiredScopes)
+}
+
+// Then use your custom checker:
+middleware.WithScopeChecker(&CustomScopeChecker{})
+```
+
+## Security Considerations
+
+1. **Token Validation**: Always validate tokens for integrity, expiration, issuer, and audience.
+2. **HTTPS**: Use HTTPS for all communication to protect tokens.
+3. **Proper Scopes**: Grant minimal necessary scopes to each client.
+4. **Token Storage**: Securely store tokens client-side, and never in localStorage.
+5. **Token Expiration**: Use short-lived access tokens with refresh token rotation.
+6. **CORS**: Configure CORS properly to restrict access to trusted domains.
+
+## Troubleshooting
+
+### Common Issues
+
+1. **401 Unauthorized**: Indicates invalid or expired token, or missing token.
+2. **403 Forbidden**: Valid token but insufficient scopes for the requested action.
+3. **JWKS Key Issues**: If the key ID (kid) in the token doesn't match any key in the JWKS.
+
+### Debugging Tips
+
+- Use the `Authorization` header debug logs for token extraction issues.
+- Check token expiration and issuer if validation fails.
+- Verify the token has the proper scopes for the requested action.
+- Ensure your JWKS endpoint is accessible and returns the correct keys.
+
+## Example Configuration
+
+Here's a complete example of setting up OAuth 2.1 with JWT validation:
+
+```go
+func SetupOAuth() *server.OAuthMiddleware {
+    // Create OAuth configuration
+    config := &server.OAuthConfig{
+        Issuer:            "https://auth.example.com",
+        Audience:          []string{"cortex-api"},
+        JWKSUrl:           "https://auth.example.com/.well-known/jwks.json",
+        RequiredScopes:    []string{"cortex:api"},
+        TokenLookupScheme: "header,query,cookie",
+        TokenHeaderName:   "Authorization",
+        TokenQueryParam:   "access_token",
+    }
+
+    // Create key provider
+    keyProvider := server.NewJWKSKeyProvider(config.JWKSUrl)
+
+    // Create JWT validator
+    validator := server.NewJWTTokenValidator(config, keyProvider)
+
+    // Create and return OAuth middleware
+    return server.NewOAuthMiddlewareWithConfig(validator, config)
+}
+```