feat(auth): implement OAuth 2.0 authorization code flow and dynamic client registration

- Add OAuth 2.0 support with authorization code flow and dynamic client registration.
- Introduce new handlers for OAuth metadata, client registration, authorization, and token issuance.
- Enhance authentication middleware to support OAuth client credentials.
- Create in-memory stores for authorization codes and tokens.
- Update configuration to include OAuth client details.
- Ensure validation checks for OAuth clients in the configuration.

internal/auth/auth_code_store.go

@@ -0,0 +1,76 @@
+package auth
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"sync"
+	"time"
+)
+
+const authCodeTTL = 10 * time.Minute
+
+// AuthCode holds a pending authorization code and its associated PKCE data.
+type AuthCode struct {
+	ClientID            string
+	RedirectURI         string
+	Scope               string
+	CodeChallenge       string
+	CodeChallengeMethod string
+	KeyID               string
+	ExpiresAt           time.Time
+}
+
+// AuthCodeStore issues single-use authorization codes for the OAuth 2.0
+// Authorization Code flow.
+type AuthCodeStore struct {
+	mu    sync.Mutex
+	codes map[string]AuthCode
+}
+
+func NewAuthCodeStore() *AuthCodeStore {
+	s := &AuthCodeStore{codes: make(map[string]AuthCode)}
+	go s.sweepLoop()
+	return s
+}
+
+// Issue stores the entry and returns the raw authorization code.
+func (s *AuthCodeStore) Issue(entry AuthCode) (string, error) {
+	b := make([]byte, 16)
+	if _, err := rand.Read(b); err != nil {
+		return "", err
+	}
+	raw := hex.EncodeToString(b)
+	entry.ExpiresAt = time.Now().Add(authCodeTTL)
+	s.mu.Lock()
+	s.codes[raw] = entry
+	s.mu.Unlock()
+	return raw, nil
+}
+
+// Consume validates and removes the code, returning the associated entry.
+func (s *AuthCodeStore) Consume(code string) (AuthCode, bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	entry, ok := s.codes[code]
+	if !ok || time.Now().After(entry.ExpiresAt) {
+		delete(s.codes, code)
+		return AuthCode{}, false
+	}
+	delete(s.codes, code)
+	return entry, true
+}
+
+func (s *AuthCodeStore) sweepLoop() {
+	ticker := time.NewTicker(5 * time.Minute)
+	defer ticker.Stop()
+	for range ticker.C {
+		now := time.Now()
+		s.mu.Lock()
+		for code, entry := range s.codes {
+			if now.After(entry.ExpiresAt) {
+				delete(s.codes, code)
+			}
+		}
+		s.mu.Unlock()
+	}
+}

internal/auth/dynamic_client_store.go

@@ -0,0 +1,62 @@
+package auth
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"sync"
+	"time"
+)
+
+// DynamicClient holds a dynamically registered OAuth client (RFC 7591).
+type DynamicClient struct {
+	ClientID     string
+	ClientName   string
+	RedirectURIs []string
+	CreatedAt    time.Time
+}
+
+// HasRedirectURI reports whether uri is registered for this client.
+func (c *DynamicClient) HasRedirectURI(uri string) bool {
+	for _, u := range c.RedirectURIs {
+		if u == uri {
+			return true
+		}
+	}
+	return false
+}
+
+// DynamicClientStore holds dynamically registered OAuth clients in memory.
+type DynamicClientStore struct {
+	mu      sync.RWMutex
+	clients map[string]DynamicClient
+}
+
+func NewDynamicClientStore() *DynamicClientStore {
+	return &DynamicClientStore{clients: make(map[string]DynamicClient)}
+}
+
+// Register creates a new client and returns it.
+func (s *DynamicClientStore) Register(name string, redirectURIs []string) (DynamicClient, error) {
+	b := make([]byte, 16)
+	if _, err := rand.Read(b); err != nil {
+		return DynamicClient{}, err
+	}
+	client := DynamicClient{
+		ClientID:     hex.EncodeToString(b),
+		ClientName:   name,
+		RedirectURIs: append([]string(nil), redirectURIs...),
+		CreatedAt:    time.Now(),
+	}
+	s.mu.Lock()
+	s.clients[client.ClientID] = client
+	s.mu.Unlock()
+	return client, nil
+}
+
+// Lookup returns the client for the given client_id.
+func (s *DynamicClientStore) Lookup(clientID string) (DynamicClient, bool) {
+	s.mu.RLock()
+	client, ok := s.clients[clientID]
+	s.mu.RUnlock()
+	return client, ok
+}

internal/auth/keyring_test.go

@@ -39,7 +39,7 @@ func TestMiddlewareAllowsHeaderAuthAndSetsContext(t *testing.T) {
 		t.Fatalf("NewKeyring() error = %v", err)
 	}
 
-	handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		keyID, ok := KeyIDFromContext(r.Context())
 		if !ok || keyID != "client-a" {
 			t.Fatalf("KeyIDFromContext() = (%q, %v), want (client-a, true)", keyID, ok)
@@ -63,7 +63,7 @@ func TestMiddlewareAllowsBearerAuthAndSetsContext(t *testing.T) {
 		t.Fatalf("NewKeyring() error = %v", err)
 	}
 
-	handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		keyID, ok := KeyIDFromContext(r.Context())
 		if !ok || keyID != "client-a" {
 			t.Fatalf("KeyIDFromContext() = (%q, %v), want (client-a, true)", keyID, ok)
@@ -90,7 +90,7 @@ func TestMiddlewarePrefersExplicitHeaderOverBearerAuth(t *testing.T) {
 		t.Fatalf("NewKeyring() error = %v", err)
 	}
 
-	handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		keyID, ok := KeyIDFromContext(r.Context())
 		if !ok || keyID != "client-a" {
 			t.Fatalf("KeyIDFromContext() = (%q, %v), want (client-a, true)", keyID, ok)
@@ -119,7 +119,7 @@ func TestMiddlewareAllowsQueryParamWhenEnabled(t *testing.T) {
 		HeaderName:      "x-brain-key",
 		QueryParam:      "key",
 		AllowQueryParam: true,
-	}, keyring, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	}, keyring, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusNoContent)
 	}))
 
@@ -138,7 +138,7 @@ func TestMiddlewareRejectsMissingOrInvalidKey(t *testing.T) {
 		t.Fatalf("NewKeyring() error = %v", err)
 	}
 
-	handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	handler := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, nil, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		t.Fatal("next handler should not be called")
 	}))
 

internal/auth/middleware.go

@@ -2,6 +2,7 @@ package auth
 
 import (
 	"context"
+	"encoding/base64"
 	"log/slog"
 	"net/http"
 	"strings"
@@ -13,32 +14,77 @@ type contextKey string
 
 const keyIDContextKey contextKey = "auth.key_id"
 
-func Middleware(cfg config.AuthConfig, keyring *Keyring, log *slog.Logger) func(http.Handler) http.Handler {
+func Middleware(cfg config.AuthConfig, keyring *Keyring, oauthRegistry *OAuthRegistry, tokenStore *TokenStore, log *slog.Logger) func(http.Handler) http.Handler {
 	headerName := cfg.HeaderName
 	if headerName == "" {
 		headerName = "x-brain-key"
 	}
-
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			token := extractToken(r, headerName)
-			if token == "" && cfg.AllowQueryParam {
-				token = strings.TrimSpace(r.URL.Query().Get(cfg.QueryParam))
+			// 1. Custom header → keyring only.
+			if keyring != nil {
+				if token := strings.TrimSpace(r.Header.Get(headerName)); token != "" {
+					keyID, ok := keyring.Lookup(token)
+					if !ok {
+						log.Warn("authentication failed", slog.String("remote_addr", r.RemoteAddr))
+						http.Error(w, "invalid API key", http.StatusUnauthorized)
+						return
+					}
+					next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID)))
+					return
+				}
 			}
 
-			if token == "" {
-				http.Error(w, "missing API key", http.StatusUnauthorized)
+			// 2. Bearer token → tokenStore (OAuth), then keyring (API key).
+			if bearer := extractBearer(r); bearer != "" {
+				if tokenStore != nil {
+					if keyID, ok := tokenStore.Lookup(bearer); ok {
+						next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID)))
+						return
+					}
+				}
+				if keyring != nil {
+					if keyID, ok := keyring.Lookup(bearer); ok {
+						next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID)))
+						return
+					}
+				}
+				log.Warn("bearer token rejected", slog.String("remote_addr", r.RemoteAddr))
+				http.Error(w, "invalid token or API key", http.StatusUnauthorized)
 				return
 			}
 
-			keyID, ok := keyring.Lookup(token)
-			if !ok {
-				log.Warn("authentication failed", slog.String("remote_addr", r.RemoteAddr))
-				http.Error(w, "invalid API key", http.StatusUnauthorized)
+			// 3. HTTP Basic → oauthRegistry (direct client credentials).
+			if clientID, clientSecret := extractOAuthClientCredentials(r); clientID != "" {
+				if oauthRegistry == nil {
+					http.Error(w, "authentication is not configured", http.StatusUnauthorized)
+					return
+				}
+				keyID, ok := oauthRegistry.Lookup(clientID, clientSecret)
+				if !ok {
+					log.Warn("oauth client authentication failed", slog.String("remote_addr", r.RemoteAddr))
+					http.Error(w, "invalid OAuth client credentials", http.StatusUnauthorized)
+					return
+				}
+				next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID)))
 				return
 			}
 
-			next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID)))
+			// 4. Query param → keyring.
+			if keyring != nil && cfg.AllowQueryParam {
+				if token := strings.TrimSpace(r.URL.Query().Get(cfg.QueryParam)); token != "" {
+					keyID, ok := keyring.Lookup(token)
+					if !ok {
+						log.Warn("authentication failed", slog.String("remote_addr", r.RemoteAddr))
+						http.Error(w, "invalid API key", http.StatusUnauthorized)
+						return
+					}
+					next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), keyIDContextKey, keyID)))
+					return
+				}
+			}
+
+			http.Error(w, "authentication required", http.StatusUnauthorized)
 		})
 	}
 }
@@ -58,6 +104,30 @@ func extractToken(r *http.Request, headerName string) string {
 	return strings.TrimSpace(credentials)
 }
 
+func extractBearer(r *http.Request) string {
+	authHeader := strings.TrimSpace(r.Header.Get("Authorization"))
+	scheme, credentials, ok := strings.Cut(authHeader, " ")
+	if !ok || !strings.EqualFold(scheme, "Bearer") {
+		return ""
+	}
+	return strings.TrimSpace(credentials)
+}
+
+func extractOAuthClientCredentials(r *http.Request) (string, string) {
+	authHeader := strings.TrimSpace(r.Header.Get("Authorization"))
+	scheme, credentials, ok := strings.Cut(authHeader, " ")
+	if ok && strings.EqualFold(scheme, "Basic") {
+		decoded, err := base64.StdEncoding.DecodeString(strings.TrimSpace(credentials))
+		if err == nil {
+			clientID, clientSecret, found := strings.Cut(string(decoded), ":")
+			if found {
+				return strings.TrimSpace(clientID), strings.TrimSpace(clientSecret)
+			}
+		}
+	}
+	return "", ""
+}
+
 func KeyIDFromContext(ctx context.Context) (string, bool) {
 	value, ok := ctx.Value(keyIDContextKey).(string)
 	return value, ok

internal/auth/oauth_registry.go

@@ -0,0 +1,33 @@
+package auth
+
+import (
+	"crypto/subtle"
+	"fmt"
+
+	"git.warky.dev/wdevs/amcs/internal/config"
+)
+
+type OAuthRegistry struct {
+	clients []config.OAuthClient
+}
+
+func NewOAuthRegistry(clients []config.OAuthClient) (*OAuthRegistry, error) {
+	if len(clients) == 0 {
+		return nil, fmt.Errorf("oauth registry requires at least one client")
+	}
+
+	return &OAuthRegistry{clients: append([]config.OAuthClient(nil), clients...)}, nil
+}
+
+func (o *OAuthRegistry) Lookup(clientID string, clientSecret string) (string, bool) {
+	for _, client := range o.clients {
+		if subtle.ConstantTimeCompare([]byte(client.ClientID), []byte(clientID)) == 1 &&
+			subtle.ConstantTimeCompare([]byte(client.ClientSecret), []byte(clientSecret)) == 1 {
+			if client.ID != "" {
+				return client.ID, true
+			}
+			return client.ClientID, true
+		}
+	}
+	return "", false
+}

internal/auth/oauth_registry_test.go

@@ -0,0 +1,92 @@
+package auth
+
+import (
+	"encoding/base64"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"git.warky.dev/wdevs/amcs/internal/config"
+)
+
+func TestNewOAuthRegistryAndLookup(t *testing.T) {
+	_, err := NewOAuthRegistry(nil)
+	if err == nil {
+		t.Fatal("NewOAuthRegistry(nil) error = nil, want error")
+	}
+
+	registry, err := NewOAuthRegistry([]config.OAuthClient{{
+		ID:           "oauth-client",
+		ClientID:     "client-id",
+		ClientSecret: "client-secret",
+	}})
+	if err != nil {
+		t.Fatalf("NewOAuthRegistry() error = %v", err)
+	}
+
+	if got, ok := registry.Lookup("client-id", "client-secret"); !ok || got != "oauth-client" {
+		t.Fatalf("Lookup(client-id, client-secret) = (%q, %v), want (oauth-client, true)", got, ok)
+	}
+	if _, ok := registry.Lookup("client-id", "wrong"); ok {
+		t.Fatal("Lookup(client-id, wrong) = true, want false")
+	}
+}
+
+func TestMiddlewareAllowsOAuthBasicAuthAndSetsContext(t *testing.T) {
+	oauthRegistry, err := NewOAuthRegistry([]config.OAuthClient{{
+		ID:           "oauth-client",
+		ClientID:     "client-id",
+		ClientSecret: "client-secret",
+	}})
+	if err != nil {
+		t.Fatalf("NewOAuthRegistry() error = %v", err)
+	}
+
+	handler := Middleware(config.AuthConfig{}, nil, oauthRegistry, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		keyID, ok := KeyIDFromContext(r.Context())
+		if !ok || keyID != "oauth-client" {
+			t.Fatalf("KeyIDFromContext() = (%q, %v), want (oauth-client, true)", keyID, ok)
+		}
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	req := httptest.NewRequest(http.MethodGet, "/mcp", nil)
+	req.Header.Set("Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte("client-id:client-secret")))
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	if rec.Code != http.StatusOK {
+		t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK)
+	}
+}
+
+
+func TestMiddlewareRejectsOAuthMissingOrInvalidCredentials(t *testing.T) {
+	oauthRegistry, err := NewOAuthRegistry([]config.OAuthClient{{
+		ID:           "oauth-client",
+		ClientID:     "client-id",
+		ClientSecret: "client-secret",
+	}})
+	if err != nil {
+		t.Fatalf("NewOAuthRegistry() error = %v", err)
+	}
+
+	handler := Middleware(config.AuthConfig{}, nil, oauthRegistry, nil, testLogger())(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		t.Fatal("next handler should not be called")
+	}))
+
+	req := httptest.NewRequest(http.MethodGet, "/mcp", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+	if rec.Code != http.StatusUnauthorized {
+		t.Fatalf("missing credentials status = %d, want %d", rec.Code, http.StatusUnauthorized)
+	}
+
+	req = httptest.NewRequest(http.MethodGet, "/mcp", nil)
+	req.Header.Set("Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte("client-id:wrong")))
+	rec = httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+	if rec.Code != http.StatusUnauthorized {
+		t.Fatalf("invalid credentials status = %d, want %d", rec.Code, http.StatusUnauthorized)
+	}
+}

internal/auth/token_store.go

@@ -0,0 +1,74 @@
+package auth
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"sync"
+	"time"
+)
+
+const defaultTokenTTL = time.Hour
+
+type tokenEntry struct {
+	keyID     string
+	expiresAt time.Time
+}
+
+// TokenStore issues and validates short-lived opaque access tokens for OAuth
+// client credentials flow.
+type TokenStore struct {
+	mu     sync.RWMutex
+	tokens map[string]tokenEntry
+	ttl    time.Duration
+}
+
+func NewTokenStore(ttl time.Duration) *TokenStore {
+	if ttl <= 0 {
+		ttl = defaultTokenTTL
+	}
+	s := &TokenStore{
+		tokens: make(map[string]tokenEntry),
+		ttl:    ttl,
+	}
+	go s.sweepLoop()
+	return s
+}
+
+// Issue generates a new token for the given keyID and returns the token and its TTL.
+func (s *TokenStore) Issue(keyID string) (string, time.Duration, error) {
+	b := make([]byte, 32)
+	if _, err := rand.Read(b); err != nil {
+		return "", 0, err
+	}
+	token := hex.EncodeToString(b)
+	s.mu.Lock()
+	s.tokens[token] = tokenEntry{keyID: keyID, expiresAt: time.Now().Add(s.ttl)}
+	s.mu.Unlock()
+	return token, s.ttl, nil
+}
+
+// Lookup validates a token and returns the associated keyID.
+func (s *TokenStore) Lookup(token string) (string, bool) {
+	s.mu.RLock()
+	entry, ok := s.tokens[token]
+	s.mu.RUnlock()
+	if !ok || time.Now().After(entry.expiresAt) {
+		return "", false
+	}
+	return entry.keyID, true
+}
+
+func (s *TokenStore) sweepLoop() {
+	ticker := time.NewTicker(5 * time.Minute)
+	defer ticker.Stop()
+	for range ticker.C {
+		now := time.Now()
+		s.mu.Lock()
+		for token, entry := range s.tokens {
+			if now.After(entry.expiresAt) {
+				delete(s.tokens, token)
+			}
+		}
+		s.mu.Unlock()
+	}
+}