From ea8601a73ad7c848efa4a8f6e2fe8e3abee661b4 Mon Sep 17 00:00:00 2001 From: kinzhi Date: Tue, 21 Jul 2026 14:57:09 +0800 Subject: [PATCH 1/2] feat(runtime): add automatic credentials --- README.md | 11 +- src/openlinker/runtime/credentials.py | 309 ++++++++++++++++++++++++++ src/openlinker/runtime/transport.py | 113 ++++++++-- src/openlinker/runtime/types.py | 6 +- src/openlinker/runtime/worker.py | 107 +++++++-- tests/test_runtime.py | 4 +- tests/test_runtime_credentials.py | 32 +++ 7 files changed, 533 insertions(+), 49 deletions(-) create mode 100644 src/openlinker/runtime/credentials.py create mode 100644 tests/test_runtime_credentials.py diff --git a/README.md b/README.md index 08fd81e..b543c82 100644 --- a/README.md +++ b/README.md @@ -235,15 +235,7 @@ async def handle(context: runtime.RuntimeContext) -> runtime.RuntimeResult: async def main() -> None: worker = runtime.RuntimeWorker( platform_url=os.environ["OPENLINKER_URL"], - runtime_url=os.environ.get("OPENLINKER_RUNTIME_URL", ""), - node_id=os.environ["OPENLINKER_NODE_ID"], - agent_id=os.environ["OPENLINKER_AGENT_ID"], agent_token=os.environ["OPENLINKER_AGENT_TOKEN"], - mtls=runtime.RuntimeMTLS( - cert_file=os.environ["OPENLINKER_RUNTIME_MTLS_CERT_FILE"], - key_file=os.environ["OPENLINKER_RUNTIME_MTLS_KEY_FILE"], - ca_file=os.environ["OPENLINKER_RUNTIME_MTLS_CA_FILE"], - ), data_dir=os.environ.get( "OPENLINKER_RUNTIME_DATA_DIR", "./.openlinker-runtime", @@ -261,6 +253,9 @@ A worker is async and single-use. Events and results are encrypted and fsynced before upload. Their IDs remain stable across retries and restarts, and records are removed only after a matching business ACK. The file store keeps a stable worker identity while rotating the Session identity for each process start. +The SDK also generates the Node private key in `data_dir`, enrolls it with the +Agent Token, and renews its 24-hour client certificate automatically. Explicit +`RuntimeMTLS` file paths are only needed for external-PKI compatibility. `MemoryRuntimeStore` is for explicit tests only and requires `allow_unsafe_memory_store=True`. Production workers should use `data_dir` or diff --git a/src/openlinker/runtime/credentials.py b/src/openlinker/runtime/credentials.py new file mode 100644 index 0000000..9dbe807 --- /dev/null +++ b/src/openlinker/runtime/credentials.py @@ -0,0 +1,309 @@ +from __future__ import annotations + +import asyncio +import hashlib +import hmac +import json +import os +import tempfile +import uuid +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +import httpx +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ec +from cryptography.x509.oid import NameOID + +from .types import ( + RUNTIME_CONTRACT_DIGEST, + RUNTIME_CONTRACT_ID, + RUNTIME_PROTOCOL_VERSION, + RUNTIME_REQUIRED_FEATURES, + RuntimeMTLS, +) + + +_CREDENTIAL_FILE = "runtime-credential.json" +_CERT_FILE = "runtime-client-chain.pem" +_KEY_FILE = "runtime-client-key.pem" +_CA_FILE = "runtime-ca.pem" +_FORMAT = 1 + + +@dataclass +class RuntimeCredentialIdentity: + node_id: str + agent_id: str + + +class RuntimeCredentialManager: + def __init__( + self, + *, + data_dir: Path, + credential_endpoint: str, + agent_token: str, + node_id: str, + agent_id: str, + node_version: str, + capacity: int, + logger: Any = None, + ) -> None: + self.data_dir = data_dir.resolve() + self.endpoint = _validate_endpoint(credential_endpoint) + self.agent_token = agent_token + self.configured_node_id = node_id + self.configured_agent_id = agent_id + self.node_version = node_version + self.capacity = capacity + self.logger = logger + self._disk: dict[str, Any] = {} + self._lock = asyncio.Lock() + self._task: asyncio.Task[None] | None = None + + async def open(self) -> None: + self.data_dir.mkdir(parents=True, exist_ok=True, mode=0o700) + os.chmod(self.data_dir, 0o700) + self._disk = await asyncio.to_thread(self._load_or_create) + + @property + def identity(self) -> RuntimeCredentialIdentity: + agent_id = str(self._disk.get("agent_id", "")) + if not agent_id: + raise RuntimeError("Runtime credential has not been enrolled") + return RuntimeCredentialIdentity(str(self._disk["node_id"]), agent_id) + + @property + def mtls(self) -> RuntimeMTLS: + if not self._disk.get("certificate_chain_pem"): + raise RuntimeError("Runtime mTLS credential is unavailable") + return RuntimeMTLS( + cert_file=str(self.data_dir / _CERT_FILE), + key_file=str(self.data_dir / _KEY_FILE), + ca_file=str(self.data_dir / _CA_FILE), + ) + + async def ensure(self, force: bool = False) -> None: + async with self._lock: + if not force and not self._needs_renewal(): + return + await self._issue() + + def start(self) -> None: + if self._task is None: + self._task = asyncio.create_task(self._renew_loop()) + + async def close(self) -> None: + if self._task is not None: + self._task.cancel() + await asyncio.gather(self._task, return_exceptions=True) + self._task = None + + async def _renew_loop(self) -> None: + while True: + renew_after = _parse_time(self._disk.get("renew_after")) + wait = max(1.0, (renew_after - datetime.now(timezone.utc)).total_seconds()) + await asyncio.sleep(wait) + try: + await self.ensure() + except Exception as exc: # pragma: no cover - logging branch + if self.logger is not None: + self.logger.warning("Runtime certificate renewal failed; retrying: %s", exc) + await asyncio.sleep(300) + + def _needs_renewal(self) -> bool: + if not self._disk.get("certificate_chain_pem"): + return True + now = datetime.now(timezone.utc) + return now + timedelta(minutes=5) >= _parse_time( + self._disk.get("not_after") + ) or now >= _parse_time(self._disk.get("renew_after")) + + async def _issue(self) -> None: + private_key = serialization.load_pem_private_key( + str(self._disk["private_key_pem"]).encode(), password=None + ) + csr = ( + x509.CertificateSigningRequestBuilder() + .subject_name( + x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "OpenLinker Runtime Node")]) + ) + .sign(private_key, hashes.SHA256()) + ) + payload = { + "node_id": self._disk["node_id"], + "display_name": f"runtime-{str(self._disk['node_id']).replace('-', '')[:12]}", + "node_version": self.node_version, + "protocol_version": RUNTIME_PROTOCOL_VERSION, + "runtime_contract_id": RUNTIME_CONTRACT_ID, + "runtime_contract_digest": RUNTIME_CONTRACT_DIGEST, + "features": list(RUNTIME_REQUIRED_FEATURES), + "capacity": self.capacity, + "csr_pem": csr.public_bytes(serialization.Encoding.PEM).decode(), + } + async with httpx.AsyncClient(timeout=15, follow_redirects=False, trust_env=False) as client: + response = await client.post( + self.endpoint, + json=payload, + headers={ + "Authorization": f"Bearer {self.agent_token}", + "Accept": "application/json", + "X-OpenLinker-SDK": "openlinker-python/runtime-worker", + }, + ) + if response.status_code != 200: + raise RuntimeError( + f"Runtime certificate request failed with HTTP {response.status_code}" + ) + if len(response.content) > 64 * 1024: + raise RuntimeError("Runtime certificate response exceeds 64 KiB") + issued = response.json() + self._validate_response(issued) + self._disk.update( + agent_id=issued["agent_id"], + certificate_chain_pem=issued["certificate_chain_pem"], + trust_bundle_pem=issued["trust_bundle_pem"], + certificate_serial=str(issued["certificate_serial"]).lower(), + not_before=issued["not_before"], + not_after=issued["not_after"], + renew_after=issued["renew_after"], + ) + await asyncio.to_thread(self._persist) + + def _validate_response(self, issued: Any) -> None: + if not isinstance(issued, dict): + raise RuntimeError("Runtime certificate response is invalid") + not_before = _parse_time(issued.get("not_before")) + not_after = _parse_time(issued.get("not_after")) + renew_after = _parse_time(issued.get("renew_after")) + if ( + issued.get("node_id") != self._disk["node_id"] + or not _uuid(str(issued.get("agent_id", ""))) + or issued.get("public_key_thumbprint") != self._disk["public_key_thumbprint"] + or "BEGIN CERTIFICATE" not in str(issued.get("certificate_chain_pem", "")) + or "BEGIN CERTIFICATE" not in str(issued.get("trust_bundle_pem", "")) + or not_after <= not_before + or not timedelta(hours=23, minutes=50) + <= not_after - not_before + <= timedelta(hours=24, minutes=10) + or not (not_before < renew_after < not_after) + ): + raise RuntimeError("Runtime certificate response is invalid") + + def _load_or_create(self) -> dict[str, Any]: + path = self.data_dir / _CREDENTIAL_FILE + if path.exists(): + stat = path.lstat() + if not path.is_file() or stat.st_mode & 0o077 or not 0 < stat.st_size <= 64 * 1024: + raise RuntimeError("Runtime credential file is corrupt or not private") + disk = json.loads(path.read_text()) + checksum = str(disk.pop("checksum", "")) + if disk.get("format") != _FORMAT or not _uuid(str(disk.get("node_id", ""))): + raise RuntimeError("Runtime credential file is corrupt") + if not hmac.compare_digest(checksum, _checksum(disk)): + raise RuntimeError("Runtime credential file checksum is invalid") + disk["checksum"] = checksum + if self.configured_node_id and disk["node_id"] != self.configured_node_id: + raise RuntimeError("configured node_id differs from the key bound to data_dir") + if self.configured_agent_id and disk.get("agent_id") not in ( + None, + "", + self.configured_agent_id, + ): + raise RuntimeError( + "configured agent_id differs from the credential bound to data_dir" + ) + return disk + node_id = self.configured_node_id or str(uuid.uuid4()) + if not _uuid(node_id): + raise ValueError("node_id must be a non-zero lowercase UUID") + key = ec.generate_private_key(ec.SECP256R1()) + private_pem = key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ).decode() + spki = key.public_key().public_bytes( + serialization.Encoding.DER, serialization.PublicFormat.SubjectPublicKeyInfo + ) + disk = { + "format": _FORMAT, + "node_id": node_id, + "private_key_pem": private_pem, + "public_key_thumbprint": hashlib.sha256(spki).hexdigest(), + "checksum": "", + } + self._disk = disk + self._persist() + return disk + + def _persist(self) -> None: + disk = dict(self._disk) + disk.pop("checksum", None) + disk["checksum"] = _checksum(disk) + self._disk["checksum"] = disk["checksum"] + _atomic_private( + self.data_dir / _CREDENTIAL_FILE, json.dumps(disk, separators=(",", ":")).encode() + ) + _atomic_private(self.data_dir / _KEY_FILE, str(disk["private_key_pem"]).encode()) + if disk.get("certificate_chain_pem"): + _atomic_private(self.data_dir / _CERT_FILE, str(disk["certificate_chain_pem"]).encode()) + _atomic_private(self.data_dir / _CA_FILE, str(disk["trust_bundle_pem"]).encode()) + + +def _atomic_private(path: Path, value: bytes) -> None: + fd, temporary = tempfile.mkstemp(prefix=".openlinker-", dir=path.parent) + try: + os.fchmod(fd, 0o600) + with os.fdopen(fd, "wb", closefd=True) as output: + output.write(value) + output.flush() + os.fsync(output.fileno()) + os.replace(temporary, path) + directory_fd = os.open(path.parent, os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + finally: + if os.path.exists(temporary): + os.unlink(temporary) + + +def _checksum(value: dict[str, Any]) -> str: + return hashlib.sha256( + json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + + +def _parse_time(value: Any) -> datetime: + if not isinstance(value, str) or not value: + return datetime.min.replace(tzinfo=timezone.utc) + return datetime.fromisoformat(value.replace("Z", "+00:00")).astimezone(timezone.utc) + + +def _validate_endpoint(value: str) -> str: + parsed = urlparse(value.strip()) + loopback = parsed.hostname in {"localhost", "127.0.0.1", "::1"} + if ( + not parsed.hostname + or parsed.username + or parsed.password + or parsed.fragment + or not (parsed.scheme == "https" or (parsed.scheme == "http" and loopback)) + ): + raise ValueError("Runtime credential endpoint must use HTTPS") + return parsed.geturl() + + +def _uuid(value: str) -> bool: + try: + parsed = uuid.UUID(value) + except ValueError: + return False + return parsed.int != 0 and str(parsed) == value diff --git a/src/openlinker/runtime/transport.py b/src/openlinker/runtime/transport.py index 21d4c95..f5b22b3 100644 --- a/src/openlinker/runtime/transport.py +++ b/src/openlinker/runtime/transport.py @@ -94,6 +94,9 @@ class RuntimeTransportPolicy: class RuntimeDiscoveryConnection: runtime_origin: str policy: RuntimeTransportPolicy + mtls_required: bool = True + credential_endpoint: str = "" + trust_bundle_endpoint: str = "" class RuntimeTransport(Protocol): @@ -188,14 +191,19 @@ def decode_runtime_discovery_manifest(manifest: dict[str, Any]) -> RuntimeDiscov runtime = manifest.get("runtime") if not isinstance(base_urls, dict) or not isinstance(runtime, dict): raise RuntimeProtocolError("OpenLinker does not provide Runtime discovery") - if runtime.get("enabled") is not True or runtime.get("mtls_required") is not True: - raise RuntimeProtocolError("OpenLinker Runtime is disabled or does not require mTLS") + if runtime.get("enabled") is not True or not isinstance(runtime.get("mtls_required"), bool): + raise RuntimeProtocolError("OpenLinker Runtime is disabled or has no transport policy") discovered = base_urls.get("runtime") if not isinstance(discovered, str) or not discovered: raise RuntimeProtocolError("OpenLinker does not provide a Runtime origin") return RuntimeDiscoveryConnection( - runtime_origin=validate_runtime_origin(discovered), + runtime_origin=validate_runtime_origin( + discovered, allow_loopback_http=runtime["mtls_required"] is False + ), policy=decode_runtime_transport_policy(runtime), + mtls_required=runtime["mtls_required"], + credential_endpoint=str(runtime.get("credential_endpoint", "")), + trust_bundle_endpoint=str(runtime.get("trust_bundle_endpoint", "")), ) @@ -312,8 +320,8 @@ def validate_platform_origin(value: str) -> str: return _validate_origin(value, runtime=False) -def validate_runtime_origin(value: str) -> str: - return _validate_origin(value, runtime=True) +def validate_runtime_origin(value: str, *, allow_loopback_http: bool = False) -> str: + return _validate_origin(value, runtime=not allow_loopback_http) def build_runtime_ssl_context(config: RuntimeMTLS) -> ssl.SSLContext: @@ -338,12 +346,19 @@ def __init__( agent_token: str, mtls: RuntimeMTLS, *, + mtls_required: bool = True, + credential_manager: Any = None, _client: httpx.AsyncClient | None = None, ) -> None: - self.runtime_origin = validate_runtime_origin(runtime_origin) + self.runtime_origin = validate_runtime_origin( + runtime_origin, allow_loopback_http=not mtls_required + ) if not agent_token or agent_token != agent_token.strip(): raise ValueError("Agent Token is required") self._agent_token = agent_token + self._mtls = mtls + self._mtls_required = mtls_required + self._credential_manager = credential_manager self._attachment_id = "" self._attachment_generation = 0 self._attachment_transition = asyncio.Lock() @@ -351,7 +366,9 @@ def __init__( if _client is not None: self._client = _client else: - ssl_context = build_runtime_ssl_context(mtls) + ssl_context: ssl.SSLContext | bool = ( + build_runtime_ssl_context(mtls) if mtls_required else True + ) if mtls.server_name: hostname = urlparse(self.runtime_origin).hostname or "" if mtls.server_name != hostname: @@ -582,9 +599,32 @@ async def _request( if raw_body is not None else (wire_json_bytes(body) if body is not None else None) ) - response = await self._client.request( - method, url, content=content, headers=request_headers, follow_redirects=False - ) + if self._credential_manager is not None: + await self._credential_manager.ensure(False) + try: + response = await self._client.request( + method, url, content=content, headers=request_headers, follow_redirects=False + ) + except httpx.TransportError as exc: + if self._credential_manager is None or not _credential_tls_failure(exc): + raise + await self._credential_manager.ensure(True) + if self._owns_client: + await self._client.aclose() + verify: ssl.SSLContext | bool = ( + build_runtime_ssl_context(self._credential_manager.mtls) + if self._mtls_required + else True + ) + self._client = httpx.AsyncClient( + verify=verify, + timeout=httpx.Timeout(35.0, connect=10.0, write=10.0, pool=5.0), + follow_redirects=False, + trust_env=False, + ) + response = await self._client.request( + method, url, content=content, headers=request_headers, follow_redirects=False + ) if attachment_generation is not None and ( attachment_generation != self._attachment_generation or attachment_id != self._attachment_id @@ -612,10 +652,17 @@ def __init__( agent_token: str, mtls: RuntimeMTLS, http_transport: HTTPRuntimeTransport, + *, + mtls_required: bool = True, + credential_manager: Any = None, ) -> None: - self.runtime_origin = validate_runtime_origin(runtime_origin) + self.runtime_origin = validate_runtime_origin( + runtime_origin, allow_loopback_http=not mtls_required + ) self._agent_token = agent_token self._mtls = mtls + self._mtls_required = mtls_required + self._credential_manager = credential_manager self._http = http_transport self._socket: Any = None self._hello: dict[str, Any] | None = None @@ -635,8 +682,15 @@ async def connect(self, hello: dict[str, Any], *, fallback_reason: str = "") -> raise RuntimeError("Runtime WebSocket is already connected") _validate_fallback_reason(fallback_reason) parsed = urlparse(self.runtime_origin) - uri = urlunparse(parsed._replace(scheme="wss", path=RUNTIME_WEBSOCKET_PATH)) - ssl_context = build_runtime_ssl_context(self._mtls) + uri = urlunparse( + parsed._replace( + scheme="wss" if self._mtls_required else "ws", + path=RUNTIME_WEBSOCKET_PATH, + ) + ) + if self._credential_manager is not None: + await self._credential_manager.ensure(False) + ssl_context = build_runtime_ssl_context(self._mtls) if self._mtls_required else None kwargs: dict[str, Any] = { "ssl": ssl_context, "max_size": RUNTIME_MAX_MESSAGE_BYTES, @@ -655,15 +709,31 @@ async def connect(self, hello: dict[str, Any], *, fallback_reason: str = "") -> } if fallback_reason: kwargs[header_name][RUNTIME_FALLBACK_REASON_HEADER] = fallback_reason - if self._mtls.server_name: + if self._mtls_required and self._mtls.server_name: kwargs["server_hostname"] = self._mtls.server_name try: self._socket = await websockets.connect(uri, **kwargs) except Exception as exc: - translated = _websocket_upgrade_error(exc) - if translated is not None: - raise translated from exc - raise + if ( + self._credential_manager is not None + and self._mtls_required + and _credential_tls_failure(exc) + ): + try: + await self._credential_manager.ensure(True) + self._mtls = self._credential_manager.mtls + kwargs["ssl"] = build_runtime_ssl_context(self._mtls) + self._socket = await websockets.connect(uri, **kwargs) + except Exception as retry_exc: + translated = _websocket_upgrade_error(retry_exc) + if translated is not None: + raise translated from retry_exc + raise + else: + translated = _websocket_upgrade_error(exc) + if translated is not None: + raise translated from exc + raise self._hello = dict(hello) hello_id = await self._send_envelope("runtime.hello", hello) try: @@ -1205,6 +1275,13 @@ def _websocket_upgrade_error(exc: BaseException) -> RuntimeRemoteError | None: return None +def _credential_tls_failure(exc: BaseException) -> bool: + message = f"{type(exc).__name__}: {exc}".lower() + return any( + marker in message for marker in ("tls", "ssl", "x509", "certificate", "unknown authority") + ) + + __all__ = [ "ClaimedAssignment", "HTTPRuntimeTransport", diff --git a/src/openlinker/runtime/types.py b/src/openlinker/runtime/types.py index 0808b64..8f889fe 100644 --- a/src/openlinker/runtime/types.py +++ b/src/openlinker/runtime/types.py @@ -108,9 +108,9 @@ def __init__(self, timeout: float, spool: RuntimeSpoolStatus) -> None: @dataclass(frozen=True) class RuntimeMTLS: - cert_file: str - key_file: str - ca_file: str + cert_file: str = "" + key_file: str = "" + ca_file: str = "" server_name: str = "" diff --git a/src/openlinker/runtime/worker.py b/src/openlinker/runtime/worker.py index 2c01b2f..072eee1 100644 --- a/src/openlinker/runtime/worker.py +++ b/src/openlinker/runtime/worker.py @@ -15,6 +15,8 @@ from websockets.exceptions import ConnectionClosed +from .credentials import RuntimeCredentialManager + from .store import ( ASSIGNMENT_ACK_SENT, ASSIGNMENT_CONFIRMED, @@ -199,10 +201,10 @@ def __init__( self, *, platform_url: str, - node_id: str, - agent_id: str, + node_id: str = "", + agent_id: str = "", agent_token: str, - mtls: RuntimeMTLS, + mtls: RuntimeMTLS | None = None, handler: RuntimeHandler | RuntimeHandlerCallable, data_dir: str | Path | None = None, store: RuntimeStore | None = None, @@ -229,7 +231,7 @@ def __init__( self.node_id = node_id self.agent_id = agent_id self.agent_token = agent_token.strip() - self.mtls = mtls + self.mtls = mtls or RuntimeMTLS() self.handler = handler self.data_dir = Path(data_dir).expanduser() if data_dir is not None else None self.store = store @@ -293,6 +295,8 @@ def __init__( self._policy_last_error: BaseException | None = None self._policy_terminal_error: RuntimePolicyRecoveryError | None = None self._attachment_reason = "explicit" + self._credential_manager: RuntimeCredentialManager | None = None + self._mtls_required = True async def run(self) -> None: if self._started: @@ -530,13 +534,51 @@ async def _start(self) -> None: async def _setup_transport(self) -> None: if self._transport is not None: return + explicit_mtls = bool(self.mtls.cert_file and self.mtls.key_file and self.mtls.ca_file) + connection = None origin = self.runtime_url - if not origin: + if not origin or not explicit_mtls: + if not self.platform_url: + raise ValueError("automatic Runtime credentials require platform_url") connection = await discover_runtime_connection(self.platform_url) - origin = connection.runtime_origin + if not origin: + origin = connection.runtime_origin self._apply_transport_policy(connection.policy) - self.runtime_url = validate_runtime_origin(origin) - self._http_transport = HTTPRuntimeTransport(self.runtime_url, self.agent_token, self.mtls) + self._mtls_required = connection.mtls_required if connection is not None else True + self.runtime_url = validate_runtime_origin( + origin, allow_loopback_http=not self._mtls_required + ) + if not explicit_mtls or not self._mtls_required: + if self.data_dir is None: + raise ValueError("automatic Runtime credentials require data_dir") + platform_origin = validate_platform_origin(self.platform_url) + endpoint = ( + connection.credential_endpoint if connection is not None else "" + ) or platform_origin + "/api/v1/runtime-credentials" + manager = RuntimeCredentialManager( + data_dir=self.data_dir, + credential_endpoint=endpoint, + agent_token=self.agent_token, + node_id=self.node_id, + agent_id=self.agent_id, + node_version=self.node_version, + capacity=self.capacity, + logger=self.logger, + ) + await manager.open() + await manager.ensure() + identity = manager.identity + self.node_id, self.agent_id = identity.node_id, identity.agent_id + self.mtls = manager.mtls + self._credential_manager = manager + manager.start() + self._http_transport = HTTPRuntimeTransport( + self.runtime_url, + self.agent_token, + self.mtls, + mtls_required=self._mtls_required, + credential_manager=self._credential_manager, + ) selected_reason = resolve_runtime_fallback_reason( self._configured_transport_mode, "policy_selected" ) @@ -560,6 +602,8 @@ async def _setup_transport(self) -> None: self.agent_token, self.mtls, self._http_transport, + mtls_required=self._mtls_required, + credential_manager=self._credential_manager, ) await websocket.connect(self._hello(), fallback_reason=selected_reason) self._transport = websocket @@ -606,6 +650,8 @@ async def _connect_websocket_with_retry( self.agent_token, self.mtls, self._http_transport, + mtls_required=self._mtls_required, + credential_manager=self._credential_manager, ) try: await websocket.connect(self._hello(), fallback_reason=fallback_reason) @@ -1222,9 +1268,7 @@ async def _handle_cancel( try: await stopping_ack except Exception as exc: - self.logger.warning( - "Runtime cancel stopping ACK was not confirmed: %s", exc - ) + self.logger.warning("Runtime cancel stopping ACK was not confirmed: %s", exc) if not handler_stopped: await self._ack_cancel_best_effort( @@ -1788,6 +1832,8 @@ async def _shutdown(self) -> None: transports = {id(item): item for item in (self._transport, self._http_transport) if item} for transport in transports.values(): await transport.close() + if self._credential_manager is not None: + await self._credential_manager.close() if self._store is not None: self._store.close() @@ -1930,8 +1976,17 @@ async def _recover_runtime_policy_once(self, *, resume_durable: bool) -> Runtime ) connection = await discover_runtime_connection(self.platform_url) self._apply_transport_policy(connection.policy) - runtime_url = validate_runtime_origin(connection.runtime_origin) - replacement_http = HTTPRuntimeTransport(runtime_url, self.agent_token, self.mtls) + self._mtls_required = connection.mtls_required + runtime_url = validate_runtime_origin( + connection.runtime_origin, allow_loopback_http=not self._mtls_required + ) + replacement_http = HTTPRuntimeTransport( + runtime_url, + self.agent_token, + self.mtls, + mtls_required=self._mtls_required, + credential_manager=self._credential_manager, + ) previous_transport = self._transport previous_http = self._http_transport replacement: RuntimeTransport = replacement_http @@ -1947,6 +2002,8 @@ async def _recover_runtime_policy_once(self, *, resume_durable: bool) -> Runtime self.agent_token, self.mtls, replacement_http, + mtls_required=self._mtls_required, + credential_manager=self._credential_manager, ) try: await websocket.connect(self._hello(), fallback_reason=reason) @@ -2161,6 +2218,8 @@ async def _probe_websocket_transport(self) -> RuntimeTransport: self.agent_token, self.mtls, self._http_transport, + mtls_required=self._mtls_required, + credential_manager=self._credential_manager, ) try: await websocket.connect(self._hello(), fallback_reason="recovery") @@ -2215,14 +2274,26 @@ def _auto_allows_pull_fallback(self) -> bool: return self._auto_prefers_websocket() and "pull" in self._transport_order[1:] def _validate_config(self) -> None: - _canonical_uuid(self.node_id, "node_id") - _canonical_uuid(self.agent_id, "agent_id") + if self.node_id: + _canonical_uuid(self.node_id, "node_id") + if self.agent_id: + _canonical_uuid(self.agent_id, "agent_id") if not self.agent_token or self.agent_token.startswith("ol_user_"): raise ValueError("Agent Token is required and must not be a User Token") - if not self.mtls.cert_file or not self.mtls.key_file or not self.mtls.ca_file: - raise ValueError("Runtime mTLS cert, key, and CA files are required") + mtls_fields = sum( + bool(value) for value in (self.mtls.cert_file, self.mtls.key_file, self.mtls.ca_file) + ) + if mtls_fields not in {0, 3}: + raise ValueError("Runtime mTLS files must be configured together") + if mtls_fields == 3 and (not self.node_id or not self.agent_id): + raise ValueError("node_id and agent_id are required with explicit mTLS files") + if mtls_fields == 0 and not self.platform_url: + raise ValueError("automatic Runtime credentials require platform_url") if self.runtime_url: - validate_runtime_origin(self.runtime_url) + # The discovery manifest remains authoritative. This preliminary + # validation permits only loopback HTTP until discovery confirms + # that the operator explicitly disabled mTLS. + validate_runtime_origin(self.runtime_url, allow_loopback_http=mtls_fields == 0) else: validate_platform_origin(self.platform_url) if self._configured_transport_mode not in {"auto", "ws", "pull"}: diff --git a/tests/test_runtime.py b/tests/test_runtime.py index cdc8132..9b29082 100644 --- a/tests/test_runtime.py +++ b/tests/test_runtime.py @@ -1737,7 +1737,7 @@ def test_memory_store_requires_an_explicit_unsafe_opt_in(): ) -def test_worker_rejects_user_token_and_missing_mtls(): +def test_worker_rejects_user_token_and_partial_explicit_mtls(): common = { "platform_url": "https://platform.example.test", "node_id": NODE_ID, @@ -1756,5 +1756,5 @@ def test_worker_rejects_user_token_and_missing_mtls(): runtime.RuntimeWorker( **common, agent_token="ol_agent_test", - mtls=runtime.RuntimeMTLS("", "", ""), + mtls=runtime.RuntimeMTLS("client.crt", "", ""), ) diff --git a/tests/test_runtime_credentials.py b/tests/test_runtime_credentials.py new file mode 100644 index 0000000..356dc30 --- /dev/null +++ b/tests/test_runtime_credentials.py @@ -0,0 +1,32 @@ +import json +import stat + +import pytest + +from openlinker.runtime.credentials import RuntimeCredentialManager + + +@pytest.mark.asyncio +async def test_runtime_credential_manager_generates_one_private_node_key(tmp_path): + options = { + "data_dir": tmp_path, + "credential_endpoint": "http://127.0.0.1:8080/api/v1/runtime-credentials", + "agent_token": "ol_agent_test", + "node_id": "", + "agent_id": "", + "node_version": "openlinker-python/runtime-worker", + "capacity": 1, + } + first = RuntimeCredentialManager(**options) + await first.open() + initial = json.loads((tmp_path / "runtime-credential.json").read_text()) + assert "BEGIN PRIVATE KEY" in initial["private_key_pem"] + assert stat.S_IMODE((tmp_path / "runtime-credential.json").stat().st_mode) == 0o600 + assert stat.S_IMODE((tmp_path / "runtime-client-key.pem").stat().st_mode) == 0o600 + + second = RuntimeCredentialManager(**options) + await second.open() + reopened = json.loads((tmp_path / "runtime-credential.json").read_text()) + assert reopened["node_id"] == initial["node_id"] + assert reopened["private_key_pem"] == initial["private_key_pem"] + assert reopened["public_key_thumbprint"] == initial["public_key_thumbprint"] From dda10cdd31584ce92a68d1e5409902de2aa500fc Mon Sep 17 00:00:00 2001 From: MacMini2 Date: Tue, 21 Jul 2026 16:34:46 +0800 Subject: [PATCH 2/2] fix(runtime): honor token-only transport --- src/openlinker/runtime/transport.py | 20 ++++++ src/openlinker/runtime/worker.py | 16 ++++- tests/test_runtime.py | 99 +++++++++++++++++++++++++++++ tests/test_runtime_transport.py | 13 ++++ 4 files changed, 146 insertions(+), 2 deletions(-) diff --git a/src/openlinker/runtime/transport.py b/src/openlinker/runtime/transport.py index f5b22b3..b398d95 100644 --- a/src/openlinker/runtime/transport.py +++ b/src/openlinker/runtime/transport.py @@ -35,6 +35,7 @@ _SDK_AGENT = "openlinker-python/runtime-worker" _ATTACHMENT_HEADER = "OpenLinker-Runtime-Attachment" RUNTIME_FALLBACK_REASON_HEADER = "OpenLinker-Runtime-Fallback-Reason" +RUNTIME_NODE_ID_HEADER = "OpenLinker-Runtime-Node" RUNTIME_FALLBACK_REASONS = frozenset( {"explicit", "websocket_unavailable", "policy_forced", "recovery"} ) @@ -346,6 +347,7 @@ def __init__( agent_token: str, mtls: RuntimeMTLS, *, + node_id: str, mtls_required: bool = True, credential_manager: Any = None, _client: httpx.AsyncClient | None = None, @@ -355,6 +357,7 @@ def __init__( ) if not agent_token or agent_token != agent_token.strip(): raise ValueError("Agent Token is required") + self.node_id = _runtime_node_id(node_id) self._agent_token = agent_token self._mtls = mtls self._mtls_required = mtls_required @@ -594,6 +597,9 @@ async def _request( attachment_generation = self._attachment_generation request_headers[_ATTACHMENT_HEADER] = attachment_id request_headers.update(headers or {}) + # Node identity is transport-owned; callers cannot select a different + # Node by supplying an operation-specific header. + request_headers[RUNTIME_NODE_ID_HEADER] = self.node_id content = ( raw_body if raw_body is not None @@ -664,6 +670,7 @@ def __init__( self._mtls_required = mtls_required self._credential_manager = credential_manager self._http = http_transport + self._node_id = http_transport.node_id self._socket: Any = None self._hello: dict[str, Any] | None = None self._ready: RuntimeReady | None = None @@ -706,6 +713,7 @@ async def connect(self, hello: dict[str, Any], *, fallback_reason: str = "") -> kwargs[header_name] = { "Authorization": f"Bearer {self._agent_token}", "X-OpenLinker-SDK": _SDK_AGENT, + RUNTIME_NODE_ID_HEADER: self._node_id, } if fallback_reason: kwargs[header_name][RUNTIME_FALLBACK_REASON_HEADER] = fallback_reason @@ -1247,6 +1255,18 @@ def _runtime_session_id(value: Any) -> str: return value +def _runtime_node_id(value: Any) -> str: + if not isinstance(value, str): + raise ValueError("Runtime Node identity must be a lowercase non-zero UUID") + try: + parsed = uuid.UUID(value) + except ValueError as exc: + raise ValueError("Runtime Node identity must be a lowercase non-zero UUID") from exc + if parsed.int == 0 or str(parsed) != value: + raise ValueError("Runtime Node identity must be a lowercase non-zero UUID") + return value + + def _websocket_upgrade_error(exc: BaseException) -> RuntimeRemoteError | None: response = getattr(exc, "response", None) status_code = getattr(response, "status_code", None) diff --git a/src/openlinker/runtime/worker.py b/src/openlinker/runtime/worker.py index 072eee1..ddf9b03 100644 --- a/src/openlinker/runtime/worker.py +++ b/src/openlinker/runtime/worker.py @@ -548,7 +548,12 @@ async def _setup_transport(self) -> None: self.runtime_url = validate_runtime_origin( origin, allow_loopback_http=not self._mtls_required ) - if not explicit_mtls or not self._mtls_required: + if not self._mtls_required: + if not self.node_id or not self.agent_id: + raise ValueError( + "node_id and agent_id are required for token-only Runtime transport" + ) + elif not explicit_mtls: if self.data_dir is None: raise ValueError("automatic Runtime credentials require data_dir") platform_origin = validate_platform_origin(self.platform_url) @@ -576,6 +581,7 @@ async def _setup_transport(self) -> None: self.runtime_url, self.agent_token, self.mtls, + node_id=self.node_id, mtls_required=self._mtls_required, credential_manager=self._credential_manager, ) @@ -1975,8 +1981,13 @@ async def _recover_runtime_policy_once(self, *, resume_durable: bool) -> Runtime ) ) connection = await discover_runtime_connection(self.platform_url) + if connection.mtls_required != self._mtls_required: + raise RuntimePolicyRecoveryError( + RuntimeError( + "Runtime mTLS requirement changed; restart the Worker to apply the new security mode" + ) + ) self._apply_transport_policy(connection.policy) - self._mtls_required = connection.mtls_required runtime_url = validate_runtime_origin( connection.runtime_origin, allow_loopback_http=not self._mtls_required ) @@ -1984,6 +1995,7 @@ async def _recover_runtime_policy_once(self, *, resume_durable: bool) -> Runtime runtime_url, self.agent_token, self.mtls, + node_id=self.node_id, mtls_required=self._mtls_required, credential_manager=self._credential_manager, ) diff --git a/tests/test_runtime.py b/tests/test_runtime.py index 9b29082..3a4a89c 100644 --- a/tests/test_runtime.py +++ b/tests/test_runtime.py @@ -347,6 +347,70 @@ def make_policy_worker( ) +@pytest.mark.asyncio +async def test_worker_token_only_mode_skips_runtime_credentials(monkeypatch): + captured: dict[str, Any] = {} + transport = FakeTransport() + + async def discover(_platform_url): + return RuntimeDiscoveryConnection( + "https://runtime.example.test", + RuntimeTransportPolicy(("pull",), "pull"), + mtls_required=False, + credential_endpoint="not-a-valid-credential-endpoint", + ) + + def reject_credentials(*_args, **_kwargs): + raise AssertionError("token-only startup opened Runtime mTLS credentials") + + def build_http(*_args, **kwargs): + captured.update(kwargs) + return transport + + monkeypatch.setattr(runtime_worker_module, "discover_runtime_connection", discover) + monkeypatch.setattr(runtime_worker_module, "RuntimeCredentialManager", reject_credentials) + monkeypatch.setattr(runtime_worker_module, "HTTPRuntimeTransport", build_http) + worker = runtime.RuntimeWorker( + platform_url="https://platform.example.test", + node_id=NODE_ID, + agent_id=AGENT_ID, + agent_token="ol_agent_token_only", + store=runtime.MemoryRuntimeStore(), + allow_unsafe_memory_store=True, + handler=lambda _context: {}, + transport="pull", + ) + + await worker._setup_transport() + assert worker._credential_manager is None + assert captured["node_id"] == NODE_ID + assert captured["mtls_required"] is False + assert captured["credential_manager"] is None + + +@pytest.mark.asyncio +async def test_worker_token_only_mode_requires_configured_identity(monkeypatch): + async def discover(_platform_url): + return RuntimeDiscoveryConnection( + "https://runtime.example.test", + RuntimeTransportPolicy(("pull",), "pull"), + mtls_required=False, + ) + + monkeypatch.setattr(runtime_worker_module, "discover_runtime_connection", discover) + worker = runtime.RuntimeWorker( + platform_url="https://platform.example.test", + agent_token="ol_agent_token_only", + store=runtime.MemoryRuntimeStore(), + allow_unsafe_memory_store=True, + handler=lambda _context: {}, + transport="pull", + ) + + with pytest.raises(ValueError, match="required for token-only"): + await worker._setup_transport() + + @pytest.mark.asyncio async def test_assignment_is_durable_and_confirmed_before_handler_runs(): store = runtime.MemoryRuntimeStore() @@ -1306,6 +1370,41 @@ def build_http(*_args, **_kwargs): assert factory_calls == 1 +@pytest.mark.asyncio +async def test_worker_policy_recovery_rejects_mtls_requirement_change(monkeypatch): + discovery_calls = 0 + factory_calls = 0 + + class PolicyTransport(FakeTransport): + async def claim_assignment(self, _wait, _request): + raise runtime.RuntimeRemoteError( + "FORBIDDEN", "RUNTIME_POLICY_CHANGED", status_code=403 + ) + + async def discover(_platform_url): + nonlocal discovery_calls + discovery_calls += 1 + return RuntimeDiscoveryConnection( + f"https://runtime-{discovery_calls}.example.test", + RuntimeTransportPolicy(("pull",), "pull"), + mtls_required=discovery_calls == 1, + ) + + def build_http(*_args, **_kwargs): + nonlocal factory_calls + factory_calls += 1 + return PolicyTransport() + + monkeypatch.setattr(runtime_worker_module, "discover_runtime_connection", discover) + monkeypatch.setattr(runtime_worker_module, "HTTPRuntimeTransport", build_http) + worker = make_policy_worker(runtime.MemoryRuntimeStore(), mode="pull") + + with pytest.raises(RuntimeError, match="mTLS requirement changed"): + await asyncio.wait_for(worker.run(), timeout=1) + assert discovery_calls == 2 + assert factory_calls == 1 + + @pytest.mark.asyncio async def test_worker_rediscovers_once_on_established_websocket_policy_close(monkeypatch): discovery_calls = 0 diff --git a/tests/test_runtime_transport.py b/tests/test_runtime_transport.py index 0896c45..d4b34f7 100644 --- a/tests/test_runtime_transport.py +++ b/tests/test_runtime_transport.py @@ -32,6 +32,7 @@ ATTACHMENT_ID = "88888888-8888-4888-8888-888888888888" NEXT_ATTACHMENT_ID = "99999999-9999-4999-8999-999999999999" RUNTIME_SESSION_ID = "33333333-3333-4333-8333-333333333333" +NODE_ID = "11111111-1111-4111-8111-111111111111" def test_runtime_discovery_policy_fixtures_are_language_consistent(): @@ -199,6 +200,7 @@ async def handler(request: httpx.Request) -> httpx.Response: "https://runtime.example.test", "ol_agent_secret", runtime.RuntimeMTLS("client.crt", "client.key", "ca.crt"), + node_id=NODE_ID, _client=client, ) await transport.create_session( @@ -215,6 +217,7 @@ async def handler(request: httpx.Request) -> httpx.Response: assert seen[0].url.path == "/api/v1/agent-runtime/sessions" assert "/v2/" not in seen[0].url.path.lower() assert seen[0].headers["Authorization"] == "Bearer ol_agent_secret" + assert seen[0].headers["OpenLinker-Runtime-Node"] == NODE_ID assert seen[0].headers["OpenLinker-Runtime-Fallback-Reason"] == "explicit" assert "OpenLinker-Runtime-Attachment" not in seen[0].headers assert seen[1].headers["OpenLinker-Runtime-Attachment"] == ATTACHMENT_ID @@ -255,6 +258,7 @@ async def handler(request: httpx.Request) -> httpx.Response: "https://runtime.example.test", "ol_agent_secret", runtime.RuntimeMTLS("client.crt", "client.key", "ca.crt"), + node_id=NODE_ID, _client=client, ) await transport.create_session({"runtime_session_id": RUNTIME_SESSION_ID}) @@ -322,6 +326,7 @@ async def close(self) -> None: "https://runtime.example.test", "ol_agent_secret", runtime.RuntimeMTLS("client.crt", "client.key", "ca.crt"), + node_id=NODE_ID, _client=client, ) websocket = WebSocketRuntimeTransport( @@ -364,6 +369,7 @@ async def handler(request: httpx.Request) -> httpx.Response: "https://runtime.example.test", "ol_agent_secret", runtime.RuntimeMTLS("client.crt", "client.key", "ca.crt"), + node_id=NODE_ID, _client=client, ) with pytest.raises(runtime.RuntimeProtocolError, match="redirect"): @@ -388,6 +394,7 @@ async def handler(request: httpx.Request) -> httpx.Response: "https://runtime.example.test", "ol_agent_secret", runtime.RuntimeMTLS("client.crt", "client.key", "ca.crt"), + node_id=NODE_ID, _client=client, ) await transport.create_session({"runtime_session_id": "session"}) @@ -420,6 +427,7 @@ async def handler(request: httpx.Request) -> httpx.Response: "https://runtime.example.test", "ol_agent_secret", runtime.RuntimeMTLS("client.crt", "client.key", "ca.crt"), + node_id=NODE_ID, _client=client, ) await transport.create_session({"runtime_session_id": "session"}) @@ -459,6 +467,7 @@ async def handler(request: httpx.Request) -> httpx.Response: "https://runtime.example.test", "ol_agent_secret", runtime.RuntimeMTLS("client.crt", "client.key", "ca.crt"), + node_id=NODE_ID, _client=client, ) await transport.create_session({"runtime_session_id": "session"}) @@ -496,6 +505,7 @@ async def test_websocket_heartbeat_allows_capacity_change_but_not_identity_chang "https://runtime.example.test", "ol_agent_secret", runtime.RuntimeMTLS("client.crt", "client.key", "ca.crt"), + node_id=NODE_ID, _client=client, ) websocket = WebSocketRuntimeTransport( @@ -540,6 +550,7 @@ async def fake_connect(uri: str, *, additional_headers=None, **kwargs): "https://runtime.example.test", "ol_agent_secret", runtime.RuntimeMTLS("client.crt", "client.key", "ca.crt"), + node_id=NODE_ID, _client=client, ) websocket = WebSocketRuntimeTransport( @@ -594,6 +605,7 @@ async def reject_upgrade(_uri: str, *, additional_headers=None, **_kwargs): "https://runtime.example.test", "ol_agent_secret", runtime.RuntimeMTLS("client.crt", "client.key", "ca.crt"), + node_id=NODE_ID, _client=client, ) websocket = WebSocketRuntimeTransport( @@ -631,6 +643,7 @@ async def handler(request: httpx.Request) -> httpx.Response: "https://runtime.example.test", "ol_agent_secret", runtime.RuntimeMTLS("client.crt", "client.key", "ca.crt"), + node_id=NODE_ID, _client=client, ) response = await transport.call_agent(