nihalashetty
ae67bff5a3
Self-hosted platform for building, testing, and shipping LangChain/LangGraph agents. Deep-agent sub-agents on the canvas, a live tracing/observability timeline, auto-provisioned built-in tools with import/export, per-environment tool variables, streamed evaluations, and per-user auth token forwarding.
118 lines
6.0 KiB
Python
118 lines
6.0 KiB
Python
"""OAuth 2.1 for MCP: discovery, dynamic client registration, authorization-code + PKCE, single-use
|
|
codes, audience binding, and token validation on the MCP endpoint. Gated by mcp_oauth_enabled."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import hashlib
|
|
import os
|
|
import urllib.parse
|
|
import uuid
|
|
|
|
import httpx
|
|
|
|
from forge.config import settings
|
|
from forge.main import create_app
|
|
|
|
|
|
def _pkce() -> tuple[str, str]:
|
|
verifier = base64.urlsafe_b64encode(os.urandom(32)).rstrip(b"=").decode()
|
|
challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode()
|
|
return verifier, challenge
|
|
|
|
|
|
async def _authorize_code(c, *, email, password, client_id, redirect_uri, challenge, resource) -> str:
|
|
form = {
|
|
"email": email, "password": password, "workspace_id": "", "client_id": client_id,
|
|
"redirect_uri": redirect_uri, "code_challenge": challenge, "state": "s", "resource": resource, "scope": "",
|
|
}
|
|
sub = await c.post("/v1/oauth/authorize", data=form, follow_redirects=False)
|
|
assert sub.status_code == 302, sub.text
|
|
q = urllib.parse.parse_qs(urllib.parse.urlparse(sub.headers["location"]).query)
|
|
assert q.get("state") == ["s"]
|
|
return q["code"][0]
|
|
|
|
|
|
async def test_mcp_oauth_disabled_returns_404():
|
|
app = create_app()
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
|
|
assert (await c.get("/.well-known/oauth-authorization-server")).status_code == 404
|
|
|
|
|
|
async def test_mcp_oauth_end_to_end(monkeypatch):
|
|
monkeypatch.setattr(settings, "mcp_oauth_enabled", True)
|
|
app = create_app()
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
|
|
# a real user + a project (in the user's tenant) with one tool
|
|
email = f"o{uuid.uuid4().hex[:10]}@example.com"
|
|
reg = await c.post("/v1/auth/register", json={"email": email, "password": "supersecret1"})
|
|
assert reg.status_code == 201, reg.text
|
|
c.headers["Authorization"] = f"Bearer {reg.json()['access_token']}"
|
|
pid = (await c.post("/v1/projects", json={"name": "OAuth", "slug": "oauth-proj"})).json()["id"]
|
|
tid = (await c.post(f"/v1/projects/{pid}/tools", json={"name": "calc", "kind": "builtin", "config": {"builtin": "calculator", "description": "c"}})).json()["id"]
|
|
# publish the tool via an exposed tool set (the MCP surface = exposed sets' enabled tools)
|
|
await c.post(f"/v1/projects/{pid}/tool-sets", json={"name": "General", "tool_ids": [tid]})
|
|
|
|
# discovery is live when enabled
|
|
asm = (await c.get("/.well-known/oauth-authorization-server")).json()
|
|
assert asm["code_challenge_methods_supported"] == ["S256"]
|
|
prm = (await c.get(f"/.well-known/oauth-protected-resource/v1/mcp/{pid}")).json()
|
|
assert prm["resource"].endswith(f"/v1/mcp/{pid}") and prm["authorization_servers"]
|
|
|
|
# dynamic client registration (RFC 7591)
|
|
redirect_uri = "http://localhost/callback"
|
|
rc = await c.post("/v1/oauth/register", json={"redirect_uris": [redirect_uri], "client_name": "Test client"})
|
|
assert rc.status_code == 201, rc.text
|
|
client_id = rc.json()["client_id"]
|
|
|
|
# the consent form renders
|
|
verifier, challenge = _pkce()
|
|
resource = f"{settings.public_base_url.rstrip('/')}/v1/mcp/{pid}"
|
|
gf = await c.get("/v1/oauth/authorize", params={
|
|
"response_type": "code", "client_id": client_id, "redirect_uri": redirect_uri,
|
|
"code_challenge": challenge, "code_challenge_method": "S256", "resource": resource, "state": "s",
|
|
})
|
|
assert gf.status_code == 200 and "Authorize" in gf.text
|
|
|
|
# authorization-code exchange (PKCE verifier) -> access token
|
|
code = await _authorize_code(c, email=email, password="supersecret1", client_id=client_id,
|
|
redirect_uri=redirect_uri, challenge=challenge, resource=resource)
|
|
tok = await c.post("/v1/oauth/token", data={
|
|
"grant_type": "authorization_code", "code": code, "redirect_uri": redirect_uri,
|
|
"client_id": client_id, "code_verifier": verifier,
|
|
})
|
|
assert tok.status_code == 200, tok.text
|
|
access = tok.json()["access_token"]
|
|
assert tok.json()["token_type"] == "Bearer"
|
|
|
|
# the audience-bound token authorizes the MCP endpoint
|
|
body = {"jsonrpc": "2.0", "id": 1, "method": "tools/list"}
|
|
oauth_headers = {"Authorization": f"Bearer {access}"}
|
|
r = await c.post(f"/v1/mcp/{pid}", headers=oauth_headers, json=body)
|
|
assert r.status_code == 200 and any(t["name"] == "calc" for t in r.json()["result"]["tools"])
|
|
|
|
# no credential -> 401 with an RFC 9728 discovery pointer
|
|
no = await c.post(f"/v1/mcp/{pid}", json=body)
|
|
assert no.status_code == 401 and "resource_metadata" in no.headers.get("www-authenticate", "")
|
|
|
|
# audience binding: the token must not work on a different project
|
|
pid2 = (await c.post("/v1/projects", json={"name": "Other", "slug": "oauth-other"})).json()["id"]
|
|
assert (await c.post(f"/v1/mcp/{pid2}", headers=oauth_headers, json=body)).status_code == 401
|
|
|
|
# single-use: replaying the same authorization code is rejected
|
|
replay = await c.post("/v1/oauth/token", data={
|
|
"grant_type": "authorization_code", "code": code, "redirect_uri": redirect_uri,
|
|
"client_id": client_id, "code_verifier": verifier,
|
|
})
|
|
assert replay.status_code == 400
|
|
|
|
# PKCE is enforced: a fresh code with the wrong verifier fails
|
|
v2, ch2 = _pkce()
|
|
code2 = await _authorize_code(c, email=email, password="supersecret1", client_id=client_id,
|
|
redirect_uri=redirect_uri, challenge=ch2, resource=resource)
|
|
bad = await c.post("/v1/oauth/token", data={
|
|
"grant_type": "authorization_code", "code": code2, "redirect_uri": redirect_uri,
|
|
"client_id": client_id, "code_verifier": verifier, # wrong verifier
|
|
})
|
|
assert bad.status_code == 400
|