143 lines
5.4 KiB
Python
143 lines
5.4 KiB
Python
"""Private Unix-socket caller projection and rehydration client.
|
|
|
|
The broker response body is deliberately public metadata only. The short-lived
|
|
credential projection is delivered in a private response header, consumed
|
|
inside the model API process, and never returned by an HTTP API or persisted.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import hashlib
|
|
import http.client
|
|
import json
|
|
from pathlib import Path
|
|
import socket
|
|
from typing import Any, Mapping
|
|
|
|
CALLER = "chatgpt-governed-implementer-singlehost"
|
|
SCOPE = "release:candidate"
|
|
GOVERNED_TASK_SCOPE = "governed-task:project"
|
|
PROJECTED_SCOPES = frozenset({GOVERNED_TASK_SCOPE, SCOPE})
|
|
CREDENTIAL_HEADER = "X-AIOS-Caller-Projection-Credential"
|
|
|
|
|
|
class CallerProjectionError(RuntimeError):
|
|
def __init__(self, code: str, status: int = 503) -> None:
|
|
super().__init__(code)
|
|
self.code = code
|
|
self.status = status
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ProjectionResult:
|
|
public: dict[str, Any]
|
|
credential_loaded: bool
|
|
|
|
|
|
class _UnixHTTPConnection(http.client.HTTPConnection):
|
|
def __init__(self, socket_path: Path, timeout: float) -> None:
|
|
super().__init__("localhost", timeout=timeout)
|
|
self.socket_path = socket_path
|
|
|
|
def connect(self) -> None:
|
|
sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
sock.settimeout(self.timeout)
|
|
try:
|
|
sock.connect(str(self.socket_path))
|
|
except Exception:
|
|
sock.close()
|
|
raise
|
|
self.sock = sock
|
|
|
|
|
|
class CallerProjectionClient:
|
|
def __init__(self, socket_path: Path, *, timeout: float = 2.0) -> None:
|
|
self.socket_path = Path(socket_path)
|
|
self.timeout = timeout
|
|
|
|
@staticmethod
|
|
def _idempotency_key(caller: str, scope: str) -> str:
|
|
digest = hashlib.sha256((caller + "\0" + scope).encode()).hexdigest()
|
|
return "caller-projection-" + digest[:32] + "-v1"
|
|
|
|
def rehydrate(
|
|
self,
|
|
caller: str,
|
|
scope: str,
|
|
request: Mapping[str, Any] | None = None,
|
|
) -> ProjectionResult:
|
|
if caller != CALLER:
|
|
raise CallerProjectionError("CALLER_MISMATCH", 403)
|
|
if scope not in PROJECTED_SCOPES:
|
|
raise CallerProjectionError("SCOPE_MISMATCH", 403)
|
|
body = {
|
|
"caller": caller,
|
|
"scope": scope,
|
|
"idempotency_key": self._idempotency_key(caller, scope),
|
|
"request": dict(request or {}),
|
|
}
|
|
connection = _UnixHTTPConnection(self.socket_path, self.timeout)
|
|
try:
|
|
connection.request(
|
|
"POST",
|
|
"/v1/project-caller",
|
|
body=json.dumps(body, sort_keys=True, separators=(",", ":")).encode(),
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"Accept": "application/json",
|
|
"X-AIOS-Account": caller,
|
|
},
|
|
)
|
|
response = connection.getresponse()
|
|
status = response.status
|
|
raw = response.read()
|
|
credential = response.getheader(CREDENTIAL_HEADER, "")
|
|
except (OSError, TimeoutError, http.client.HTTPException) as exc:
|
|
raise CallerProjectionError("PROJECTION_BROKER_UNAVAILABLE", 503) from exc
|
|
finally:
|
|
connection.close()
|
|
|
|
try:
|
|
value = json.loads(raw.decode("utf-8"))
|
|
except Exception as exc:
|
|
raise CallerProjectionError("PROJECTION_RESPONSE_INVALID", 503) from exc
|
|
if not isinstance(value, Mapping):
|
|
raise CallerProjectionError("PROJECTION_RESPONSE_INVALID", 503)
|
|
if status != 200:
|
|
code = str(value.get("error") or "PROJECTION_DENIED")
|
|
raise CallerProjectionError(code, 403 if status < 500 else 503)
|
|
if any(
|
|
str(key).lower() in {
|
|
"credential",
|
|
"credential_material",
|
|
"token",
|
|
"bearer",
|
|
"secret",
|
|
"password",
|
|
}
|
|
for key in value
|
|
):
|
|
raise CallerProjectionError("PROJECTION_SECRET_FIELD_REJECTED", 503)
|
|
if value.get("caller") != CALLER or value.get("scope") != scope:
|
|
raise CallerProjectionError("PROJECTION_BINDING_INVALID", 503)
|
|
if value.get("projection_state") != "ACKED":
|
|
raise CallerProjectionError("PROJECTION_STATE_INVALID", 503)
|
|
if value.get("credential_state") != "REHYDRATED":
|
|
raise CallerProjectionError("CREDENTIAL_REHYDRATION_INVALID", 503)
|
|
if value.get("rehydration_state") not in {"GENERATED", "RECOVERED"}:
|
|
raise CallerProjectionError("REHYDRATION_STATE_INVALID", 503)
|
|
effective_scopes = value.get("effective_scopes")
|
|
if (
|
|
not isinstance(effective_scopes, list)
|
|
or len(effective_scopes) != len(set(effective_scopes))
|
|
or not all(isinstance(item, str) for item in effective_scopes)
|
|
or scope not in effective_scopes
|
|
or not set(effective_scopes).issubset(PROJECTED_SCOPES)
|
|
):
|
|
raise CallerProjectionError("PROJECTION_SCOPE_INVALID", 503)
|
|
if value.get("secret_exposed") is not False:
|
|
raise CallerProjectionError("PROJECTION_SECRET_EXPOSURE", 503)
|
|
if not isinstance(credential, str) or len(credential) < 32:
|
|
raise CallerProjectionError("CREDENTIAL_REHYDRATION_MISSING", 503)
|
|
return ProjectionResult(public=dict(value), credential_loaded=True)
|