Add OAuth 2.1 resource server for external provider auth

CHANGELOG.md

@@ -8,6 +8,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
+- **OAuth 2.1 resource server**: provider-agnostic JWKS-based token validation for external OAuth providers (WorkOS, Auth0, Cloudflare Access, Keycloak, etc.)
+- **Dual auth**: self-signed JWTs (via CLI) and OAuth provider tokens both accepted — OAuth for interactive clients, self-signed for edge providers/scripts
+- **User auto-provisioning**: auto-create user record on first valid OAuth login (`AWARENESS_OAUTH_AUTO_PROVISION`, default: false for tighter control during early access)
+- **Well-known metadata**: `/.well-known/oauth-protected-resource` (RFC 9728) for OAuth discovery by MCP clients
+- **OAuth env vars**: `AWARENESS_OAUTH_ISSUER`, `AWARENESS_OAUTH_AUDIENCE`, `AWARENESS_OAUTH_JWKS_URI`, `AWARENESS_OAUTH_USER_CLAIM`, `AWARENESS_OAUTH_AUTO_PROVISION`
 - **JWT auth middleware**: opt-in via `AWARENESS_AUTH_REQUIRED=true`, validates Bearer tokens, extracts owner_id from `sub` claim
 - **Row-level security**: Postgres RLS policies on all data tables as defense-in-depth alongside application-level owner_id filtering
 - **CLI: `mcp-awareness-user`**: add/list/set-password/export/delete users with email normalization, E.164 phone validation, argon2id password hashing

alembic/versions/i4d5e6f7g8h9_add_oauth_columns.py

@@ -0,0 +1,58 @@
+# mcp-awareness — ambient system awareness for AI agents
+# Copyright (C) 2026 Chris Means
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""add OAuth identity columns to users table
+
+Revision ID: i4d5e6f7g8h9
+Revises: h3c4d5e6f7g8
+Create Date: 2026-03-29 20:00:00.000000
+
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "i4d5e6f7g8h9"
+down_revision: str | Sequence[str] | None = "h3c4d5e6f7g8"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    # OAuth identity: sub claim + issuer for provider-agnostic lookup
+    op.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS oauth_subject TEXT")
+    op.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS oauth_issuer TEXT")
+    # Unique constraint: one OAuth identity per provider per user
+    op.execute(
+        "CREATE UNIQUE INDEX IF NOT EXISTS ix_users_oauth_identity "
+        "ON users (oauth_issuer, oauth_subject) WHERE oauth_issuer IS NOT NULL"
+    )
+    # Fast lookup index for every authenticated request
+    op.execute(
+        "CREATE INDEX IF NOT EXISTS ix_users_oauth_subject "
+        "ON users (oauth_subject) WHERE oauth_subject IS NOT NULL"
+    )
+
+
+def downgrade() -> None:
+    op.execute("DROP INDEX IF EXISTS ix_users_oauth_subject")
+    op.execute("DROP INDEX IF EXISTS ix_users_oauth_identity")
+    op.execute("ALTER TABLE users DROP COLUMN IF EXISTS oauth_issuer")
+    op.execute("ALTER TABLE users DROP COLUMN IF EXISTS oauth_subject")

pyproject.toml

@@ -11,7 +11,7 @@ dependencies = [
     "psycopg_pool>=3.2",
     "alembic>=1.13",
     "sqlalchemy>=2.0",
-    "PyJWT>=2.8",
+    "PyJWT[crypto]>=2.8",
     "argon2-cffi>=23.1",
     "phonenumbers>=8.13",
     "zxcvbn>=4.4",

src/mcp_awareness/middleware.py

@@ -93,55 +93,94 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         await self.app(scope, receive, send)
 
 
+class WellKnownMiddleware:
+    """Serve /.well-known/oauth-protected-resource (RFC 9728)."""
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        oauth_issuer: str,
+        host: str,
+        port: int,
+        mount_path: str = "",
+    ) -> None:
+        self.app = app
+        self.oauth_issuer = oauth_issuer.rstrip("/")
+        self.host = host
+        self.port = port
+        self.mount_path = mount_path
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] == "http":
+            path = scope.get("path", "")
+            if path == "/.well-known/oauth-protected-resource":
+                metadata = {
+                    "resource": f"https://{self.host}:{self.port}{self.mount_path}/mcp",
+                    "authorization_servers": [self.oauth_issuer],
+                    "token_methods": ["Bearer"],
+                }
+                resp = JSONResponse(metadata)
+                await resp(scope, receive, send)
+                return
+        await self.app(scope, receive, send)
+
+
 class AuthMiddleware:
-    """Validate JWT Bearer token and set owner context."""
+    """Validate JWT Bearer token and set owner context.
 
-    def __init__(self, app: ASGIApp, jwt_secret: str, algorithm: str = "HS256") -> None:
+    Supports dual auth: self-signed JWTs (via shared secret) and OAuth provider
+    tokens (via JWKS). Self-signed is tried first; if it fails and an OAuth
+    validator is configured, the token is validated against the provider's keys.
+    """
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        jwt_secret: str,
+        algorithm: str = "HS256",
+        oauth_validator: object | None = None,
+        auto_provision: bool = True,
+        resource_metadata_url: str = "",
+    ) -> None:
         self.app = app
         self.jwt_secret = jwt_secret
         self.algorithm = algorithm
+        self.oauth_validator = oauth_validator
+        self.auto_provision = auto_provision
+        self.resource_metadata_url = resource_metadata_url
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if scope["type"] != "http":
             await self.app(scope, receive, send)
             return
 
         path = scope.get("path", "")
-        # Skip auth for health, favicon, and non-MCP paths
-        if path in ("/health", "/favicon.ico"):
+        # Skip auth for health, favicon, well-known, and non-MCP paths
+        if path in ("/health", "/favicon.ico") or path.startswith("/.well-known/"):
             await self.app(scope, receive, send)
             return
 
         # Extract Bearer token
         headers = dict(scope.get("headers", []))
         auth_header = headers.get(b"authorization", b"").decode()
         if not auth_header.startswith("Bearer "):
-            resp = JSONResponse(
-                {"error": "Missing or invalid Authorization header"}, status_code=401
-            )
+            resp = self._unauthorized("Missing or invalid Authorization header")
             await resp(scope, receive, send)
             return
 
         token = auth_header[7:]  # Strip "Bearer "
-        try:
-            import jwt
 
-            payload = jwt.decode(token, self.jwt_secret, algorithms=[self.algorithm])
-            owner_id: str | None = payload.get("sub")
-            if not owner_id:
-                resp = JSONResponse({"error": "JWT missing 'sub' claim"}, status_code=401)
-                await resp(scope, receive, send)
-                return
-        except Exception as exc:
-            # Handle both ExpiredSignatureError and InvalidTokenError
-            import jwt as jwt_mod
+        # Try self-signed JWT first
+        owner_id, error = self._try_self_signed(token)
 
-            if isinstance(exc, jwt_mod.ExpiredSignatureError):
-                resp = JSONResponse({"error": "Token expired"}, status_code=401)
-            elif isinstance(exc, jwt_mod.InvalidTokenError):
-                resp = JSONResponse({"error": "Invalid token"}, status_code=401)
-            else:
-                raise
+        # Fall back to OAuth provider validation
+        if owner_id is None and self.oauth_validator is not None:
+            owner_id = await self._try_oauth(token)
+            if owner_id is not None:
+                error = None
+
+        if owner_id is None:
+            resp = self._unauthorized(error or "Invalid token")
             await resp(scope, receive, send)
             return
 
@@ -153,3 +192,103 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
             await self.app(scope, receive, send)
         finally:
             _owner_ctx.reset(token_reset)
+
+    def _try_self_signed(self, token: str) -> tuple[str | None, str | None]:
+        """Validate a self-signed JWT (from mcp-awareness-token CLI).
+
+        Returns (owner_id, error_message). On success error is None.
+        On failure owner_id is None and error carries the reason.
+        """
+        if not self.jwt_secret:
+            return None, None
+        try:
+            import jwt
+
+            payload = jwt.decode(token, self.jwt_secret, algorithms=[self.algorithm])
+            owner_id: str | None = payload.get("sub")
+            if owner_id:
+                return owner_id, None
+            return None, "JWT missing 'sub' claim"
+        except Exception as exc:
+            import jwt as jwt_mod
+
+            if isinstance(exc, jwt_mod.ExpiredSignatureError):
+                return None, "Token expired"
+            if isinstance(exc, jwt_mod.InvalidTokenError):
+                return None, "Invalid token"
+            return None, None
+
+    async def _try_oauth(self, token: str) -> str | None:
+        """Validate an OAuth token against the external provider's JWKS."""
+        from .oauth import OAuthTokenValidator
+
+        validator: OAuthTokenValidator = self.oauth_validator  # type: ignore[assignment]
+        try:
+            claims = validator.validate(token)
+        except Exception:
+            return None
+
+        owner_id = claims["owner_id"]
+        oauth_subject = claims.get("oauth_subject")
+        oauth_issuer = claims.get("oauth_issuer")
+        email = claims.get("email")
+
+        # Resolve user identity: OAuth lookup → email link → auto-provision
+        resolved_id = self._resolve_user(
+            owner_id, email, claims.get("name"), oauth_subject, oauth_issuer
+        )
+
+        return resolved_id or owner_id
+
+    def _resolve_user(
+        self,
+        owner_id: str,
+        email: str | None,
+        display_name: str | None,
+        oauth_subject: str | None,
+        oauth_issuer: str | None,
+    ) -> str | None:
+        """Resolve OAuth token to a local user, linking or creating as needed.
+
+        Resolution order:
+        1. Look up by OAuth identity (issuer + subject) — already linked user
+        2. If email present, try to link to a pre-provisioned user by email
+        3. If auto_provision enabled, create a new user
+        4. Otherwise return None (use owner_id from token as-is)
+        """
+        try:
+            from .server import store
+
+            # 1. Already linked?
+            if oauth_issuer and oauth_subject:
+                existing = store.get_user_by_oauth(oauth_issuer, oauth_subject)
+                if existing:
+                    return str(existing["id"])
+
+            # 2. Pre-provisioned user with matching email? Link on first login.
+            if email and oauth_subject and oauth_issuer:
+                linked_id = store.link_oauth_identity(oauth_subject, oauth_issuer, email)
+                if linked_id:
+                    return str(linked_id)
+
+            # 3. Auto-provision new user
+            if self.auto_provision:
+                store.create_user_if_not_exists(
+                    owner_id, email, display_name, oauth_subject, oauth_issuer
+                )
+                return owner_id
+
+        except Exception:
+            # Don't fail the request if user resolution fails
+            pass
+
+        return None
+
+    def _unauthorized(self, message: str) -> JSONResponse:
+        """Build a 401 response with proper WWW-Authenticate header."""
+        headers: dict[str, str] = {}
+        if self.resource_metadata_url:
+            headers["WWW-Authenticate"] = f'Bearer resource_metadata="{self.resource_metadata_url}"'
+        else:
+            headers["WWW-Authenticate"] = "Bearer"
+        return JSONResponse({"error": message}, status_code=401, headers=headers)