mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 05:23:41 +08:00
firestarsdog
96 lines
2.9 KiB
Python
96 lines
2.9 KiB
Python
"""Bounded local replay operations shared by native UI and Galaxy."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
import socket
|
|
import struct
|
|
import time
|
|
|
|
from openpilot.starpilot.storage import starpilot_storage_root
|
|
|
|
|
|
MAX_REQUEST = 4096
|
|
MAX_RESPONSE = 262144
|
|
MAX_WORKER_STATUS = 65536
|
|
|
|
|
|
def replay_root() -> Path:
|
|
return starpilot_storage_root() / 'replay'
|
|
|
|
|
|
def socket_path() -> Path:
|
|
return replay_root() / 'owner.sock'
|
|
|
|
|
|
class ReplayUnavailable(RuntimeError):
|
|
pass
|
|
|
|
|
|
def _unique(pairs):
|
|
value = {}
|
|
for key, item in pairs:
|
|
if key in value:
|
|
raise ValueError('duplicate field')
|
|
value[key] = item
|
|
return value
|
|
|
|
|
|
def parse(raw: bytes, limit: int) -> dict:
|
|
if len(raw) > limit:
|
|
raise ValueError('message exceeds limit')
|
|
value = json.loads(raw.decode('utf-8'), object_pairs_hook=_unique,
|
|
parse_constant=lambda _: (_ for _ in ()).throw(ValueError('nonfinite value')))
|
|
if type(value) is not dict:
|
|
raise ValueError('message must be an object')
|
|
return value
|
|
|
|
|
|
def encode(value: dict, limit: int) -> bytes:
|
|
raw = json.dumps(value, separators=(',', ':'), allow_nan=False).encode('utf-8')
|
|
if not raw or len(raw) > limit:
|
|
raise ValueError('message exceeds limit')
|
|
return struct.pack('!I', len(raw)) + raw
|
|
|
|
|
|
def receive(sock: socket.socket, limit: int, timeout: float) -> dict:
|
|
deadline = time.monotonic() + timeout
|
|
|
|
def exact(size):
|
|
output = bytearray()
|
|
while len(output) < size:
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
raise ReplayUnavailable('Replay response timed out')
|
|
sock.settimeout(remaining)
|
|
chunk = sock.recv(size - len(output))
|
|
if not chunk:
|
|
raise ReplayUnavailable('Replay connection closed')
|
|
output.extend(chunk)
|
|
return bytes(output)
|
|
|
|
length = struct.unpack('!I', exact(4))[0]
|
|
if not 0 < length <= limit:
|
|
raise ValueError('message size invalid')
|
|
return parse(exact(length), limit)
|
|
|
|
|
|
def request_operation(request: dict, *, socket_path: Path | None = None, timeout: float = 2.0) -> dict:
|
|
if type(request) is not dict or type(timeout) not in (int, float) or not 0 < timeout <= 4:
|
|
raise ReplayUnavailable('Invalid replay operation')
|
|
path = replay_root() / 'owner.sock' if socket_path is None else Path(socket_path)
|
|
try:
|
|
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as peer:
|
|
peer.settimeout(timeout)
|
|
peer.connect(str(path))
|
|
peer.sendall(encode(request, MAX_REQUEST))
|
|
response = receive(peer, MAX_RESPONSE, timeout)
|
|
except (OSError, ValueError, TypeError, UnicodeError, RecursionError) as error:
|
|
raise ReplayUnavailable('Replay owner is unavailable') from error
|
|
if (type(response.get('ok')) is not bool or
|
|
set(response) != ({'ok', 'result'} if response['ok'] else {'ok', 'error'}) or
|
|
type(response.get('result' if response['ok'] else 'error')) is not dict):
|
|
raise ReplayUnavailable('Replay response is invalid')
|
|
return response
|