98 lines
3.9 KiB
Python
98 lines
3.9 KiB
Python
"""Private loopback transport over the operator's dedicated SSH reverse tunnel."""
|
|
|
|
import hashlib
|
|
import http.client
|
|
import json
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
from urllib.parse import urlsplit
|
|
|
|
PREFIX = "/api/v1/ai-polygon"
|
|
|
|
|
|
class CoreClient:
|
|
def __init__(self, origin: str, token_file: Path, instance: str):
|
|
url = urlsplit(origin)
|
|
if (
|
|
url.scheme != "http"
|
|
or url.hostname != "127.0.0.1"
|
|
or url.path
|
|
or url.query
|
|
or url.fragment
|
|
or url.username
|
|
or not url.port
|
|
):
|
|
raise ValueError("Core must use an explicit loopback SSH tunnel origin")
|
|
self.port = url.port
|
|
self.token = token_file.read_text().strip()
|
|
if not 32 <= len(self.token) <= 512:
|
|
raise ValueError("Invalid simulation credential")
|
|
self.instance = instance
|
|
self._connections = threading.local()
|
|
|
|
def request(self, path: str, body=None):
|
|
connection = getattr(self._connections, "connection", None)
|
|
if (
|
|
connection is not None
|
|
and time.monotonic() - getattr(self._connections, "last_used", 0) > 2
|
|
):
|
|
connection.close()
|
|
connection = None
|
|
if connection is None:
|
|
connection = http.client.HTTPConnection("127.0.0.1", self.port, timeout=10)
|
|
self._connections.connection = connection
|
|
try:
|
|
connection.request(
|
|
"GET" if body is None else "POST",
|
|
PREFIX + path,
|
|
body=None if body is None else json.dumps(body, allow_nan=False).encode(),
|
|
headers={
|
|
"Authorization": "Bearer " + self.token,
|
|
"Worker-Instance": self.instance,
|
|
"Content-Type": "application/json",
|
|
},
|
|
)
|
|
response = connection.getresponse()
|
|
raw = response.read(4 * 1024**2 + 1)
|
|
if response.status not in (200, 201) or len(raw) > 4 * 1024**2:
|
|
raise RuntimeError("Core rejected simulation operation: " + str(response.status))
|
|
self._connections.last_used = time.monotonic()
|
|
return json.loads(raw)
|
|
except Exception:
|
|
connection.close()
|
|
self._connections.connection = None
|
|
# Never retry an uncertain mutation automatically.
|
|
raise
|
|
|
|
def download(self, world: dict, target: Path):
|
|
if target.exists():
|
|
with target.open("rb") as stream:
|
|
if hashlib.file_digest(stream, "sha256").hexdigest() == world["sha256"]:
|
|
return
|
|
if world.get("storage", {}).get("kind") == "worker":
|
|
raise RuntimeError("The admitted Worker-local scene is missing or changed")
|
|
connection = http.client.HTTPConnection("127.0.0.1", self.port, timeout=30)
|
|
temporary = target.with_suffix(".part")
|
|
try:
|
|
connection.request("GET", PREFIX + "/worlds/" + world["world_id"] + "/source.ply")
|
|
response = connection.getresponse()
|
|
if (
|
|
response.status != 200
|
|
or int(response.getheader("Content-Length") or 0) != world["byte_length"]
|
|
):
|
|
raise RuntimeError("Scene transfer contract changed")
|
|
digest, size = hashlib.sha256(), 0
|
|
with temporary.open("wb") as stream:
|
|
while block := response.read(4 * 1024**2):
|
|
size += len(block)
|
|
if size > world["byte_length"]:
|
|
raise RuntimeError("Scene transfer exceeds admitted size")
|
|
digest.update(block)
|
|
stream.write(block)
|
|
if size != world["byte_length"] or digest.hexdigest() != world["sha256"]:
|
|
raise RuntimeError("Scene transfer identity changed")
|
|
temporary.replace(target)
|
|
finally:
|
|
connection.close()
|