382 lines
15 KiB
Python
382 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
from collections.abc import AsyncIterator, Awaitable, Callable
|
|
from contextlib import asynccontextmanager
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from uuid import UUID
|
|
|
|
import yaml # type: ignore[import-untyped]
|
|
from fastapi import FastAPI, Request, Response
|
|
from sqlalchemy import select
|
|
|
|
from packages.hub_policy import Governance, PolicyEngine, PolicyRule, SlidingWindowRateLimiter
|
|
from packages.hub_runtime import read_secret_file
|
|
from services.control_api.auth import OIDCProviderClient, OIDCSettings, PrincipalMapper
|
|
from services.control_api.auth.errors import InvalidToken, UnknownMachinePrincipal
|
|
from services.control_api.auth.machine import MachineClientBinding
|
|
from services.control_api.contextforge import ContextForgeClient, ContextForgeSettings
|
|
from services.control_api.persistence import (
|
|
SqlAlchemyAliasSource,
|
|
create_engine_from_url,
|
|
session_factory,
|
|
)
|
|
from services.control_api.persistence.demo_client import DEMO_CLIENT_ID
|
|
from services.control_api.persistence.models import ClientRow, WorkspaceRow
|
|
from services.dataplane_router import (
|
|
AliasRegistry,
|
|
DataplaneProxy,
|
|
InvocationRecorder,
|
|
RouterSettings,
|
|
build_router,
|
|
)
|
|
from services.dataplane_router.models import InvocationMetadata
|
|
|
|
from .audit import SqlAlchemyAuditLog
|
|
from .common import (
|
|
hub_database_url,
|
|
identity_assertion,
|
|
load_connector_identity_key,
|
|
run_asgi,
|
|
)
|
|
from .demo_token import TOKEN_PREFIX, DemoTokenAuthenticator, load_demo_token_signing_key
|
|
from .mcp_skills import StaticSkills
|
|
from .oauth_metadata import OAuthMetadataSettings, bearer_challenge
|
|
|
|
|
|
class JsonLineInvocationRecorder(InvocationRecorder):
|
|
"""Emit bounded response-contract metadata without request data or credentials."""
|
|
|
|
def record(self, metadata: InvocationMetadata) -> None:
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"event": "mcp.response.completed",
|
|
"traceId": metadata.trace_id,
|
|
"project": metadata.project_slug,
|
|
"mcpMethod": metadata.mcp_method,
|
|
"mcpName": metadata.mcp_name,
|
|
"protocolRevision": metadata.protocol_revision,
|
|
"responseStatus": metadata.response_status,
|
|
"responseBytes": metadata.response_bytes,
|
|
"responseContentType": metadata.response_content_type,
|
|
"jsonrpcResponse": metadata.jsonrpc_response,
|
|
"responseIdMatches": metadata.response_id_matches,
|
|
"responseShape": metadata.response_shape,
|
|
"durationMs": metadata.duration_ms,
|
|
},
|
|
separators=(",", ":"),
|
|
sort_keys=True,
|
|
),
|
|
flush=True,
|
|
)
|
|
|
|
|
|
def _policy_engine(catalog_root: Path) -> PolicyEngine:
|
|
policies: list[PolicyRule] = []
|
|
for path in sorted((catalog_root / "policies").glob("*.yaml")):
|
|
document = yaml.safe_load(path.read_text(encoding="utf-8"))
|
|
if not isinstance(document, dict):
|
|
raise ValueError("policy catalog document is invalid")
|
|
policies.append(PolicyRule.from_document(document))
|
|
if not policies or not any(policy.slug == "global-deny" for policy in policies):
|
|
raise ValueError("policy catalog must include global-deny")
|
|
return PolicyEngine(policies)
|
|
|
|
|
|
def _tool_governance_metadata(catalog_root: Path) -> dict[str, tuple[str, str, str]]:
|
|
metadata: dict[str, tuple[str, str, str]] = {}
|
|
for path in sorted((catalog_root / "tools").glob("*.yaml")):
|
|
document = yaml.safe_load(path.read_text(encoding="utf-8"))
|
|
if not isinstance(document, dict):
|
|
raise ValueError("tool catalog document is invalid")
|
|
spec = document.get("spec")
|
|
if not isinstance(spec, dict):
|
|
raise ValueError("tool catalog spec is invalid")
|
|
name = spec.get("canonicalName")
|
|
connector = spec.get("connectorType")
|
|
risk_class = spec.get("riskClass")
|
|
if not isinstance(name, str) or not name:
|
|
raise ValueError("tool governance metadata is incomplete")
|
|
if not isinstance(connector, str) or not connector:
|
|
raise ValueError("tool governance metadata is incomplete")
|
|
if not isinstance(risk_class, str) or not risk_class:
|
|
raise ValueError("tool governance metadata is incomplete")
|
|
published_name = f"itworx-{connector}-{name.replace('.', '-')}"
|
|
configured_names = spec.get("exposedNames", [])
|
|
if not isinstance(configured_names, list) or not all(
|
|
isinstance(item, str) for item in configured_names
|
|
):
|
|
raise ValueError("tool governance exposed names are invalid")
|
|
for exposed_name in (name, published_name, *configured_names):
|
|
if exposed_name in metadata:
|
|
raise ValueError(f"duplicate tool governance metadata: {exposed_name}")
|
|
metadata[exposed_name] = (connector, risk_class, name)
|
|
if not metadata:
|
|
raise ValueError("tool governance metadata is empty")
|
|
return metadata
|
|
|
|
|
|
class DatabaseAliasSource:
|
|
def __init__(self, sessions: Any, gateway_url: str) -> None:
|
|
self.sessions = sessions
|
|
self.gateway_url = gateway_url
|
|
|
|
def aliases(self) -> tuple[Any, ...]:
|
|
with self.sessions() as session:
|
|
workspace_id = session.scalar(
|
|
select(WorkspaceRow.id).where(WorkspaceRow.slug == "itworx")
|
|
)
|
|
if workspace_id is None:
|
|
return ()
|
|
return SqlAlchemyAliasSource(
|
|
session,
|
|
workspace_id=workspace_id,
|
|
gateway_base_url=self.gateway_url,
|
|
).aliases()
|
|
|
|
|
|
class DatabaseMachineAuthenticator:
|
|
_bearer = re.compile(r"^Bearer ([A-Za-z0-9._~+/=-]{1,65536})$")
|
|
|
|
def __init__(
|
|
self,
|
|
providers: OIDCProviderClient | tuple[OIDCProviderClient, ...],
|
|
sessions: Any,
|
|
) -> None:
|
|
if not isinstance(providers, tuple):
|
|
providers = (providers,)
|
|
if not providers:
|
|
raise ValueError("at least one OIDC provider is required")
|
|
self.providers = providers
|
|
self.sessions = sessions
|
|
|
|
def __call__(self, authorization_header: str | None) -> Any:
|
|
if authorization_header is None:
|
|
raise InvalidToken()
|
|
match = self._bearer.fullmatch(authorization_header.strip())
|
|
if match is None:
|
|
raise InvalidToken()
|
|
claims = None
|
|
for provider in self.providers:
|
|
try:
|
|
claims = provider.validate_token(match.group(1))
|
|
break
|
|
except InvalidToken:
|
|
continue
|
|
if claims is None:
|
|
raise InvalidToken()
|
|
with self.sessions() as session:
|
|
client = session.scalar(
|
|
select(ClientRow).where(
|
|
ClientRow.issuer == claims.issuer,
|
|
ClientRow.subject == claims.subject,
|
|
ClientRow.status == "active",
|
|
)
|
|
)
|
|
if client is None:
|
|
raise UnknownMachinePrincipal()
|
|
client.last_seen_at = datetime.now(UTC)
|
|
session.commit()
|
|
binding = MachineClientBinding(
|
|
subject=client.subject,
|
|
client_id=UUID(str(client.id)),
|
|
enabled=True,
|
|
)
|
|
return PrincipalMapper(machine_subjects={binding.subject: binding}).map_claims(claims)
|
|
|
|
|
|
class CompositeAuthenticator:
|
|
"""Routes demo-token-shaped bearer credentials to the demo path, else OIDC."""
|
|
|
|
def __init__(
|
|
self, demo: Callable[[str | None], Any], oidc: Callable[[str | None], Any]
|
|
) -> None:
|
|
self.demo = demo
|
|
self.oidc = oidc
|
|
|
|
def __call__(self, authorization_header: str | None) -> Any:
|
|
if authorization_header is not None and authorization_header.startswith(
|
|
f"Bearer {TOKEN_PREFIX}"
|
|
):
|
|
return self.demo(authorization_header)
|
|
return self.oidc(authorization_header)
|
|
|
|
|
|
class DemoAwareRateLimiter(SlidingWindowRateLimiter):
|
|
"""Applies a stricter shared ceiling to the one fixed demo client identity.
|
|
|
|
Every other client_id behaves exactly as the base SlidingWindowRateLimiter
|
|
it subclasses — this only special-cases the demo client so Governance's
|
|
`rate_limiter: SlidingWindowRateLimiter | None` type still holds.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
default_limit: int,
|
|
demo_limit: int,
|
|
demo_client_id: UUID,
|
|
window_seconds: float = 60.0,
|
|
) -> None:
|
|
super().__init__(limit=default_limit, window_seconds=window_seconds)
|
|
self._demo_limiter = SlidingWindowRateLimiter(
|
|
limit=demo_limit, window_seconds=window_seconds
|
|
)
|
|
self._demo_client_id = demo_client_id
|
|
|
|
def allow(self, *, client_id: UUID | str, project: str, tool: str) -> bool:
|
|
if str(client_id) == str(self._demo_client_id):
|
|
return self._demo_limiter.allow(client_id=client_id, project=project, tool=tool)
|
|
return super().allow(client_id=client_id, project=project, tool=tool)
|
|
|
|
|
|
def _oidc_providers() -> tuple[OIDCProviderClient, ...]:
|
|
providers = [OIDCProviderClient(OIDCSettings.from_env())]
|
|
additional = {
|
|
"issuer": os.environ.get("HUB_MCP_ADDITIONAL_OIDC_ISSUER", "").strip(),
|
|
"audience": os.environ.get("HUB_MCP_ADDITIONAL_OIDC_AUDIENCE", "").strip(),
|
|
"client_id": os.environ.get("HUB_MCP_ADDITIONAL_OIDC_CLIENT_ID", "").strip(),
|
|
"redirect_uri": os.environ.get("HUB_MCP_ADDITIONAL_OIDC_REDIRECT_URI", "").strip(),
|
|
}
|
|
configured = [bool(value) for value in additional.values()]
|
|
if any(configured) and not all(configured):
|
|
raise ValueError("additional MCP OIDC provider configuration is incomplete")
|
|
if all(configured):
|
|
settings = OIDCSettings.from_env(
|
|
{
|
|
"HUB_OIDC_ISSUER": additional["issuer"],
|
|
"HUB_OIDC_AUDIENCE": additional["audience"],
|
|
"HUB_OIDC_CLIENT_ID": additional["client_id"],
|
|
"HUB_OIDC_REDIRECT_URI": additional["redirect_uri"],
|
|
"HUB_OIDC_ALLOWED_ALGORITHMS": os.environ.get(
|
|
"HUB_OIDC_ALLOWED_ALGORITHMS", "RS256"
|
|
),
|
|
}
|
|
)
|
|
providers.append(OIDCProviderClient(settings))
|
|
return tuple(providers)
|
|
|
|
|
|
def build_app() -> FastAPI:
|
|
oauth_metadata = OAuthMetadataSettings.from_environment()
|
|
engine = create_engine_from_url(hub_database_url())
|
|
sessions = session_factory(engine)
|
|
oidc_providers = _oidc_providers()
|
|
oidc_provider = oidc_providers[0]
|
|
contextforge = ContextForgeClient(ContextForgeSettings.from_env())
|
|
catalog_root = Path(os.environ.get("HUB_CATALOG_ROOT", "/opt/hub/catalog")).resolve()
|
|
rate_limit = int(os.environ.get("HUB_RATE_LIMIT_PER_MINUTE", "600"))
|
|
if not 1 <= rate_limit <= 10_000:
|
|
raise ValueError("HUB_RATE_LIMIT_PER_MINUTE must be between 1 and 10000")
|
|
demo_rate_limit = int(os.environ.get("DEMO_TOKEN_RATE_LIMIT_PER_MINUTE", "10"))
|
|
if not 1 <= demo_rate_limit <= rate_limit:
|
|
raise ValueError(
|
|
"DEMO_TOKEN_RATE_LIMIT_PER_MINUTE must be between 1 and HUB_RATE_LIMIT_PER_MINUTE"
|
|
)
|
|
governance = Governance(
|
|
_policy_engine(catalog_root),
|
|
SqlAlchemyAuditLog(sessions, issuer=str(oidc_provider.settings.issuer)),
|
|
rate_limiter=DemoAwareRateLimiter(
|
|
default_limit=rate_limit, demo_limit=demo_rate_limit, demo_client_id=DEMO_CLIENT_ID
|
|
),
|
|
)
|
|
connector_identity_key = load_connector_identity_key()
|
|
intake_key_path = os.environ.get("KNOWLEDGE_HUB_INGESTION_IDENTITY_KEY_FILE")
|
|
intake_identity_key: bytes | None = None
|
|
if intake_key_path:
|
|
intake_identity_key = read_secret_file(Path(intake_key_path)).encode("utf-8")
|
|
if len(intake_identity_key) < 32:
|
|
raise ValueError("knowledge intake identity key must contain at least 32 bytes")
|
|
if (
|
|
os.environ.get("INGESTION_CONNECTOR_ENABLED", "false").lower() == "true"
|
|
and not intake_identity_key
|
|
):
|
|
raise ValueError("knowledge intake identity key is required when intake is enabled")
|
|
alias_cache_ttl = int(os.environ.get("HUB_ALIAS_CACHE_TTL_SECONDS", "5"))
|
|
if not 1 <= alias_cache_ttl <= 60:
|
|
raise ValueError("HUB_ALIAS_CACHE_TTL_SECONDS must be between 1 and 60")
|
|
router_settings = RouterSettings(
|
|
alias_cache_ttl_seconds=alias_cache_ttl,
|
|
max_upstream_concurrency=int(os.environ.get("HUB_UPSTREAM_MAX_CONCURRENCY", "24")),
|
|
upstream_queue_timeout_seconds=float(
|
|
os.environ.get("HUB_UPSTREAM_QUEUE_TIMEOUT_SECONDS", "5")
|
|
),
|
|
)
|
|
# Use ContextForge's unique control-network alias. The generic alias is
|
|
# present on several Docker networks and can resolve outside the service
|
|
# token's source-IP allowlist, causing intermittent 403 responses.
|
|
source = DatabaseAliasSource(sessions, "http://contextforge-control:4444/mcp")
|
|
registry = AliasRegistry(source, settings=router_settings)
|
|
demo_revocations_path = Path(
|
|
os.environ.get("DEMO_TOKEN_REVOCATIONS_FILE", "/etc/itworx/demo-token-revocations.json")
|
|
)
|
|
authenticator = CompositeAuthenticator(
|
|
demo=DemoTokenAuthenticator(
|
|
None, demo_revocations_path, key_loader=load_demo_token_signing_key
|
|
),
|
|
oidc=DatabaseMachineAuthenticator(oidc_providers, sessions),
|
|
)
|
|
proxy = DataplaneProxy(
|
|
registry,
|
|
authenticator,
|
|
recorder=JsonLineInvocationRecorder(),
|
|
settings=router_settings,
|
|
governance=governance,
|
|
tool_metadata=_tool_governance_metadata(catalog_root),
|
|
upstream_authorization=contextforge.authorization_header,
|
|
identity_assertion=lambda project, environment, client_id, trace_id: identity_assertion(
|
|
(intake_identity_key or b"") if project == "knowledge-hub" else connector_identity_key,
|
|
project=project,
|
|
environment=environment,
|
|
client_id=client_id,
|
|
trace_id=trace_id,
|
|
),
|
|
mcp_extension=StaticSkills(serve_public_widgets=True),
|
|
)
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(_: FastAPI) -> AsyncIterator[None]:
|
|
try:
|
|
yield
|
|
finally:
|
|
await proxy.close()
|
|
for provider in oidc_providers:
|
|
provider.close()
|
|
contextforge.close()
|
|
engine.dispose()
|
|
|
|
app = FastAPI(
|
|
title="ITWorx MCP Hub dataplane",
|
|
docs_url=None,
|
|
redoc_url=None,
|
|
openapi_url=None,
|
|
lifespan=lifespan,
|
|
)
|
|
|
|
@app.middleware("http")
|
|
async def advertise_chatgpt_oauth(
|
|
request: Request, call_next: Callable[[Request], Awaitable[Response]]
|
|
) -> Response:
|
|
response = await call_next(request)
|
|
path_parts = request.url.path.removeprefix("/mcp/").split("/")
|
|
if (
|
|
len(path_parts) == 2
|
|
and all(part.replace("-", "").isalnum() for part in path_parts)
|
|
and response.status_code == 401
|
|
):
|
|
resource_path = "/".join(path_parts)
|
|
response.headers["WWW-Authenticate"] = bearer_challenge(resource_path, oauth_metadata)
|
|
return response
|
|
|
|
app.include_router(build_router(proxy))
|
|
return app
|
|
|
|
|
|
def run() -> None:
|
|
run_asgi(build_app())
|