diff --git a/backend/druks/mcp/constants.py b/backend/druks/mcp/constants.py index dbfe7294..974ca98f 100644 --- a/backend/druks/mcp/constants.py +++ b/backend/druks/mcp/constants.py @@ -17,26 +17,16 @@ REGISTRY_SEARCH_CACHE_PREFIX = "mcp:registry:search:" REGISTRY_CACHE_TTL_SECONDS = 300 -# OAuth connect + mint plumbing. The callback path is public API surface — the +# OAuth connect + mint plumbing rides the shared engine (druks.services' +# OauthClient) under this provider namespace. The namespace keys the engine's +# connect-state, token-cache, and refresh-lock Redis entries, so it is pinned: +# a rolling deploy's old and new processes must elect one refresher per grant. +# The prefixes spell out the derived keys; the token and lock keys append +# {name}:{account_id}. The callback path is public API surface — the # authorization server redirects the operator's browser to -# {urls.endpoint}{OAUTH_CALLBACK_PATH} after consent. Access tokens cache in -# Redis under the token key prefix for their lifetime minus the skew (so a -# token injected into a run never expires moments after delivery); pending -# connect state (PKCE verifier + endpoints) lives under the connect prefix for -# its short TTL, single-use. +# {urls.endpoint}{OAUTH_CALLBACK_PATH} after consent. +OAUTH_PROVIDER = "mcp:oauth" OAUTH_CALLBACK_PATH = "/api/mcp-servers/oauth/callback" -OAUTH_CONNECT_STATE_PREFIX = "mcp:oauth:connect:" -OAUTH_ACCESS_TOKEN_PREFIX = "mcp:oauth:access_token:" -OAUTH_CONNECT_STATE_TTL_SECONDS = 600 -OAUTH_TOKEN_TTL_SKEW_SECONDS = 60 - -# Mint's mutual exclusion, in the Redis that fronts the token cache (the run -# lock's SET NX idiom): a rotating grant tolerates exactly one refresher per -# grant. The server name and the grant's account are part of both keys. The lock TTL -# is a crash backstop at three times the HTTP client's timeout — a live refresh -# cannot outlive it. Losers poll the cache on the interval for about one -# token-endpoint round trip, then fail loudly. -OAUTH_REFRESH_LOCK_PREFIX = "mcp:oauth:refresh_lock:" -OAUTH_REFRESH_LOCK_TTL_SECONDS = 90 -OAUTH_MINT_WAIT_INTERVAL_SECONDS = 0.2 -OAUTH_MINT_WAIT_ATTEMPTS = 150 +OAUTH_CONNECT_STATE_PREFIX = f"{OAUTH_PROVIDER}:connect:" +OAUTH_ACCESS_TOKEN_PREFIX = f"{OAUTH_PROVIDER}:access_token:" +OAUTH_REFRESH_LOCK_PREFIX = f"{OAUTH_PROVIDER}:refresh_lock:" diff --git a/backend/druks/mcp/oauth.py b/backend/druks/mcp/oauth.py index 360fc471..2b6a5468 100644 --- a/backend/druks/mcp/oauth.py +++ b/backend/druks/mcp/oauth.py @@ -1,31 +1,16 @@ -import asyncio -import base64 -import hashlib -import json -import secrets -from typing import cast -from urllib.parse import urlencode, urlparse +from urllib.parse import urlparse import httpx from sqlalchemy import select, update from sqlalchemy.dialects.postgresql import insert as pg_insert -from druks.database import db_session -from druks.mcp.constants import ( - OAUTH_ACCESS_TOKEN_PREFIX, - OAUTH_CALLBACK_PATH, - OAUTH_CONNECT_STATE_PREFIX, - OAUTH_CONNECT_STATE_TTL_SECONDS, - OAUTH_MINT_WAIT_ATTEMPTS, - OAUTH_MINT_WAIT_INTERVAL_SECONDS, - OAUTH_REFRESH_LOCK_PREFIX, - OAUTH_REFRESH_LOCK_TTL_SECONDS, - OAUTH_TOKEN_TTL_SKEW_SECONDS, -) +from druks.database import db_session, get_session +from druks.mcp.constants import OAUTH_CALLBACK_PATH, OAUTH_PROVIDER from druks.mcp.enums import IdentityMode from druks.mcp.exceptions import GrantRefreshError, MissingGrantError, OauthConnectError from druks.mcp.models import McpOauthGrant, McpServer -from druks.redis import get_client +from druks.services import OauthClient, OauthExchangeError, OauthRefreshError +from druks.services.constants import OAUTH_MINT_WAIT_ATTEMPTS, OAUTH_MINT_WAIT_INTERVAL_SECONDS def _http() -> httpx.AsyncClient: @@ -167,10 +152,10 @@ async def begin_connect( identity_mode: IdentityMode, ) -> str: """Start the operator's authorization-code + PKCE flow for one server: - discover the authorization server, register druks as a public client, stash - the pending exchange (verifier + endpoints) in Redis under the state, and - return the consent URL to open. Nothing durable is written here — an - abandoned consent simply expires.""" + discover the authorization server, register druks as a public client, and + hand the engine the resulting client to stash the pending exchange and + render the consent URL. Nothing durable is written here — an abandoned + consent simply expires.""" redirect_uri = f"{endpoint.rstrip('/')}{OAUTH_CALLBACK_PATH}" async with _http() as client: metadata = await _discover(client, name, server_url) @@ -180,81 +165,40 @@ async def begin_connect( if methods is not None and "S256" not in methods: raise OauthConnectError(name, "the authorization server does not support PKCE S256") registration = await _register_client(client, name, metadata, redirect_uri) - state = secrets.token_urlsafe(32) - code_verifier = secrets.token_urlsafe(64) - code_challenge = ( - base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest()) - .rstrip(b"=") - .decode() + return await OauthClient( + provider=OAUTH_PROVIDER, + authorization_endpoint=metadata["authorization_endpoint"], + token_endpoint=metadata["token_endpoint"], + client_id=registration["client_id"], + client_secret=registration.get("client_secret", ""), + # RFC 8707: bind the tokens to the MCP server they are for. + extra_token_params={"resource": server_url}, + ).begin_connect( + redirect_uri=redirect_uri, + context={ + "name": name, + "server_url": server_url, + "account_id": account_id, + "identity_mode": identity_mode, + }, + extra_authorize_params={"resource": server_url}, ) - pending = { - "name": name, - "server_url": server_url, - "account_id": account_id, - "identity_mode": identity_mode, - "code_verifier": code_verifier, - "token_endpoint": metadata["token_endpoint"], - "client_id": registration["client_id"], - "client_secret": registration.get("client_secret", ""), - "redirect_uri": redirect_uri, - } - await get_client().set( - f"{OAUTH_CONNECT_STATE_PREFIX}{state}", - json.dumps(pending), - ex=OAUTH_CONNECT_STATE_TTL_SECONDS, - ) - query = urlencode( - { - "response_type": "code", - "client_id": registration["client_id"], - "redirect_uri": redirect_uri, - "state": state, - "code_challenge": code_challenge, - "code_challenge_method": "S256", - # RFC 8707: bind the token to the MCP server it is for. - "resource": server_url, - } - ) - return f"{metadata['authorization_endpoint']}?{query}" async def complete_connect(*, state: str, code: str) -> str: - """The callback half: consume the pending state (single-use), exchange the - code + verifier for tokens, and store the grant. Returns the server name. - The grant is the only outcome — nothing is cached here, because it becomes - real only when this request's transaction commits, and a cache filled - ahead of that would outlive its failure. The first delivery mints from the - committed grant.""" - raw = await get_client().getdel(f"{OAUTH_CONNECT_STATE_PREFIX}{state}") - if not raw: - raise OauthConnectError("unknown", "unknown or expired state; start the connect flow again") - pending = json.loads(raw) - name = pending["name"] - data = { - "grant_type": "authorization_code", - "code": code, - "redirect_uri": pending["redirect_uri"], - "client_id": pending["client_id"], - "code_verifier": pending["code_verifier"], - "resource": pending["server_url"], - } - if pending["client_secret"]: - data["client_secret"] = pending["client_secret"] - async with _http() as client: - try: - response = await client.post(pending["token_endpoint"], data=data) - except httpx.HTTPError as error: - raise OauthConnectError(name, f"code exchange failed: {error}") from error - if response.status_code != 200: - raise OauthConnectError(name, f"code exchange failed: HTTP {response.status_code}") + """The callback half: the shared exchange (single-use state, code + + verifier), then the durable outcome — claim the server's identity mode and + store the grant. Returns the server name. The grant is the only outcome — + nothing is cached here, because it becomes real only when this request's + transaction commits, and a cache filled ahead of that would outlive its + failure. The first delivery mints from the committed grant.""" try: - tokens = response.json() - except ValueError as error: - raise OauthConnectError(name, "the token endpoint returned malformed JSON") from error - if not isinstance(tokens, dict) or not tokens.get("refresh_token"): - raise OauthConnectError( - name, "the authorization server granted no refresh token; druks needs offline access" - ) + tokens, pending = await OauthClient( + provider=OAUTH_PROVIDER, http_factory=_http + ).complete_connect(state=state, code=code) + except OauthExchangeError as error: + raise OauthConnectError(error.context.get("name", "unknown"), error.reason) from error + name = pending["name"] # The first completed connect claims the mode: insert the row if absent, # fill the mode if unclaimed. A concurrent claim wins the row lock; the # select reads whichever choice landed, and the grant goes under it. @@ -290,80 +234,68 @@ async def complete_connect(*, state: str, code: str) -> str: async def evict_access_token(name: str, account_id: str) -> None: - await get_client().delete(f"{OAUTH_ACCESS_TOKEN_PREFIX}{name}:{account_id}") + await OauthClient(provider=OAUTH_PROVIDER).evict_access_token(f"{name}:{account_id}") async def mint_access_token(name: str, account_id: str) -> str: - """The delivery-side token for a connected server: the cached access token - while it lives, else one refreshed from the grant. The provider may rotate - the refresh token on use — two concurrent refreshes trip its reuse - detection and can revoke the whole grant — so the Redis that fronts the - cache also elects one refresher per grant (the run lock's SET NX idiom; - the TTL is a crash backstop a live refresh cannot outlive). Losers poll - for the winner's cache fill, for about one token-endpoint round trip, then - fail loudly — delivery never ships a server the agent can't authenticate - to.""" - redis = get_client() - token_key = f"{OAUTH_ACCESS_TOKEN_PREFIX}{name}:{account_id}" - lock_key = f"{OAUTH_REFRESH_LOCK_PREFIX}{name}:{account_id}" - for _ in range(OAUTH_MINT_WAIT_ATTEMPTS): - cached = await redis.get(token_key) - if cached: - return cast(bytes, cached).decode() - if await redis.set(lock_key, "1", nx=True, ex=OAUTH_REFRESH_LOCK_TTL_SECONDS): - break - await asyncio.sleep(OAUTH_MINT_WAIT_INTERVAL_SECONDS) - else: - raise GrantRefreshError(name, "timed out waiting for a concurrent refresh to finish") - try: - grant = McpOauthGrant.get_for_account(name, account_id) - if not grant: + """The delivery-side token for a connected server, minted by the shared + engine from this server's grant — delivery never ships a server the agent + can't authenticate to.""" + grant = McpOauthGrant.get_for_account(name, account_id) + if not grant: + raise MissingGrantError(name, account_id) + + def load_refresh_token() -> str: + # Under the refresh lock: another process may have rotated and + # committed, and this transaction may already hold the row — + # populate_existing re-reads it past the identity map. + fresh = ( + db_session() + .scalars( + select(McpOauthGrant) + .where(McpOauthGrant.id == grant.id) + .execution_options(populate_existing=True) + ) + .one_or_none() + ) + if not fresh: raise MissingGrantError(name, account_id) # The grant's secret halves are ciphertext at rest; the plaintext - # exists only in this request body. - data = { - "grant_type": "refresh_token", - "refresh_token": grant.refresh_token.decrypt(), - "client_id": grant.client_id, - # RFC 8707: an audience-binding server expects the refresh to carry - # the same resource the code exchange was bound to. - "resource": grant.resource, - } - if grant.client_secret: - data["client_secret"] = grant.client_secret.decrypt() - async with _http() as client: - try: - response = await client.post(grant.token_endpoint, data=data) - except httpx.HTTPError as error: - raise GrantRefreshError(name, str(error)) from error - if response.status_code != 200: - await evict_access_token(name, account_id) - raise GrantRefreshError(name, f"HTTP {response.status_code} from the token endpoint") - try: - tokens = response.json() - except ValueError as error: - raise GrantRefreshError(name, "the token endpoint returned malformed JSON") from error - if not isinstance(tokens, dict) or not tokens.get("access_token"): - raise GrantRefreshError(name, "the token endpoint returned no access token") - if tokens.get("refresh_token"): - # Rotation: the provider invalidated the old refresh token on use. - # The write rides the enclosing transaction, so until its commit a - # crash loses it and a concurrent minter in another session still - # reads the spent token (only reachable when the provider's - # expires_in undercuts the run's remaining duration — the cache - # covers the window otherwise). Either way the next mint fails - # loudly and re-connecting replaces the grant; that recovery path - # is the accepted cost of not committing mid-step. - grant.refresh_token = tokens["refresh_token"] - db_session().flush() - try: - ttl = int(tokens.get("expires_in", 3600)) - OAUTH_TOKEN_TTL_SKEW_SECONDS - except (TypeError, ValueError) as error: - raise GrantRefreshError( - name, "the token endpoint returned a malformed expires_in" - ) from error - if ttl > 0: - await redis.set(token_key, tokens["access_token"], ex=ttl) - return tokens["access_token"] - finally: - await redis.delete(lock_key) + # exists only in the refresh request body. + return fresh.refresh_token.decrypt() + + def save_refresh_token(refresh_token: str) -> None: + # The provider invalidated the old token the moment it rotated, so + # the write commits on its own session, never the enclosing + # transaction — a step that rolls back later must not brick the grant. + with get_session(db_session().get_bind()) as session: + session.execute( + update(McpOauthGrant) + .where(McpOauthGrant.id == grant.id) + .values(refresh_token=refresh_token) + ) + session.commit() + # Keep the enclosing transaction's copy true as well. + grant.refresh_token = refresh_token + db_session().flush() + + client = OauthClient( + provider=OAUTH_PROVIDER, + token_endpoint=grant.token_endpoint, + client_id=grant.client_id, + client_secret=grant.client_secret.decrypt(), + # RFC 8707: an audience-binding server expects the refresh to carry + # the same resource the code exchange was bound to. + extra_token_params={"resource": grant.resource}, + mint_wait_interval_seconds=OAUTH_MINT_WAIT_INTERVAL_SECONDS, + mint_wait_attempts=OAUTH_MINT_WAIT_ATTEMPTS, + http_factory=_http, + ) + try: + return await client.mint_access_token( + key=f"{name}:{account_id}", + load_refresh_token=load_refresh_token, + save_refresh_token=save_refresh_token, + ) + except OauthRefreshError as error: + raise GrantRefreshError(name, error.reason) from error diff --git a/backend/druks/services/__init__.py b/backend/druks/services/__init__.py index 17238057..0ee081b9 100644 --- a/backend/druks/services/__init__.py +++ b/backend/druks/services/__init__.py @@ -1,4 +1,17 @@ from .base import Service -from .exceptions import ServiceConnectError, ServiceNotConnectedError +from .exceptions import ( + OauthExchangeError, + OauthRefreshError, + ServiceConnectError, + ServiceNotConnectedError, +) +from .oauth import OauthClient -__all__ = ["Service", "ServiceConnectError", "ServiceNotConnectedError"] +__all__ = [ + "OauthClient", + "OauthExchangeError", + "OauthRefreshError", + "Service", + "ServiceConnectError", + "ServiceNotConnectedError", +] diff --git a/backend/druks/services/base.py b/backend/druks/services/base.py index 067644f4..6b0b2377 100644 --- a/backend/druks/services/base.py +++ b/backend/druks/services/base.py @@ -8,6 +8,7 @@ from .exceptions import ServiceConnectError, ServiceNotConnectedError from .models import ServiceIdentity +from .oauth import OauthClient class Service: @@ -28,6 +29,14 @@ class Service: # Whether doctor fails when this service is not connected. required: ClassVar[bool] = True settings_model: ClassVar[type[BaseModel]] + # Set both endpoints when the registered app is an OAuth client; + # ``get_oauth_client()`` then hands back the connected identity as a + # configured ``OauthClient``. Scopes are not declared here — each + # ``begin_connect`` asks for its own. + authorization_endpoint: ClassVar[str] = "" + token_endpoint: ClassVar[str] = "" + # HTTP Basic on the token endpoint; False sends the secret in the body. + basic_auth: ClassVar[bool] = False def __init_subclass__(cls, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) @@ -44,6 +53,13 @@ def __init_subclass__(cls, **kwargs: Any) -> None: declared = cls.__dict__.get("Settings") if not isinstance(declared, type) or not issubclass(declared, BaseModel): raise TypeError(f"{cls.__name__}.Settings must be a pydantic model") + if bool(cls.authorization_endpoint) != bool(cls.token_endpoint): + raise TypeError(f"{cls.__name__} must declare both OAuth endpoints or neither") + if cls.token_endpoint and not {"client_id", "client_secret"} <= set(declared.model_fields): + raise TypeError( + f"{cls.__name__}.Settings must declare client_id and client_secret " + "fields — get_oauth_client() reads the OAuth client from them" + ) cls.settings_model = declared services.register(cls) @@ -74,6 +90,23 @@ def connect_fields(cls) -> list[dict[str, Any]]: def get(cls) -> ServiceIdentity: return ServiceIdentity.get(cls.name) + @classmethod + def get_oauth_client(cls) -> OauthClient: + """The connected identity as a configured ``OauthClient``, keyed by + the service name. Raises ``ServiceNotConnectedError`` until the + operator connects the service.""" + if not cls.token_endpoint: + raise TypeError(f"{cls.__name__} declares no OAuth endpoints") + connected = cls.get() + return OauthClient( + provider=cls.name, + authorization_endpoint=cls.authorization_endpoint, + token_endpoint=cls.token_endpoint, + client_id=connected.identity["client_id"], + client_secret=connected.secrets["client_secret"], + basic_auth=cls.basic_auth, + ) + @classmethod def is_connected(cls) -> bool: try: diff --git a/backend/druks/services/constants.py b/backend/druks/services/constants.py new file mode 100644 index 00000000..916db940 --- /dev/null +++ b/backend/druks/services/constants.py @@ -0,0 +1,16 @@ +# OAuth connect + mint plumbing shared by every OauthClient consumer. Pending +# connect state (the PKCE verifier plus the begun flow's client identity) +# lives in Redis for its short TTL, single-use. Access tokens cache for their +# lifetime minus the skew, so a token handed out never expires moments after +# delivery. +OAUTH_CONNECT_STATE_TTL_SECONDS = 600 +OAUTH_TOKEN_TTL_SKEW_SECONDS = 60 + +# Mint's mutual exclusion, in the Redis that fronts the token cache (SET NX): +# a rotating grant tolerates exactly one refresher. The lock TTL is a crash +# backstop at three times the HTTP client's timeout — a live refresh cannot +# outlive it. Losers poll the cache on the interval for about one +# token-endpoint round trip, then fail loudly. +OAUTH_REFRESH_LOCK_TTL_SECONDS = 90 +OAUTH_MINT_WAIT_INTERVAL_SECONDS = 0.2 +OAUTH_MINT_WAIT_ATTEMPTS = 150 diff --git a/backend/druks/services/exceptions.py b/backend/druks/services/exceptions.py index f09bdadd..0f2d2ff1 100644 --- a/backend/druks/services/exceptions.py +++ b/backend/druks/services/exceptions.py @@ -10,3 +10,28 @@ def __init__(self, service: str) -> None: class ServiceConnectError(Exception): """A rejected connect. The message is authored by the service's ``verify`` and safe to show; it never quotes anything the operator pasted.""" + + +class OauthExchangeError(Exception): + """Completing an OAuth connect flow failed — an unknown or expired state, + or a rejected code exchange. Nothing is stored on failure, so re-running + the connect flow is always safe. ``context`` is the begun flow's stash + when the state resolved, and empty when it did not.""" + + def __init__(self, provider: str, reason: str, *, context: dict) -> None: + super().__init__(f"OAuth exchange for {provider!r} failed: {reason}") + self.provider = provider + self.reason = reason + self.context = context + + +class OauthRefreshError(Exception): + """Minting an access token from a stored grant failed — the provider + rejected the refresh token, the token endpoint is unreachable, or a + concurrent refresh never freed the lock. Re-connecting replaces the + grant.""" + + def __init__(self, provider: str, reason: str) -> None: + super().__init__(f"OAuth refresh for {provider!r} failed: {reason}") + self.provider = provider + self.reason = reason diff --git a/backend/druks/services/oauth.py b/backend/druks/services/oauth.py new file mode 100644 index 00000000..ea3b1fe0 --- /dev/null +++ b/backend/druks/services/oauth.py @@ -0,0 +1,303 @@ +import asyncio +import base64 +import hashlib +import json +import secrets +from collections.abc import Callable +from typing import Any, cast +from urllib.parse import urlencode + +import httpx + +from druks.redis import get_client + +from .constants import ( + OAUTH_CONNECT_STATE_TTL_SECONDS, + OAUTH_MINT_WAIT_ATTEMPTS, + OAUTH_MINT_WAIT_INTERVAL_SECONDS, + OAUTH_REFRESH_LOCK_TTL_SECONDS, + OAUTH_TOKEN_TTL_SKEW_SECONDS, +) +from .exceptions import OauthExchangeError, OauthRefreshError + + +def _http() -> httpx.AsyncClient: + # One construction point so a suite can swap in a MockTransport client. + return httpx.AsyncClient(timeout=30.0, follow_redirects=True) + + +async def _post_token( + http: httpx.AsyncClient, + token_endpoint: str, + data: dict[str, Any], + *, + client_id: str, + client_secret: str, + basic_auth: bool, +) -> httpx.Response: + # RFC 6749 client authentication: HTTP Basic keeps the credentials out of + # the form body; a public or body-authenticating client sends them in it. + if basic_auth: + return await http.post(token_endpoint, data=data, auth=(client_id, client_secret)) + data["client_id"] = client_id + if client_secret: + data["client_secret"] = client_secret + return await http.post(token_endpoint, data=data) + + +class OauthClient: + """One provider's OAuth 2.0 authorization-code + PKCE flow, with + rotation-safe refresh — for a provider with fixed endpoints and a + pre-registered client:: + + client = OauthClient( + provider="acme", + authorization_endpoint="https://acme.example/oauth/authorize", + token_endpoint="https://acme.example/oauth/token", + client_id=..., client_secret=..., + basic_auth=True, + ) + + ``begin_connect`` returns the consent URL to open; ``complete_connect`` + consumes the callback's single-use state and exchanges the code; + ``mint_access_token`` serves delivery from the Redis token cache, electing + one refresher per grant. Grants live on the caller's own rows, reached + through the two callables mint takes. A ``Service`` with declared OAuth + endpoints hands back a configured client via ``get_oauth_client()`` — + construct directly only when no service holds the client credentials. + + ``provider`` keys every Redis entry — connect state, token cache, refresh + lock — so all clients constructed with one provider name share them, and + across a rolling deploy old and new processes elect the same single + refresher. Completion needs only ``provider``: the begun flow's endpoints + and client identity ride the stashed state, pinned at begin time so a + configuration change mid-consent cannot mismatch the PKCE verifier. + + ``basic_auth`` picks HTTP Basic on the token endpoint, for both the code + exchange and refresh; without it the client credentials travel in the form + body. ``extra_token_params`` land in both bodies (RFC 8707's ``resource`` + audience binding). Scopes are per authorization, not per client — each + ``begin_connect`` asks for its own, and the grant keeps what was approved. + """ + + def __init__( + self, + *, + provider: str, + authorization_endpoint: str = "", + token_endpoint: str = "", + client_id: str = "", + client_secret: str = "", + basic_auth: bool = False, + extra_token_params: dict[str, str] | None = None, + mint_wait_interval_seconds: float = OAUTH_MINT_WAIT_INTERVAL_SECONDS, + mint_wait_attempts: int = OAUTH_MINT_WAIT_ATTEMPTS, + http_factory: Callable[[], httpx.AsyncClient] | None = None, + ) -> None: + self.provider = provider + self.authorization_endpoint = authorization_endpoint + self.token_endpoint = token_endpoint + self.client_id = client_id + self.client_secret = client_secret + self.basic_auth = basic_auth + self.extra_token_params = dict(extra_token_params or {}) + self.mint_wait_interval_seconds = mint_wait_interval_seconds + self.mint_wait_attempts = mint_wait_attempts + self._http = http_factory or _http + + async def begin_connect( + self, + *, + redirect_uri: str, + scopes: tuple[str, ...] = (), + context: dict[str, Any] | None = None, + extra_authorize_params: dict[str, str] | None = None, + ) -> str: + """Stash the pending exchange in Redis under a fresh single-use state + and return the consent URL to open. ``scopes`` render into the consent + query — this authorization's ask, within whatever ceiling the provider + registration allows. ``context`` rides the stash and comes back from + ``complete_connect``; ``extra_authorize_params`` land in the consent + query. Nothing durable is written here — an abandoned consent simply + expires.""" + state = secrets.token_urlsafe(32) + code_verifier = secrets.token_urlsafe(64) + code_challenge = ( + base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest()) + .rstrip(b"=") + .decode() + ) + pending = { + **(context or {}), + "code_verifier": code_verifier, + "redirect_uri": redirect_uri, + "token_endpoint": self.token_endpoint, + "client_id": self.client_id, + "client_secret": self.client_secret, + "basic_auth": self.basic_auth, + "extra_token_params": self.extra_token_params, + } + await get_client().set( + f"{self.provider}:connect:{state}", + json.dumps(pending), + ex=OAUTH_CONNECT_STATE_TTL_SECONDS, + ) + query = { + "response_type": "code", + "client_id": self.client_id, + "redirect_uri": redirect_uri, + "state": state, + "code_challenge": code_challenge, + "code_challenge_method": "S256", + } + if scopes: + query["scope"] = " ".join(scopes) + query.update(extra_authorize_params or {}) + return f"{self.authorization_endpoint}?{urlencode(query)}" + + async def complete_connect(self, *, state: str, code: str) -> tuple[dict, dict]: + """The callback half: consume the pending state (single-use, GETDEL) + and exchange the code + verifier for tokens. Returns ``(tokens, + context)`` — the token response and the begun flow's stash, the + caller's begin-time context with the flow's ``token_endpoint``, + ``client_id``, and ``client_secret`` merged in. A response without a + ``refresh_token`` is rejected: a grant must survive offline. Nothing + is cached here — the caller's grant becomes real only when its own + write commits, and the first mint refreshes from it.""" + raw = await get_client().getdel(f"{self.provider}:connect:{state}") + if not raw: + raise OauthExchangeError( + self.provider, + "unknown or expired state; start the connect flow again", + context={}, + ) + pending = json.loads(raw) + data = { + "grant_type": "authorization_code", + "code": code, + "redirect_uri": pending["redirect_uri"], + "code_verifier": pending["code_verifier"], + **pending["extra_token_params"], + } + async with self._http() as http: + try: + response = await _post_token( + http, + pending["token_endpoint"], + data, + client_id=pending["client_id"], + client_secret=pending["client_secret"], + basic_auth=pending["basic_auth"], + ) + except httpx.HTTPError as error: + raise OauthExchangeError( + self.provider, f"code exchange failed: {error}", context=pending + ) from error + if response.status_code != 200: + raise OauthExchangeError( + self.provider, + f"code exchange failed: HTTP {response.status_code}", + context=pending, + ) + try: + tokens = response.json() + except ValueError as error: + raise OauthExchangeError( + self.provider, "the token endpoint returned malformed JSON", context=pending + ) from error + if not isinstance(tokens, dict) or not tokens.get("refresh_token"): + raise OauthExchangeError( + self.provider, + "the authorization server granted no refresh token; druks needs offline access", + context=pending, + ) + return tokens, pending + + async def mint_access_token( + self, + *, + key: str, + load_refresh_token: Callable[[], str], + save_refresh_token: Callable[[str], None], + ) -> str: + """The delivery-side token for one grant: the cached access token + while it lives, else one refreshed through the grant's refresh token. + The provider may rotate the refresh token on use — two concurrent + refreshes trip its reuse detection and can revoke the whole grant — + so Redis elects one refresher per ``key`` (SET NX; the TTL is a crash + backstop a live refresh cannot outlive). Losers poll for the winner's + cache fill, for about one token-endpoint round trip, then fail loudly. + + ``load_refresh_token()`` runs under the refresh lock and must observe + rotations other processes committed — a naive re-select can return a + row this transaction already identity-mapped, so read with + ``populate_existing`` or on a fresh session. + + ``save_refresh_token(token)`` receives a rotated refresh token and + must have committed it before returning: the provider has already + invalidated the old token, so the write cannot ride an enclosing + transaction that may later roll back. The cache fills only after it + returns.""" + redis = get_client() + token_key = f"{self.provider}:access_token:{key}" + lock_key = f"{self.provider}:refresh_lock:{key}" + for _ in range(self.mint_wait_attempts): + cached = await redis.get(token_key) + if cached: + return cast(bytes, cached).decode() + if await redis.set(lock_key, "1", nx=True, ex=OAUTH_REFRESH_LOCK_TTL_SECONDS): + break + await asyncio.sleep(self.mint_wait_interval_seconds) + else: + raise OauthRefreshError( + self.provider, "timed out waiting for a concurrent refresh to finish" + ) + try: + data = { + "grant_type": "refresh_token", + "refresh_token": load_refresh_token(), + **self.extra_token_params, + } + async with self._http() as http: + try: + response = await _post_token( + http, + self.token_endpoint, + data, + client_id=self.client_id, + client_secret=self.client_secret, + basic_auth=self.basic_auth, + ) + except httpx.HTTPError as error: + raise OauthRefreshError(self.provider, str(error)) from error + if response.status_code != 200: + await redis.delete(token_key) + raise OauthRefreshError( + self.provider, f"HTTP {response.status_code} from the token endpoint" + ) + try: + tokens = response.json() + except ValueError as error: + raise OauthRefreshError( + self.provider, "the token endpoint returned malformed JSON" + ) from error + if not isinstance(tokens, dict) or not tokens.get("access_token"): + raise OauthRefreshError( + self.provider, "the token endpoint returned no access token" + ) + if tokens.get("refresh_token"): + save_refresh_token(tokens["refresh_token"]) + try: + ttl = int(tokens.get("expires_in", 3600)) - OAUTH_TOKEN_TTL_SKEW_SECONDS + except (TypeError, ValueError) as error: + raise OauthRefreshError( + self.provider, "the token endpoint returned a malformed expires_in" + ) from error + if ttl > 0: + await redis.set(token_key, tokens["access_token"], ex=ttl) + return tokens["access_token"] + finally: + await redis.delete(lock_key) + + async def evict_access_token(self, key: str) -> None: + await get_client().delete(f"{self.provider}:access_token:{key}") diff --git a/backend/tests/test_author_surface.py b/backend/tests/test_author_surface.py index a097c86e..ca751f8e 100644 --- a/backend/tests/test_author_surface.py +++ b/backend/tests/test_author_surface.py @@ -7,7 +7,14 @@ AUTHOR_SURFACE = { "druks.extensions": {"Extension", "ExtensionSettings", "Secret"}, "druks.browser": {"BrowserSession", "BrowserSessionSignedOutError"}, - "druks.services": {"Service", "ServiceConnectError", "ServiceNotConnectedError"}, + "druks.services": { + "OauthClient", + "OauthExchangeError", + "OauthRefreshError", + "Service", + "ServiceConnectError", + "ServiceNotConnectedError", + }, "druks.agents": {"Agent", "AgentOutput"}, "druks.workflows": { "AgentCall", diff --git a/backend/tests/test_oauth_client.py b/backend/tests/test_oauth_client.py new file mode 100644 index 00000000..ef59eff1 --- /dev/null +++ b/backend/tests/test_oauth_client.py @@ -0,0 +1,210 @@ +import asyncio +from urllib.parse import parse_qsl, urlparse + +import httpx +import pytest +from druks.redis import get_client +from druks.services import OauthClient, OauthExchangeError, OauthRefreshError + +_PROVIDER = "acme" +_AUTHORIZATION_ENDPOINT = "https://auth.acme.test/authorize" +_TOKEN_ENDPOINT = "https://auth.acme.test/token" +_REDIRECT_URI = "https://druks.example/api/acme/oauth/callback" +_TOKEN_KEY = f"{_PROVIDER}:access_token:grant-1" +_LOCK_KEY = f"{_PROVIDER}:refresh_lock:grant-1" + + +class FakeTokenEndpoint: + def __init__(self) -> None: + self.status = 200 + self.response = {"access_token": "at-1", "refresh_token": "rt-1", "expires_in": 3600} + self.requests: list[dict] = [] + self.authorizations: list[str] = [] + + def handler(self, request: httpx.Request) -> httpx.Response: + self.requests.append(dict(parse_qsl(request.content.decode()))) + self.authorizations.append(request.headers.get("Authorization", "")) + return httpx.Response(self.status, json=self.response) + + +@pytest.fixture +def token_endpoint(): + return FakeTokenEndpoint() + + +def _client(token_endpoint: FakeTokenEndpoint, **overrides) -> OauthClient: + kwargs = dict( + provider=_PROVIDER, + authorization_endpoint=_AUTHORIZATION_ENDPOINT, + token_endpoint=_TOKEN_ENDPOINT, + client_id="client-123", + client_secret="secret-123", + mint_wait_interval_seconds=0, + http_factory=lambda: httpx.AsyncClient( + transport=httpx.MockTransport(token_endpoint.handler) + ), + ) + kwargs.update(overrides) + return OauthClient(**kwargs) + + +def _fail_save(token: str) -> None: + pytest.fail("nothing rotated") + + +async def test_mint_serves_the_cache_without_a_refresh(token_endpoint): + await get_client().set(_TOKEN_KEY, "at-cached") + + token = await _client(token_endpoint).mint_access_token( + key="grant-1", + load_refresh_token=lambda: pytest.fail("cache hit reads no grant"), + save_refresh_token=_fail_save, + ) + + assert token == "at-cached" + assert not token_endpoint.requests + + +async def test_mint_refreshes_persists_rotation_and_fills_with_skewed_ttl(token_endpoint): + token_endpoint.response = {"access_token": "at-2", "refresh_token": "rt-new", "expires_in": 300} + saved: list[str] = [] + + token = await _client(token_endpoint).mint_access_token( + key="grant-1", + load_refresh_token=lambda: "rt-old", + save_refresh_token=saved.append, + ) + + assert token == "at-2" + assert saved == ["rt-new"] + refresh = token_endpoint.requests[0] + assert refresh["grant_type"] == "refresh_token" + assert refresh["refresh_token"] == "rt-old" + # Body-style client auth: the credentials travel in the form. + assert refresh["client_id"] == "client-123" + assert refresh["client_secret"] == "secret-123" + redis = get_client() + assert await redis.get(_TOKEN_KEY) == b"at-2" + assert 0 < await redis.ttl(_TOKEN_KEY) <= 240 + assert not await redis.get(_LOCK_KEY) + + +async def test_mint_fills_the_cache_only_after_the_rotation_is_saved(token_endpoint): + token_endpoint.response = {"access_token": "at-2", "refresh_token": "rt-new", "expires_in": 300} + + def save_refresh_token(token: str) -> None: + raise RuntimeError("rotation write failed") + + with pytest.raises(RuntimeError, match="rotation write failed"): + await _client(token_endpoint).mint_access_token( + key="grant-1", + load_refresh_token=lambda: "rt-old", + save_refresh_token=save_refresh_token, + ) + + redis = get_client() + assert not await redis.get(_TOKEN_KEY) + assert not await redis.get(_LOCK_KEY) + + +async def test_mint_losing_the_lock_polls_for_the_winners_token(token_endpoint): + redis = get_client() + await redis.set(_LOCK_KEY, "1") + + async def _winner_finishes(): + await redis.set(_TOKEN_KEY, "at-winner") + await redis.delete(_LOCK_KEY) + + winner = asyncio.create_task(_winner_finishes()) + token = await _client(token_endpoint).mint_access_token( + key="grant-1", + load_refresh_token=lambda: "rt-old", + save_refresh_token=_fail_save, + ) + await winner + + assert token == "at-winner" + assert not token_endpoint.requests + + +async def test_mint_times_out_loudly_when_the_lock_never_frees(token_endpoint): + await get_client().set(_LOCK_KEY, "1") + + with pytest.raises(OauthRefreshError, match="concurrent refresh"): + await _client(token_endpoint, mint_wait_attempts=3).mint_access_token( + key="grant-1", + load_refresh_token=lambda: "rt-old", + save_refresh_token=_fail_save, + ) + + +async def test_mint_refresh_rejection_evicts_and_raises(token_endpoint): + token_endpoint.status = 400 + + with pytest.raises(OauthRefreshError, match="HTTP 400"): + await _client(token_endpoint).mint_access_token( + key="grant-1", + load_refresh_token=lambda: "rt-old", + save_refresh_token=_fail_save, + ) + + redis = get_client() + assert not await redis.get(_TOKEN_KEY) + assert not await redis.get(_LOCK_KEY) + + +async def test_mint_refresh_uses_basic_auth(token_endpoint): + await _client(token_endpoint, basic_auth=True).mint_access_token( + key="grant-1", + load_refresh_token=lambda: "rt-old", + save_refresh_token=lambda token: None, + ) + + assert token_endpoint.authorizations[0].startswith("Basic ") + # Basic auth keeps the client credentials out of the form body. + assert "client_id" not in token_endpoint.requests[0] + assert "client_secret" not in token_endpoint.requests[0] + + +async def test_connect_roundtrip_exchanges_with_basic_auth(token_endpoint): + url = await _client(token_endpoint, basic_auth=True).begin_connect( + redirect_uri=_REDIRECT_URI, + scopes=("profile.read", "posts.write"), + context={"account": "a-1"}, + extra_authorize_params={"audience": "api"}, + ) + + assert url.startswith(f"{_AUTHORIZATION_ENDPOINT}?") + params = dict(parse_qsl(urlparse(url).query)) + assert params["scope"] == "profile.read posts.write" + assert params["audience"] == "api" + assert params["code_challenge_method"] == "S256" + + # Completion needs only the provider: the begun flow's client identity + # rides the stashed state. + tokens, context = await OauthClient( + provider=_PROVIDER, + http_factory=lambda: httpx.AsyncClient( + transport=httpx.MockTransport(token_endpoint.handler) + ), + ).complete_connect(state=params["state"], code="code-1") + + assert tokens["refresh_token"] == "rt-1" + assert context["account"] == "a-1" + assert context["client_id"] == "client-123" + exchange = token_endpoint.requests[0] + assert exchange["grant_type"] == "authorization_code" + assert exchange["code"] == "code-1" + assert exchange["code_verifier"] + assert "client_id" not in exchange + assert "client_secret" not in exchange + assert token_endpoint.authorizations[0].startswith("Basic ") + + +async def test_complete_connect_requires_a_refresh_token(token_endpoint): + token_endpoint.response = {"access_token": "at-1", "expires_in": 3600} + url = await _client(token_endpoint).begin_connect(redirect_uri=_REDIRECT_URI) + state = dict(parse_qsl(urlparse(url).query))["state"] + + with pytest.raises(OauthExchangeError, match="no refresh token"): + await _client(token_endpoint).complete_connect(state=state, code="code-1") diff --git a/backend/tests/test_services.py b/backend/tests/test_services.py index ccaac22e..1aa0f2fb 100644 --- a/backend/tests/test_services.py +++ b/backend/tests/test_services.py @@ -367,3 +367,84 @@ async def fake_post(self, url, **kwargs): assert "restart" in response.json()["detail"] with pytest.raises(ServiceNotConnectedError): ServiceIdentity.get("github") + + +# --- OAuth declaration -------------------------------------------------------- + + +@pytest.fixture +def declared_services(): + # Service subclasses self-register at class definition; tests declare + # inside this fixture and leave the registry as found. + from druks.extensions.registry import services + + saved = dict(services._items) + yield + services._items.clear() + services._items.update(saved) + + +def test_get_oauth_client_reads_the_connected_identity(declared_services, druks_db): + from druks.services import Service + from pydantic import BaseModel, SecretStr + + class Acme(Service): + name = "acme" + title = "Acme OAuth app" + authorization_endpoint = "https://acme.test/authorize" + token_endpoint = "https://acme.test/token" + basic_auth = True + + class Settings(BaseModel): + client_id: str + client_secret: SecretStr + + ServiceIdentity.connect( + "acme", identity={"client_id": "id-1"}, secrets={"client_secret": "sec-1"} + ) + + client = Acme.get_oauth_client() + + assert client.provider == "acme" + assert client.authorization_endpoint == "https://acme.test/authorize" + assert client.token_endpoint == "https://acme.test/token" + assert client.client_id == "id-1" + assert client.client_secret == "sec-1" + assert client.basic_auth is True + + +def test_oauth_service_declarations_fail_loudly(declared_services): + from druks.services import Service + from pydantic import BaseModel, SecretStr + + with pytest.raises(TypeError, match="client_id and client_secret"): + + class Keyless(Service): + name = "keyless" + title = "Keyless" + authorization_endpoint = "https://acme.test/authorize" + token_endpoint = "https://acme.test/token" + + class Settings(BaseModel): + api_key: SecretStr + + with pytest.raises(TypeError, match="both OAuth endpoints"): + + class HalfDeclared(Service): + name = "half_declared" + title = "Half declared" + token_endpoint = "https://acme.test/token" + + class Settings(BaseModel): + client_id: str + client_secret: SecretStr + + class Plain(Service): + name = "plain_service" + title = "Plain" + + class Settings(BaseModel): + api_key: SecretStr + + with pytest.raises(TypeError, match="no OAuth endpoints"): + Plain.get_oauth_client() diff --git a/docs/writing-an-extension.md b/docs/writing-an-extension.md index e792f219..6c6dc0bb 100644 --- a/docs/writing-an-extension.md +++ b/docs/writing-an-extension.md @@ -731,6 +731,101 @@ its own service, and the operator decides per card whether the underlying registration is shared or a narrower one — that choice is their scope and blast-radius control. +## Connect provider accounts (OAuth) + +`OauthClient` runs the OAuth 2.0 authorization-code + PKCE flow. Use it for a +provider with fixed endpoints and a registered client. It mints access tokens +and keeps refresh-token rotation safe. Your extension stores each grant on its +own rows. The platform stores no grants. + +Declare the endpoints on the service that holds the client credentials. The +`Settings` model must have `client_id` and `client_secret` fields: + +```python +class Acme(Service): + name = "acme" + title = "Acme OAuth app" + authorization_endpoint = "https://acme.example/oauth/authorize" + token_endpoint = "https://acme.example/oauth/token" + # True = HTTP Basic on the token endpoint. False = secret in the body. + basic_auth = True + + class Settings(BaseModel): + client_id: str = Field(title="Client ID") + client_secret: SecretStr = Field(title="Client secret") +``` + +`Acme.get_oauth_client()` returns a configured client for the connected +identity. Call `begin_connect` from your connect route. Call `complete_connect` +from your callback route: + +```python +url = await Acme.get_oauth_client().begin_connect( + redirect_uri="https://druks.example/api/acme/oauth/callback", + scopes=("profile.read", "posts.write"), + context={"account_id": account_id}, +) +# ... the operator consents; the provider redirects back with state + code ... +tokens, context = await Acme.get_oauth_client().complete_connect(state=state, code=code) +AcmeGrant.store(account_id=context["account_id"], refresh_token=tokens["refresh_token"]) +``` + +Scopes belong to one authorization, not to the service. Each `begin_connect` +call asks for its own scopes. The provider registration sets the ceiling. The +grant keeps the scopes the user approved. + +`begin_connect` stores the pending exchange in Redis and returns the consent +URL. The state is single-use and expires after a short time. `complete_connect` +consumes the state and exchanges the code for tokens. It rejects a token +response without a `refresh_token`, because a grant must work offline. It +raises `OauthExchangeError` when the flow is denied or expired. It returns your +`context` from begin time, with the flow's client identity merged in. + +When a run needs the provider, call `mint_access_token`. The engine serves +tokens from a Redis cache. It lets only one refresher run for each `key`. This +is necessary: two refreshes at the same time can make the provider revoke the +whole grant. The engine raises `OauthRefreshError` when the grant does not +refresh. Then ask the operator to connect again. + +```python +token = await Acme.get_oauth_client().mint_access_token( + key=account_id, + load_refresh_token=load_refresh_token, + save_refresh_token=save_refresh_token, +) +``` + +The two callables connect the engine to your storage: + +```python +def load_refresh_token() -> str: + # Runs under the refresh lock. Another process may have rotated and + # committed. Re-read the row; do not trust the identity map. + grant = db_session().scalars( + select(AcmeGrant) + .where(AcmeGrant.account_id == account_id) + .execution_options(populate_existing=True) + ).one() + return grant.secrets["refresh_token"] + + +def save_refresh_token(rotated: str) -> None: + # The provider has already invalidated the old token. Commit on an own + # session, never on the enclosing step transaction. A later rollback + # must not lose the new token. + with Session(db_session().get_bind()) as session: + grant = session.scalars( + select(AcmeGrant).where(AcmeGrant.account_id == account_id) + ).one() + grant.secrets["refresh_token"] = rotated + session.commit() +``` + +`secrets` is an `EncryptedJsonField` column on your grant row (see +[models](#models-and-migrations)). The refresh token is ciphertext at rest. +When the operator disconnects, delete your row and evict the cached access +token: `await Acme.get_oauth_client().evict_access_token(account_id)`. + ## Extension settings and checks An inner `ExtensionSettings` class defines dashboard-editable knobs and owns their @@ -906,7 +1001,8 @@ Import from concern namespaces, not from `druks.durable` or internal modules: | Namespace | Public names | | --- | --- | | `druks.extensions` | `Extension`, `ExtensionSettings`, `Secret` | -| `druks.services` | `Service`, `ServiceConnectError`, `ServiceNotConnectedError` | +| `druks.services` | `Service`, `ServiceConnectError`, `ServiceNotConnectedError`, `OauthClient`, `OauthExchangeError`, `OauthRefreshError` | +| `druks.secrets.fields` | `EncryptedJsonField`, `SecretsMapping` | | `druks.agents` | `Agent`, `AgentOutput` | | `druks.workflows` | `Workflow`, `Gate`, `step`, run/agent response types, lifecycle enums and workflow errors | | `druks.db` | `Base`, `StoredSubject`, `db_session` |