Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
275 changes: 211 additions & 64 deletions nitrostack/auth/oauth.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,20 @@
import os
import sys
import time
import json
import urllib.request
import urllib.parse
from typing import List, Optional, Dict, Any
from typing import List, Optional, Dict, Any, Tuple
from http.server import HTTPServer, BaseHTTPRequestHandler
import threading
from nitrostack.core.module import module
from nitrostack.core.di import DIContainer
from nitrostack.auth.oauth_module import (
build_authorization_server_metadata,
build_protected_resource_metadata,
build_registration_response,
is_client_registration_enabled,
)

class OAuthService:
def __init__(
Expand All @@ -22,6 +29,10 @@ def __init__(
jwks_uri: Optional[str] = None,
audience: Optional[str] = None,
issuer: Optional[str] = None,
token_cache_seconds: Optional[int] = None,
enable_client_registration: Optional[bool] = None,
static_client_id: Optional[str] = None,
static_client_secret: Optional[str] = None,
):
self.resource_uri = resource_uri
self.authorization_servers = authorization_servers
Expand All @@ -30,12 +41,34 @@ def __init__(
self.token_introspection_client_id = token_introspection_client_id
self.token_introspection_client_secret = token_introspection_client_secret
self.discovery_port = discovery_port

# Environmental fallbacks
self.jwks_uri = jwks_uri or os.environ.get("JWKS_URI")
self.audience = audience or os.environ.get("TOKEN_AUDIENCE") or resource_uri
self.issuer = issuer or os.environ.get("TOKEN_ISSUER")

self.token_cache_seconds = (
token_cache_seconds
if token_cache_seconds is not None
else int(os.environ.get("OAUTH_TOKEN_CACHE_SECONDS", "300"))
)

# Dynamic Client Registration (RFC 7591) — see oauth_module.py for why this is
# a simplified, static-credential variant rather than full per-client storage.
self.enable_client_registration = (
enable_client_registration
if enable_client_registration is not None
else os.environ.get("OAUTH_ENABLE_CLIENT_REGISTRATION", "").lower() == "true"
)
self.static_client_id = static_client_id or os.environ.get("OAUTH_CLIENT_ID")
self.static_client_secret = static_client_secret or os.environ.get("OAUTH_CLIENT_SECRET")

# Caches: JWKS client objects (keyed by jwks_uri — PyJWKClient already does its
# own internal signing-key caching, this just avoids reconstructing the client
# itself on every call) and introspection *results* (keyed by token, so repeated
# calls with the same token skip both the HTTP round-trip and JWT verification).
self._jwks_clients: Dict[str, Any] = {}
self._token_cache: Dict[str, Tuple[Dict[str, Any], float]] = {}

self._server: Optional[HTTPServer] = None
self._thread: Optional[threading.Thread] = None

Expand All @@ -45,6 +78,13 @@ def start_discovery_server(self) -> None:
return

service_instance = self
registration_path = "/oauth/v2/register"

def _write_json(handler: BaseHTTPRequestHandler, status: int, payload: Dict[str, Any]) -> None:
handler.send_response(status)
handler.send_header("Content-Type", "application/json")
handler.end_headers()
handler.wfile.write(json.dumps(payload).encode("utf-8"))

class DiscoveryHandler(BaseHTTPRequestHandler):
def log_message(self, format, *args):
Expand All @@ -53,30 +93,43 @@ def log_message(self, format, *args):

def do_GET(self):
if self.path == "/.well-known/oauth-protected-resource":
response_data = {
"resource": service_instance.resource_uri,
"authorization_servers": service_instance.authorization_servers,
"scopes_supported": service_instance.scopes_supported
}
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.end_headers()
self.wfile.write(json.dumps(response_data).encode("utf-8"))
_write_json(self, 200, build_protected_resource_metadata(service_instance))
elif self.path == "/.well-known/oauth-authorization-server":
# Mock/basic authorization server metadata if query hits this resource
response_data = {
"issuer": service_instance.authorization_servers[0] if service_instance.authorization_servers else "http://localhost",
"token_endpoint": service_instance.token_introspection_endpoint or "",
"introspection_endpoint": service_instance.token_introspection_endpoint or ""
}
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.end_headers()
self.wfile.write(json.dumps(response_data).encode("utf-8"))
registration_endpoint = (
registration_path if is_client_registration_enabled(service_instance) else None
)
_write_json(
self,
200,
build_authorization_server_metadata(service_instance, registration_endpoint),
)
else:
self.send_response(404)
self.end_headers()

def do_POST(self):
if self.path != registration_path:
self.send_response(404)
self.end_headers()
return

if not is_client_registration_enabled(service_instance):
_write_json(
self,
404,
{"error": "not_found", "error_description": "Client registration is not enabled"},
)
return

length = int(self.headers.get("Content-Length", 0))
raw_body = self.rfile.read(length) if length else b""
try:
body = json.loads(raw_body.decode("utf-8")) if raw_body else {}
except Exception:
body = {}

_write_json(self, 200, build_registration_response(service_instance, body))

def run_server():
# Try binding to OAUTH_DISCOVERY_PORT
port = int(os.environ.get("OAUTH_DISCOVERY_PORT", self.discovery_port))
Expand Down Expand Up @@ -110,48 +163,80 @@ def stop_discovery_server(self) -> None:

async def introspect_token(self, token: str) -> Dict[str, Any]:
"""
Validates token using JWKS verification or RFC 7662 token introspection.
Validate a Bearer token and return its introspection result (RFC 7662 shape:
{"active": bool, ...claims}).

Checks, in order:
1. Cache (a prior successful result for this exact token, still within TTL).
2. Token introspection endpoint (RFC 7662), if configured.
3. JWKS/JWT signature verification, if configured.
4. Neither configured -> {"active": False}. There is deliberately no "assume
valid" fallback: a server that hasn't been told how to validate tokens
must reject them, not accept everything. (This mirrors the TypeScript SDK,
which has no such fallback either.)

Every successful result is passed through `_validate_audience` (RFC 8707)
before being trusted or cached — a token that's valid but wasn't issued for
this resource is still rejected.
"""
# 1. JWKS Verification if configured
if self.jwks_uri:
try:
import jwt
# Parse JWT headers to get kid
unverified_headers = jwt.get_unverified_header(token)
jwks_client = jwt.PyJWKClient(self.jwks_uri)
signing_key = jwks_client.get_signing_key_from_jwt(token)

# Verify token signature
data = jwt.decode(
token,
signing_key.key,
algorithms=["RS256"],
audience=self.audience,
issuer=self.issuer
)
return {
"active": True,
"scope": data.get("scope", ""),
"sub": data.get("sub"),
"client_id": data.get("client_id")
}
except Exception as e:
# Log signature failure to stderr
sys.stderr.write(f"OAuth JWKS verification failed: {e}\n")
sys.stderr.flush()
return {"active": False}
cached = self._cache_get(token)
if cached is not None:
return cached

if not self.token_introspection_endpoint:
# If no introspection endpoint is configured, mock active check for local debugging
# A real deployment must provide an introspection endpoint.
sys.stderr.write("OAuth Warning: No token_introspection_endpoint configured. Assuming mock active.\n")
return {"active": True, "scope": " ".join(self.scopes_supported), "sub": "mock-user"}
if self.token_introspection_endpoint:
result = await self._introspect_via_endpoint(token)
elif self.jwks_uri:
result = self._introspect_via_jwks(token)
else:
return {"active": False}

# Perform HTTP POST request
if result.get("active") and not self._validate_audience(result):
sys.stderr.write(
f"OAuth: token rejected, audience mismatch (expected {self.audience!r})\n"
)
sys.stderr.flush()
result = {"active": False}

if result.get("active"):
self._cache_set(token, result)
return result

def _introspect_via_jwks(self, token: str) -> Dict[str, Any]:
"""JWT signature verification using a cached JWKS client (RFC 7517/7519)."""
try:
import jwt
jwks_client = self._get_jwks_client(self.jwks_uri)
signing_key = jwks_client.get_signing_key_from_jwt(token)

data = jwt.decode(
token,
signing_key.key,
algorithms=["RS256"],
audience=self.audience,
issuer=self.issuer,
)
return {
"active": True,
"scope": data.get("scope", ""),
"sub": data.get("sub"),
"client_id": data.get("client_id"),
"aud": data.get("aud"),
"exp": data.get("exp"),
"iat": data.get("iat"),
"iss": data.get("iss"),
}
except Exception as e:
sys.stderr.write(f"OAuth JWKS verification failed: {e}\n")
sys.stderr.flush()
return {"active": False}

async def _introspect_via_endpoint(self, token: str) -> Dict[str, Any]:
"""RFC 7662 token introspection: POST the token to the authorization server
and ask whether it's active."""
data = urllib.parse.urlencode({"token": token}).encode("utf-8")
req = urllib.request.Request(self.token_introspection_endpoint, data=data, method="POST")
req.add_header("Content-Type", "application/x-www-form-urlencoded")

# Add basic auth if client credentials provided
if self.token_introspection_client_id and self.token_introspection_client_secret:
import base64
Expand All @@ -160,22 +245,76 @@ async def introspect_token(self, token: str) -> Dict[str, Any]:
req.add_header("Authorization", f"Basic {encoded_auth}")

try:
# We run in a threadpool or run_in_executor to avoid blocking async loop
# But standard library urllib.request is synchronous, so let's run it synchronously in context
# (or use asyncio loop.run_in_executor if we are in async method).
# urllib.request is synchronous; run it on a thread so it doesn't block
# the event loop this async method is running on.
import asyncio
loop = asyncio.get_event_loop()

def do_request():
with urllib.request.urlopen(req, timeout=5) as response:
return json.loads(response.read().decode("utf-8"))

return await loop.run_in_executor(None, do_request)
except Exception as e:
sys.stderr.write(f"OAuth Introspection Request Failed: {e}\n")
sys.stderr.flush()
return {"active": False}

def _get_jwks_client(self, jwks_uri: str):
"""Return a cached PyJWKClient for this URI, creating it on first use.

PyJWKClient already caches individual signing keys internally; this cache
is one level up — it avoids reconstructing the client object itself (and
re-fetching the whole key set) on every single token check.
"""
import jwt
client = self._jwks_clients.get(jwks_uri)
if client is None:
client = jwt.PyJWKClient(jwks_uri)
self._jwks_clients[jwks_uri] = client
return client

def _validate_audience(self, introspection: Dict[str, Any]) -> bool:
"""
RFC 8707 resource-indicator check: does this token's `aud` claim include
the resource it's being presented to? Without this, a token minted for a
*different* service could be replayed here and accepted — the token is
legitimately signed/active, just not meant for this resource.

`aud` is legal as either a single string or a list of strings per JWT
conventions, so it's normalized to a list before comparing.
"""
expected = self.audience or self.resource_uri
if not expected:
# Nothing configured to check against — permissive, matching the
# TypeScript SDK's default when no audience is configured.
return True

raw_aud = introspection.get("aud")
if raw_aud is None:
# No audience claim on the token at all — nothing to validate against.
# Treat as permissive rather than rejecting tokens from authorization
# servers that don't emit `aud`.
return True

token_audiences = raw_aud if isinstance(raw_aud, list) else [raw_aud]
return expected in token_audiences

def _cache_get(self, token: str) -> Optional[Dict[str, Any]]:
cached = self._token_cache.get(token)
if cached is None:
return None
result, expires_at = cached
if time.monotonic() >= expires_at:
del self._token_cache[token]
return None
return result

def _cache_set(self, token: str, result: Dict[str, Any]) -> None:
if self.token_cache_seconds <= 0:
return
self._token_cache[token] = (result, time.monotonic() + self.token_cache_seconds)

@module(name="OAuthModule")
class OAuthModule:
@classmethod
Expand All @@ -191,6 +330,10 @@ def for_root(
jwks_uri: Optional[str] = None,
audience: Optional[str] = None,
issuer: Optional[str] = None,
token_cache_seconds: Optional[int] = None,
enable_client_registration: Optional[bool] = None,
static_client_id: Optional[str] = None,
static_client_secret: Optional[str] = None,
):
service = OAuthService(
resource_uri=resource_uri,
Expand All @@ -202,7 +345,11 @@ def for_root(
discovery_port=discovery_port,
jwks_uri=jwks_uri,
audience=audience,
issuer=issuer
issuer=issuer,
token_cache_seconds=token_cache_seconds,
enable_client_registration=enable_client_registration,
static_client_id=static_client_id,
static_client_secret=static_client_secret,
)
DIContainer.get_instance().register_value(OAuthService, service)
return cls
Loading