Files
aios/active-runtime/app/aios_baseline/transport.py
T

514 lines
21 KiB
Python

"""Authenticated loopback HTTP transport for the Baseline 1.1 service."""
from __future__ import annotations
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import ipaddress
import json
import math
import re
import time
import threading
from typing import Any, Mapping
from urllib.error import HTTPError, URLError
from urllib.parse import quote, unquote, urlsplit
from urllib.request import HTTPRedirectHandler, ProxyHandler, Request, build_opener
from .service import Identity, ServiceError
MAX_REQUEST_BYTES = 65_536
MAX_RESPONSE_BYTES = 1_048_576
MAX_DEADLINE_SECONDS = 30
_IDENTIFIER = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.:-]{0,127}")
_IDENTITY_FIELDS = frozenset({"tenant_id", "subject_id", "role", "worker_id", "identity"})
class TransportError(RuntimeError):
def __init__(self, code: str, status: int | None = None):
self.code = code
self.status = status
super().__init__(code)
class _RequestError(RuntimeError):
def __init__(self, code: str, status: int):
self.code = code
self.status = status
super().__init__(code)
def _safe_identifier(value: Any, code: str = "HTTP_PATH_INVALID") -> str:
if not isinstance(value, str) or _IDENTIFIER.fullmatch(value) is None:
raise _RequestError(code, 400)
return value
def _has_identity_field(value: Any) -> bool:
if isinstance(value, Mapping):
if any(key in _IDENTITY_FIELDS for key in value):
return True
return any(_has_identity_field(item) for item in value.values())
if isinstance(value, list):
return any(_has_identity_field(item) for item in value)
return False
def _loopback_host(host: str, code: str) -> str:
if not isinstance(host, str):
raise TransportError(code)
try:
address = ipaddress.ip_address(host)
except ValueError as exc:
raise TransportError(code) from exc
if not address.is_loopback or address.version != 4:
raise TransportError(code)
return str(address)
class _BaselineHTTPServer(ThreadingHTTPServer):
daemon_threads = True
allow_reuse_address = True
def __init__(self, address, handler, service, identities):
self.baseline_service = service
self.trusted_identities = identities
self.request_slots=threading.BoundedSemaphore(16)
super().__init__(address, handler)
def get_request(self):
connection, address = super().get_request()
connection.settimeout(MAX_DEADLINE_SECONDS)
return connection, address
def process_request(self,request,client_address):
if not self.request_slots.acquire(blocking=False):
try:
body=b'{"error":{"code":"HTTP_CAPACITY_EXCEEDED"}}'
request.sendall(b'HTTP/1.1 503 Service Unavailable\r\nContent-Type: application/json\r\nConnection: close\r\nContent-Length: '+str(len(body)).encode()+b'\r\n\r\n'+body)
except OSError: pass
self.shutdown_request(request)
return
try: super().process_request(request,client_address)
except Exception:
self.request_slots.release()
raise
def process_request_thread(self,request,client_address):
try: super().process_request_thread(request,client_address)
finally: self.request_slots.release()
class _Handler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
server_version = "AIOSBaselineHTTP/1"
sys_version = ""
def log_message(self, _format, *_args):
# Request paths can contain tenant-owned identifiers; do not emit them or
# bearer-related failure details through the default stderr logger.
return
def send_error(self, code, message=None, explain=None):
# BaseHTTPRequestHandler otherwise emits an HTML page which may include
# parser details. Keep pre-dispatch protocol failures in the same closed
# JSON error format.
status = code if isinstance(code, int) and 400 <= code <= 599 else 400
self._send(status, {"error": {"code": "HTTP_PROTOCOL_ERROR"}})
def _send(self, status: int, value: Any) -> None:
try:
body = json.dumps(
value, sort_keys=True, separators=(",", ":"), ensure_ascii=False,
allow_nan=False,
).encode("utf-8")
except (TypeError, ValueError):
status = 500
body = b'{"error":{"code":"HTTP_RESPONSE_INVALID"}}'
if len(body) > MAX_RESPONSE_BYTES:
status = 500
body = b'{"error":{"code":"HTTP_RESPONSE_TOO_LARGE"}}'
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.send_header("Cache-Control", "no-store")
self.send_header("X-Content-Type-Options", "nosniff")
self.send_header("Connection", "close")
self.end_headers()
try:
self.wfile.write(body)
except OSError:
# A peer can disconnect or let its deadline expire while a response
# is being written. Do not let the server print a handler traceback.
pass
self.close_connection = True
def _error(self, error: _RequestError) -> None:
self._send(error.status, {"error": {"code": error.code}})
def _deadline(self) -> float:
values = self.headers.get_all("X-AIOS-Deadline", [])
if len(values) != 1:
raise _RequestError("HTTP_DEADLINE_REQUIRED", 400)
try:
deadline = float(values[0])
except (TypeError, ValueError) as exc:
raise _RequestError("HTTP_DEADLINE_INVALID", 400) from exc
now = time.time()
if not math.isfinite(deadline) or deadline > now + MAX_DEADLINE_SECONDS:
raise _RequestError("HTTP_DEADLINE_INVALID", 400)
if deadline <= now:
raise _RequestError("HTTP_DEADLINE_EXPIRED", 408)
self.connection.settimeout(deadline - now)
return deadline
def _identity(self) -> Identity:
values = self.headers.get_all("Authorization", [])
if len(values) != 1 or not values[0].startswith("Bearer "):
raise _RequestError("HTTP_UNAUTHORIZED", 401)
token = values[0][7:]
if not token or any(ord(char) < 33 or ord(char) > 126 for char in token):
raise _RequestError("HTTP_UNAUTHORIZED", 401)
identity = self.server.trusted_identities.get(token)
if identity is None:
raise _RequestError("HTTP_UNAUTHORIZED", 401)
return identity
def _json_body(self) -> dict:
if self.headers.get("Transfer-Encoding") is not None:
raise _RequestError("HTTP_TRANSFER_ENCODING_FORBIDDEN", 400)
content_types = self.headers.get_all("Content-Type", [])
if content_types != ["application/json"]:
raise _RequestError("HTTP_CONTENT_TYPE_INVALID", 415)
lengths = self.headers.get_all("Content-Length", [])
if len(lengths) != 1 or not lengths[0].isdigit():
raise _RequestError("HTTP_CONTENT_LENGTH_INVALID", 411)
length = int(lengths[0])
if length > MAX_REQUEST_BYTES:
raise _RequestError("HTTP_BODY_TOO_LARGE", 413)
if length <= 0:
raise _RequestError("HTTP_BODY_REQUIRED", 400)
raw = self.rfile.read(length)
if len(raw) != length:
raise _RequestError("HTTP_BODY_TRUNCATED", 400)
try:
def unique_object(pairs):
result = {}
for key, item in pairs:
if key in result:
raise _RequestError("HTTP_JSON_DUPLICATE_KEY", 400)
result[key] = item
return result
def reject_constant(_value):
raise _RequestError("HTTP_JSON_NON_FINITE", 400)
value = json.loads(
raw.decode("utf-8"), object_pairs_hook=unique_object,
parse_constant=reject_constant,
)
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
raise _RequestError("HTTP_JSON_INVALID", 400) from exc
if type(value) is not dict:
raise _RequestError("HTTP_JSON_OBJECT_REQUIRED", 400)
if _has_identity_field(value):
raise _RequestError("HTTP_IDENTITY_FIELD_FORBIDDEN", 400)
return value
def _path(self) -> list[str]:
parsed = urlsplit(self.path)
if parsed.query or parsed.fragment or not parsed.path.startswith("/"):
raise _RequestError("HTTP_PATH_INVALID", 400)
raw_segments = parsed.path.split("/")[1:]
if any(not segment for segment in raw_segments):
raise _RequestError("HTTP_NOT_FOUND", 404)
segments = [unquote(segment) for segment in raw_segments]
if any("/" in segment or "\\" in segment for segment in segments):
raise _RequestError("HTTP_PATH_INVALID", 400)
return segments
@staticmethod
def _exact(body: Mapping[str, Any], fields: set[str]) -> None:
if set(body) != fields:
raise _RequestError("HTTP_BODY_SCHEMA_INVALID", 400)
def _dispatch(self, method: str, identity: Identity, segments: list[str], body: dict | None):
service = self.server.baseline_service
if method == "POST" and segments == ["v1", "goals", "submit"]:
self._exact(body, {"request", "idempotency_key"})
if type(body["request"]) is not dict:
raise _RequestError("HTTP_BODY_SCHEMA_INVALID", 400)
return 201, service.submit_goal(
identity, body["request"], idempotency_key=body["idempotency_key"]
)
if method == "GET" and len(segments) == 3 and segments[:2] == ["v1", "tasks"]:
return 200, service.get_task(identity, _safe_identifier(segments[2]))
if (
method == "POST" and len(segments) == 4
and segments[:2] == ["v1", "tasks"] and segments[3] == "cancel"
):
self._exact(body, set())
return 200, service.cancel_task(identity, _safe_identifier(segments[2]))
if (
method == "GET" and len(segments) == 4
and segments[:2] == ["v1", "missions"] and segments[3] == "explain"
):
return 200, service.explain(identity, _safe_identifier(segments[2]))
if method == "POST" and segments == ["v1", "workers", "register"]:
self._exact(body, set())
return 200, service.register_worker(identity)
if method == "POST" and segments == ["v1", "workers", "heartbeat"]:
self._exact(body, set())
return 200, service.heartbeat(identity)
if method == "POST" and segments == ["v1", "workers", "claim"]:
self._exact(body, set())
return 200, service.claim_next(identity)
if method == "POST" and segments == ["v1", "workers", "renew"]:
self._exact(body, {"task_id", "claim_id"})
return 200, service.renew(
identity, _safe_identifier(body["task_id"]), _safe_identifier(body["claim_id"])
)
if method == "POST" and segments == ["v1", "workers", "execute"]:
self._exact(body, {"task_id", "claim_id"})
return 200, service.execute_claim(
identity, _safe_identifier(body["task_id"]), _safe_identifier(body["claim_id"]),deadline=self.request_deadline
)
raise _RequestError("HTTP_NOT_FOUND", 404)
def _handle(self, method: str) -> None:
try:
self.request_deadline=self._deadline()
identity = self._identity()
segments = self._path()
body = self._json_body() if method == "POST" else None
if method == "GET":
lengths = self.headers.get_all("Content-Length", [])
if self.headers.get("Transfer-Encoding") is not None:
raise _RequestError("HTTP_GET_BODY_FORBIDDEN", 400)
if len(lengths) > 1 or (lengths and (not lengths[0].isdigit() or int(lengths[0]) != 0)):
raise _RequestError("HTTP_GET_BODY_FORBIDDEN", 400)
status, result = self._dispatch(method, identity, segments, body)
self._send(status, result)
except _RequestError as exc:
self._error(exc)
except ServiceError as exc:
status = exc.status if isinstance(exc.status, int) and 400 <= exc.status <= 599 else 500
self._send(status, {"error": {"code": exc.code}})
except TimeoutError:
self._send(408, {"error": {"code": "HTTP_DEADLINE_EXPIRED"}})
except Exception:
# No exception string, traceback, request headers, or token crosses
# this boundary.
self._send(500, {"error": {"code": "HTTP_INTERNAL_ERROR"}})
def do_GET(self):
self._handle("GET")
def do_POST(self):
self._handle("POST")
def _method_not_allowed(self):
try:
self._deadline()
self._identity()
self._send(405, {"error": {"code": "HTTP_METHOD_NOT_ALLOWED"}})
except _RequestError as exc:
self._error(exc)
do_PUT = _method_not_allowed
do_PATCH = _method_not_allowed
do_DELETE = _method_not_allowed
do_OPTIONS = _method_not_allowed
do_HEAD = _method_not_allowed
do_TRACE = _method_not_allowed
do_CONNECT = _method_not_allowed
def serve(
service: Any,
tokens: Mapping[str, Identity],
host: str = "127.0.0.1",
port: int = 0,
) -> ThreadingHTTPServer:
"""Build an unstarted real HTTP server bound only to an IPv4 loopback."""
host = _loopback_host(host, "HTTP_LOOPBACK_REQUIRED")
if type(port) is not int or not 0 <= port <= 65_535:
raise TransportError("HTTP_PORT_INVALID")
if not isinstance(tokens, Mapping) or not tokens:
raise TransportError("HTTP_TRUSTED_IDENTITIES_REQUIRED")
identities: dict[str, Identity] = {}
for token, identity in tokens.items():
if (
not isinstance(token, str) or not 32<=len(token)<=512
or any(ord(char) < 33 or ord(char) > 126 for char in token)
or not isinstance(identity, Identity)
):
raise TransportError("HTTP_TRUSTED_IDENTITY_INVALID")
identities[token] = identity
return _BaselineHTTPServer((host, port), _Handler, service, identities)
class _NoRedirect(HTTPRedirectHandler):
def redirect_request(self, req, fp, code, msg, headers, newurl):
return None
class HttpTransport:
"""Small urllib client for the loopback API; redirects and proxies are off."""
def __init__(
self,
base_url: str,
token: str,
*,
timeout_seconds: float = 15,
max_response_bytes: int = MAX_RESPONSE_BYTES,
):
try:
parsed = urlsplit(base_url)
parsed_port = parsed.port
except ValueError as exc:
raise TransportError("TRANSPORT_BASE_URL_INVALID") from exc
if (
parsed.scheme != "http" or parsed.username is not None or parsed.password is not None
or parsed.query or parsed.fragment or parsed.path not in {"", "/"}
or parsed.hostname is None or parsed_port is None
):
raise TransportError("TRANSPORT_BASE_URL_INVALID")
_loopback_host(parsed.hostname, "TRANSPORT_LOOPBACK_REQUIRED")
if not isinstance(token, str) or not 32<=len(token)<=512 or any(
ord(char) < 33 or ord(char) > 126 for char in token
):
raise TransportError("TRANSPORT_TOKEN_INVALID")
if (
isinstance(timeout_seconds, bool)
or not isinstance(timeout_seconds, (int, float))
or not math.isfinite(float(timeout_seconds))
or not 0.1 <= timeout_seconds <= MAX_DEADLINE_SECONDS
):
raise TransportError("TRANSPORT_TIMEOUT_INVALID")
if type(max_response_bytes) is not int or not 1 <= max_response_bytes <= 4 * 1024 * 1024:
raise TransportError("TRANSPORT_RESPONSE_LIMIT_INVALID")
self.base_url = f"http://{parsed.hostname}:{parsed_port}"
self.token = token
self.timeout_seconds = float(timeout_seconds)
self.max_response_bytes = max_response_bytes
self._opener = build_opener(ProxyHandler({}), _NoRedirect())
def _decode(self, response) -> Any:
content_type = response.headers.get("Content-Type", "").split(";", 1)[0]
if content_type != "application/json":
raise TransportError("TRANSPORT_RESPONSE_TYPE_INVALID")
declared = response.headers.get("Content-Length")
if declared is None or not declared.isdigit() or int(declared) > self.max_response_bytes:
raise TransportError("TRANSPORT_RESPONSE_TOO_LARGE")
raw = response.read(self.max_response_bytes + 1)
if len(raw) > self.max_response_bytes or len(raw) != int(declared):
raise TransportError("TRANSPORT_RESPONSE_SIZE_INVALID")
try:
def unique_object(pairs):
result = {}
for key, item in pairs:
if key in result:
raise ValueError("duplicate")
result[key] = item
return result
return json.loads(
raw.decode("utf-8"), object_pairs_hook=unique_object,
parse_constant=lambda _value: (_ for _ in ()).throw(ValueError("non-finite")),
)
except (UnicodeDecodeError, json.JSONDecodeError, ValueError) as exc:
raise TransportError("TRANSPORT_RESPONSE_JSON_INVALID") from exc
def _request(self, method: str, path: str, body: Mapping[str, Any] | None = None) -> Any:
if not isinstance(path, str) or not path.startswith("/") or "?" in path or "#" in path:
raise TransportError("TRANSPORT_PATH_INVALID")
data = None
headers = {
"Authorization": "Bearer " + self.token,
"Accept": "application/json",
"X-AIOS-Deadline": str(time.time() + self.timeout_seconds),
}
if method == "POST":
if not isinstance(body, Mapping):
raise TransportError("TRANSPORT_BODY_INVALID")
try:
data = json.dumps(
body, sort_keys=True, separators=(",", ":"), ensure_ascii=False,
allow_nan=False,
).encode("utf-8")
except (TypeError, ValueError) as exc:
raise TransportError("TRANSPORT_BODY_INVALID") from exc
if not data or len(data) > MAX_REQUEST_BYTES:
raise TransportError("TRANSPORT_BODY_TOO_LARGE")
headers["Content-Type"] = "application/json"
request = Request(self.base_url + path, data=data, headers=headers, method=method)
try:
with self._opener.open(request, timeout=self.timeout_seconds) as response:
return self._decode(response)
except HTTPError as exc:
if 300 <= exc.code < 400:
exc.close()
raise TransportError("TRANSPORT_REDIRECT_REJECTED", exc.code) from None
try:
payload = self._decode(exc)
except TransportError:
payload = None
except (TimeoutError, OSError):
raise TransportError("TRANSPORT_UNAVAILABLE") from None
finally:
exc.close()
code = None
if isinstance(payload, Mapping) and isinstance(payload.get("error"), Mapping):
code = payload["error"].get("code")
if not isinstance(code, str) or _IDENTIFIER.fullmatch(code) is None:
code = "TRANSPORT_HTTP_ERROR"
raise TransportError(code, exc.code) from None
except (URLError, TimeoutError, OSError):
raise TransportError("TRANSPORT_UNAVAILABLE") from None
def submit_goal(self, request: Mapping[str, Any], *, idempotency_key: str) -> dict:
return self._request(
"POST", "/v1/goals/submit",
{"request": request, "idempotency_key": idempotency_key},
)
def get_task(self, task_id: str) -> dict:
return self._request("GET", "/v1/tasks/" + quote(task_id, safe=""))
def cancel_task(self, task_id: str) -> dict:
return self._request("POST", "/v1/tasks/" + quote(task_id, safe="") + "/cancel", {})
def explain(self, mission_id: str) -> dict:
return self._request("GET", "/v1/missions/" + quote(mission_id, safe="") + "/explain")
def register_worker(self) -> dict:
return self._request("POST", "/v1/workers/register", {})
def heartbeat(self) -> dict:
return self._request("POST", "/v1/workers/heartbeat", {})
def claim_next(self) -> dict | None:
return self._request("POST", "/v1/workers/claim", {})
def renew(self, task_id: str, claim_id: str) -> dict:
return self._request(
"POST", "/v1/workers/renew", {"task_id": task_id, "claim_id": claim_id}
)
def execute_claim(self, task_id: str, claim_id: str) -> dict:
return self._request(
"POST", "/v1/workers/execute", {"task_id": task_id, "claim_id": claim_id}
)
# Goal-oriented names used by thin clients.
submit = submit_goal
query = get_task
cancel = cancel_task
Transport = HttpTransport
__all__ = ["HttpTransport", "Transport", "TransportError", "serve"]