128 lines
3.8 KiB
Python
128 lines
3.8 KiB
Python
# retoor <retoor@molodetz.nl>
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import binascii
|
|
import json
|
|
import time
|
|
from dataclasses import dataclass
|
|
|
|
from devplacepy.database import nodes as store
|
|
from devplacepy.services.nodes.policy import Target
|
|
from devplacepy.services.nodes.spec import SPEC
|
|
|
|
LIMITS = SPEC.limits
|
|
|
|
DONE = "done"
|
|
FAILED = "failed"
|
|
EXPIRED = "expired"
|
|
TERMINAL = (DONE, FAILED, EXPIRED)
|
|
|
|
MAX_HEADERS = 64
|
|
MAX_HEADER_CHARS = 1024
|
|
DROPPED_HEADERS = frozenset({"set-cookie", "set-cookie2", "authorization", "proxy-authenticate"})
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BrokerResult:
|
|
uid: str
|
|
status: str
|
|
http_status: int
|
|
headers: dict
|
|
body: bytes
|
|
truncated: bool
|
|
error: str
|
|
via: str = "fetch"
|
|
|
|
@property
|
|
def ok(self) -> bool:
|
|
return self.status == DONE
|
|
|
|
|
|
class BodyError(ValueError):
|
|
pass
|
|
|
|
|
|
def clean_headers(raw) -> dict[str, str]:
|
|
if not isinstance(raw, dict):
|
|
return {}
|
|
kept: dict[str, str] = {}
|
|
for key, value in raw.items():
|
|
if len(kept) >= MAX_HEADERS:
|
|
break
|
|
if not isinstance(key, str) or not isinstance(value, str):
|
|
continue
|
|
name = key.strip().lower()
|
|
if not name or name in DROPPED_HEADERS:
|
|
continue
|
|
kept[name[:MAX_HEADER_CHARS]] = value[:MAX_HEADER_CHARS]
|
|
return kept
|
|
|
|
|
|
def decode_body(body_b64: str) -> tuple[bytes, bool]:
|
|
try:
|
|
body = base64.b64decode(body_b64.encode(), validate=True)
|
|
except (binascii.Error, ValueError):
|
|
raise BodyError("bad_base64")
|
|
if len(body) > LIMITS.max_body_bytes:
|
|
return body[: LIMITS.max_body_bytes], True
|
|
return body, False
|
|
|
|
|
|
class JobBroker:
|
|
def enqueue(self, caller_kind: str, caller_id: str, target: Target, mode: str) -> str:
|
|
return store.enqueue_request(
|
|
caller_kind, caller_id, target.url, target.method, target.adapter, mode
|
|
)
|
|
|
|
def queued(self, limit: int = LIMITS.dispatch_batch) -> list[dict]:
|
|
return store.list_queued(limit)
|
|
|
|
def claim(self, uid: str, device_uid: str) -> bool:
|
|
return store.claim_request(uid, device_uid)
|
|
|
|
def complete(self, uid: str, http_status: int, headers: dict, body: bytes, truncated: bool) -> bool:
|
|
return store.finish_request(
|
|
uid,
|
|
result_status=http_status,
|
|
result_headers=json.dumps(headers),
|
|
result_body=base64.b64encode(body).decode(),
|
|
result_bytes=len(body),
|
|
truncated=truncated,
|
|
)
|
|
|
|
def fail(self, uid: str, error: str, *, expired: bool = False) -> bool:
|
|
return store.fail_request(uid, error, expired=expired)
|
|
|
|
def expire(self) -> dict[str, int]:
|
|
return store.expire_requests()
|
|
|
|
def read(self, uid: str) -> BrokerResult | None:
|
|
row = store.get_request(uid)
|
|
if row is None:
|
|
return None
|
|
return BrokerResult(
|
|
uid=row["uid"],
|
|
status=row["status"],
|
|
http_status=int(row.get("result_status") or 0),
|
|
headers=json.loads(row.get("result_headers") or "{}") if row.get("result_headers") else {},
|
|
body=base64.b64decode(row["result_body"]) if row.get("result_body") else b"",
|
|
truncated=bool(row.get("truncated")),
|
|
error=row.get("error") or "",
|
|
via=row.get("mode") if row.get("mode") in ("tunnel", "fetch") else "fetch",
|
|
)
|
|
|
|
async def wait(self, uid: str, timeout: float) -> BrokerResult:
|
|
deadline = time.monotonic() + timeout
|
|
pause = LIMITS.poll_interval_ms / 1000
|
|
while True:
|
|
row = store.get_request(uid)
|
|
if row is not None and row["status"] in TERMINAL:
|
|
return self.read(uid)
|
|
if time.monotonic() >= deadline:
|
|
self.fail(uid, "caller_timeout")
|
|
return self.read(uid)
|
|
await asyncio.sleep(pause)
|