Fix private production pilot recovery, screen transfer and MCP dispatch
This commit is contained in:
1 parent
f6b14c5210
commit
db5f85c922
18 files changed
+1735
-66
No files matched your search
Binary file not shown.
@@ -0,0 +1,316 @@
|
||||
"""Portable registration integration through the bundled FastMCP dispatcher.
|
||||
|
||||
Executors are bounded fixtures. No installed adapter, desktop or Fleet node is
|
||||
contacted; these checks establish dispatch and identity wiring, not acceptance.
|
||||
"""
|
||||
import ast
|
||||
import asyncio
|
||||
import base64
|
||||
from functools import partial, wraps
|
||||
import hashlib
|
||||
import inspect
|
||||
import json
|
||||
import math
|
||||
from pathlib import Path
|
||||
import re
|
||||
import shlex
|
||||
import socket
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from unittest import mock
|
||||
import uuid
|
||||
|
||||
WORK = Path(__file__).resolve().parents[1]
|
||||
VENDOR = WORK / "sources/original-geranas/FrameBridge-Fleet-Hub/vendor"
|
||||
sys.path.insert(0, str(WORK / "expanded"))
|
||||
sys.path.insert(0, str(VENDOR))
|
||||
import anyio
|
||||
from mcp.server.fastmcp import FastMCP, Image
|
||||
from mcp.types import ToolAnnotations
|
||||
from framebridge_expanded import tty_mcp
|
||||
from framebridge_expanded.real_screen.screen_frontdoor import register_framebridge_screen_tools
|
||||
from framebridge_expanded.real_screen.screen_presentation import register_screen_presentation
|
||||
from test_screen_frontdoor import Executor, NODE_ID, SESSION_ID
|
||||
|
||||
PORTABLE = WORK / "portable/front_door_mcp.py"
|
||||
|
||||
|
||||
class ChunkExecutor(Executor):
|
||||
"""Honor the typed route's optional larger chunk while retaining bindings."""
|
||||
def __call__(self, **kwargs):
|
||||
outer = super().__call__(**kwargs)
|
||||
request = json.loads(base64.b64decode(shlex.split(kwargs["script"])[-1]))
|
||||
if request["op"] == "read_chunk" and request.get("chunk_bytes"):
|
||||
size, index = request["chunk_bytes"], request["index"]
|
||||
result = json.loads(outer["stdout"])
|
||||
result["result"].update(
|
||||
total_chunks=math.ceil(len(self.data) / size),
|
||||
data=base64.b64encode(self.data[index * size:(index + 1) * size]).decode(),
|
||||
more=(index + 1) * size < len(self.data),
|
||||
)
|
||||
outer["stdout"] = json.dumps(result)
|
||||
return outer
|
||||
|
||||
|
||||
def portable(*, execute=None, screen_execute=None, threaded=True):
|
||||
"""Load the actual registration statements without bootstrap/profile effects."""
|
||||
tree = ast.parse(PORTABLE.read_text())
|
||||
names = {"_BlockingToolRegistrar", "_registered_operations", "framebridge_status",
|
||||
"framebridge_exec", "_screen_exec_request"}
|
||||
assignments = {"_blocking_tools", "_real_screen_client"}
|
||||
registrations = {"register_screen_presentation", "register_tty_tools"}
|
||||
selected = []
|
||||
for node in tree.body:
|
||||
if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in names:
|
||||
selected.append(node)
|
||||
elif isinstance(node, ast.Assign) and any(
|
||||
isinstance(target, ast.Name) and target.id in assignments for target in node.targets):
|
||||
if not threaded and node.targets[0].id == "_blocking_tools":
|
||||
node.value = ast.Name(id="mcp", ctx=ast.Load())
|
||||
selected.append(node)
|
||||
elif (isinstance(node, ast.Expr) and isinstance(node.value, ast.Call)
|
||||
and isinstance(node.value.func, ast.Name) and node.value.func.id in registrations):
|
||||
selected.append(node)
|
||||
mcp = FastMCP("dispatch fixture")
|
||||
execute = execute or (lambda *args: {"ok": True})
|
||||
screen_execute = screen_execute or (lambda *args: {"ok": False})
|
||||
namespace = dict(mcp=mcp, anyio=anyio, partial=partial, wraps=wraps, inspect=inspect,
|
||||
ToolAnnotations=ToolAnnotations, Image=Image, socket=socket,
|
||||
ADAPTER_SOCKET=Path("/nonexistent/framebridge-test.sock"),
|
||||
PRODUCT_VERSION="test", GUIDE_SCHEMA=1,
|
||||
_framebridge_exec_geranas_impl=execute, expanded_execute_node=execute,
|
||||
expanded_execute_screen_node=screen_execute,
|
||||
register_framebridge_screen_tools=register_framebridge_screen_tools,
|
||||
resolve_profile_node=lambda target: target.upper(),
|
||||
register_screen_presentation=register_screen_presentation,
|
||||
register_tty_tools=tty_mcp.register_tools)
|
||||
exec(compile(ast.fix_missing_locations(ast.Module(body=selected, type_ignores=[])),
|
||||
str(PORTABLE), "exec"), namespace)
|
||||
return mcp, namespace
|
||||
|
||||
|
||||
def structured(reply):
|
||||
if isinstance(reply, tuple):
|
||||
content, value = reply
|
||||
if value is not None:
|
||||
return value
|
||||
reply = content
|
||||
if isinstance(reply, dict):
|
||||
return reply
|
||||
return json.loads(next(item.text for item in reply if item.type == "text"))
|
||||
|
||||
|
||||
async def until_set(event):
|
||||
async def poll():
|
||||
while not event.is_set():
|
||||
await asyncio.sleep(.001)
|
||||
await asyncio.wait_for(poll(), 1)
|
||||
|
||||
|
||||
class DispatchTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_actual_sdk_returns_status_while_exec_and_tty_capture_are_running(self):
|
||||
exec_started, capture_started, release = (threading.Event() for _ in range(3))
|
||||
fixture = ChunkExecutor()
|
||||
calls = []
|
||||
|
||||
def execute(*args):
|
||||
calls.append(args)
|
||||
exec_started.set()
|
||||
release.wait(2)
|
||||
return {"ok": True, "marker": args[0]}
|
||||
|
||||
def screen(alias, script, authorized, timeout_s, job_id, sha256, include_stdout):
|
||||
capture_started.set()
|
||||
release.wait(2)
|
||||
return fixture(target=alias, script=script, authorized=authorized,
|
||||
timeout_s=timeout_s, job_id=job_id, sha256=sha256,
|
||||
include_stdout=include_stdout)
|
||||
|
||||
mcp, _ = portable(execute=execute, screen_execute=screen)
|
||||
slow = asyncio.create_task(mcp.call_tool("framebridge_exec", {
|
||||
"script": "bounded-marker", "authorized": True}))
|
||||
capture = asyncio.create_task(mcp.call_tool("cua_tty_observe", {
|
||||
"session_id": SESSION_ID, "authorized": True}))
|
||||
try:
|
||||
await until_set(exec_started)
|
||||
await until_set(capture_started)
|
||||
status = structured(await asyncio.wait_for(mcp.call_tool("framebridge_status", {}), .3))
|
||||
self.assertTrue(status["ok"])
|
||||
self.assertFalse(slow.done())
|
||||
self.assertFalse(capture.done())
|
||||
finally:
|
||||
release.set()
|
||||
exec_result, capture_result = await asyncio.gather(slow, capture)
|
||||
self.assertEqual(structured(exec_result)["marker"], "bounded-marker")
|
||||
self.assertEqual(len(calls), 1)
|
||||
self.assertTrue(any(item.type == "image" for item in capture_result))
|
||||
|
||||
async def test_fast_exec_response_is_delivered_before_slow_exec_completes(self):
|
||||
started, release = threading.Event(), threading.Event()
|
||||
calls = []
|
||||
|
||||
def execute(*args):
|
||||
marker = args[0]
|
||||
calls.append(marker)
|
||||
if marker == "slow":
|
||||
started.set()
|
||||
release.wait(2)
|
||||
return {"ok": True, "marker": marker}
|
||||
|
||||
mcp, _ = portable(execute=execute)
|
||||
slow = asyncio.create_task(mcp.call_tool("framebridge_exec", {
|
||||
"script": "slow", "authorized": True}))
|
||||
try:
|
||||
await until_set(started)
|
||||
fast = structured(await asyncio.wait_for(mcp.call_tool("framebridge_exec", {
|
||||
"script": "fast", "authorized": True}), .3))
|
||||
self.assertEqual(fast["marker"], "fast")
|
||||
self.assertFalse(slow.done())
|
||||
finally:
|
||||
release.set()
|
||||
slow_result = await slow
|
||||
self.assertEqual(structured(slow_result)["marker"], "slow")
|
||||
self.assertCountEqual(calls, ["slow", "fast"])
|
||||
|
||||
async def test_full_portable_import_registers_tools_without_bootstrap_side_effects(self):
|
||||
import importlib.util
|
||||
from framebridge_expanded import fleet
|
||||
spec = importlib.util.spec_from_file_location("frontdoor_dispatch_test", PORTABLE)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
with mock.patch.object(fleet, "bootstrap_original") as bootstrap:
|
||||
spec.loader.exec_module(module)
|
||||
bootstrap.assert_called_once_with()
|
||||
names = {tool.name for tool in await module.mcp.list_tools()}
|
||||
self.assertTrue(set(tty_mcp.TOOL_NAMES) <= names)
|
||||
self.assertTrue({"framebridge_exec", "framebridge_cua", "framebridge_screen_show"} <= names)
|
||||
self.assertFalse(inspect.iscoroutinefunction(module.framebridge_exec))
|
||||
|
||||
async def test_schemas_metadata_and_internal_sync_callbacks_are_preserved(self):
|
||||
mcp, namespace = portable()
|
||||
baseline, _ = portable(threaded=False)
|
||||
actual = {tool.name: tool.model_dump() for tool in await mcp.list_tools()}
|
||||
expected = {tool.name: tool.model_dump() for tool in await baseline.list_tools()}
|
||||
self.assertEqual(actual, expected)
|
||||
self.assertFalse(inspect.iscoroutinefunction(namespace["framebridge_exec"]))
|
||||
self.assertFalse(inspect.iscoroutinefunction(namespace["_screen_exec_request"]))
|
||||
self.assertIsInstance(namespace["framebridge_exec"]("marker", True), dict)
|
||||
|
||||
async def test_local_and_remote_tty_capture_use_identity_producing_transport(self):
|
||||
for target in ("LOCAL", "WORKER_1"):
|
||||
with self.subTest(target=target):
|
||||
fixture, calls = ChunkExecutor(), []
|
||||
|
||||
def screen(alias, script, authorized, timeout_s, job_id, sha256, include_stdout):
|
||||
calls.append(alias)
|
||||
return fixture(target=alias, script=script, authorized=authorized,
|
||||
timeout_s=timeout_s, job_id=job_id, sha256=sha256,
|
||||
include_stdout=include_stdout)
|
||||
|
||||
legacy = mock.Mock(side_effect=AssertionError("legacy EXEC must not capture"))
|
||||
mcp, _ = portable(execute=legacy, screen_execute=screen)
|
||||
result = await mcp.call_tool("cua_tty_observe", {
|
||||
"session_id": SESSION_ID, "authorized": True, "target": target})
|
||||
self.assertTrue(any(item.type == "image" for item in result))
|
||||
self.assertEqual(set(calls), {target})
|
||||
self.assertGreaterEqual(len(calls), 2)
|
||||
legacy.assert_not_called()
|
||||
|
||||
async def test_missing_uuid_and_wrong_reply_binding_still_fail_closed(self):
|
||||
for corruption, expected in (("uuid", "SCREEN_EXEC_NODE_ID_INVALID"),
|
||||
("hash", "SCREEN_EXEC_BINDING_INVALID")):
|
||||
with self.subTest(corruption=corruption):
|
||||
fixture = ChunkExecutor()
|
||||
|
||||
def screen(alias, script, authorized, timeout_s, job_id, sha256, include_stdout):
|
||||
result = fixture(target=alias, script=script, authorized=authorized,
|
||||
timeout_s=timeout_s, job_id=job_id, sha256=sha256,
|
||||
include_stdout=include_stdout)
|
||||
if corruption == "uuid":
|
||||
result.pop("node_id")
|
||||
else:
|
||||
result["mcp_script_sha256"] = "0" * 64
|
||||
return result
|
||||
|
||||
mcp, _ = portable(screen_execute=screen)
|
||||
result = structured(await mcp.call_tool("cua_tty_observe", {
|
||||
"session_id": SESSION_ID, "authorized": True}))
|
||||
self.assertFalse(result["ok"])
|
||||
self.assertEqual(result["error"], expected)
|
||||
|
||||
async def test_other_tty_rpcs_also_leave_status_responsive(self):
|
||||
started, release = threading.Event(), threading.Event()
|
||||
mcp, _ = portable()
|
||||
|
||||
def request(*args):
|
||||
started.set()
|
||||
release.wait(2)
|
||||
return {"ok": True}
|
||||
|
||||
with mock.patch.object(tty_mcp, "_remote_request", side_effect=request):
|
||||
slow = asyncio.create_task(mcp.call_tool("cua_tty_status", {
|
||||
"authorized": True, "target": "WORKER_1"}))
|
||||
try:
|
||||
await until_set(started)
|
||||
self.assertTrue(structured(await asyncio.wait_for(
|
||||
mcp.call_tool("framebridge_status", {}), .3))["ok"])
|
||||
self.assertFalse(slow.done())
|
||||
finally:
|
||||
release.set()
|
||||
await slow
|
||||
|
||||
async def test_anyio_cancellation_waits_for_single_worker_without_replay(self):
|
||||
started, release, finished = (threading.Event() for _ in range(3))
|
||||
mcp, namespace = portable()
|
||||
calls, scopes = [], []
|
||||
|
||||
@namespace["_blocking_tools"].tool()
|
||||
def bounded_operation() -> dict:
|
||||
calls.append(1)
|
||||
started.set()
|
||||
release.wait(2)
|
||||
finished.set()
|
||||
return {"ok": True}
|
||||
|
||||
async def dispatch():
|
||||
with anyio.CancelScope() as scope:
|
||||
scopes.append(scope)
|
||||
await mcp.call_tool("bounded_operation", {})
|
||||
|
||||
async with anyio.create_task_group() as group:
|
||||
group.start_soon(dispatch)
|
||||
await until_set(started)
|
||||
scopes[0].cancel()
|
||||
await anyio.sleep(.02)
|
||||
self.assertFalse(finished.is_set())
|
||||
release.set()
|
||||
self.assertTrue(finished.is_set())
|
||||
self.assertEqual(calls, [1])
|
||||
|
||||
|
||||
class AutomaticJobTests(unittest.TestCase):
|
||||
def test_identical_concurrent_auto_jobs_are_distinct_and_explicit_ids_unchanged(self):
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
tree = ast.parse(PORTABLE.read_text())
|
||||
function = next(node for node in tree.body if isinstance(node, ast.FunctionDef)
|
||||
and node.name == "_framebridge_exec_geranas_impl")
|
||||
namespace = dict(hashlib=hashlib, time=time, uuid=uuid,
|
||||
JOB_RE=re.compile(r"^[A-Za-z0-9._-]{1,80}$"),
|
||||
adapter_request=lambda payload: {"ok": True, "job_id": payload["job_id"]},
|
||||
sanitize_text=lambda text: text)
|
||||
exec(compile(ast.Module(body=[function], type_ignores=[]), str(PORTABLE), "exec"), namespace)
|
||||
execute = namespace["_framebridge_exec_geranas_impl"]
|
||||
with mock.patch.object(time, "strftime", return_value="20261028-000000"):
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
results = list(pool.map(lambda _: execute("same-command", True), range(16)))
|
||||
self.assertEqual(len({result["job_id"] for result in results}), 16)
|
||||
self.assertTrue(all(namespace["JOB_RE"].fullmatch(result["job_id"]) for result in results))
|
||||
first = execute("same-command", True, job_id="owner-idempotent-id")
|
||||
second = execute("same-command", True, job_id="owner-idempotent-id")
|
||||
self.assertEqual(first["job_id"], "owner-idempotent-id")
|
||||
self.assertEqual(first["job_id"], second["job_id"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,224 @@
|
||||
"""Deterministic worker initialization races, using only in-memory fixtures."""
|
||||
import ast
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
WORK = Path(__file__).resolve().parents[1]
|
||||
EXPANDED = WORK / "expanded/framebridge_expanded"
|
||||
|
||||
|
||||
def definitions(module, names, namespace):
|
||||
tree = ast.parse((EXPANDED / module).read_text())
|
||||
selected = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in names]
|
||||
exec(compile(ast.Module(body=selected, type_ignores=[]), str(EXPANDED / module), "exec"), namespace)
|
||||
return namespace
|
||||
|
||||
|
||||
class ObservedRLock:
|
||||
"""Signal when the competing worker reaches the lock, without timing guesses."""
|
||||
def __init__(self):
|
||||
self.lock = threading.RLock()
|
||||
self.waiting = threading.Event()
|
||||
self.watch = None
|
||||
|
||||
def __enter__(self):
|
||||
if threading.get_ident() == self.watch:
|
||||
self.waiting.set()
|
||||
self.lock.acquire()
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc):
|
||||
self.lock.release()
|
||||
|
||||
def watched(self, function):
|
||||
self.watch = threading.get_ident()
|
||||
return function()
|
||||
|
||||
|
||||
class WorkerInitializationTests(unittest.TestCase):
|
||||
def test_runtime_publishes_complete_module_once_to_competing_workers(self):
|
||||
started, release = threading.Event(), threading.Event()
|
||||
lock, module = ObservedRLock(), SimpleNamespace()
|
||||
loads = []
|
||||
|
||||
def initialize(target):
|
||||
loads.append(target)
|
||||
started.set()
|
||||
if not release.wait(2):
|
||||
raise AssertionError("fixture load was not released")
|
||||
target.settings = "complete"
|
||||
|
||||
loader = SimpleNamespace(name="fixture", exec_module=initialize)
|
||||
namespace = definitions("runtime.py", {"_legacy"}, dict(
|
||||
_LEGACY=None, _LEGACY_LOCK=lock, PREFIX=Path("/fixture"),
|
||||
importlib=SimpleNamespace(
|
||||
machinery=SimpleNamespace(SourceFileLoader=lambda *args: loader),
|
||||
util=SimpleNamespace(spec_from_loader=lambda *args: object(),
|
||||
module_from_spec=lambda spec: module))))
|
||||
with ThreadPoolExecutor(max_workers=2) as pool:
|
||||
first = pool.submit(namespace["_legacy"])
|
||||
try:
|
||||
self.assertTrue(started.wait(1))
|
||||
self.assertIsNone(namespace["_LEGACY"])
|
||||
second = pool.submit(lock.watched, namespace["_legacy"])
|
||||
self.assertTrue(lock.waiting.wait(1))
|
||||
self.assertFalse(second.done())
|
||||
finally:
|
||||
release.set()
|
||||
self.assertIs(first.result(1), module)
|
||||
self.assertIs(second.result(1), module)
|
||||
self.assertEqual(module.settings, "complete")
|
||||
self.assertEqual(loads, [module])
|
||||
|
||||
def test_runtime_failed_initialization_remains_retryable(self):
|
||||
attempts = []
|
||||
module = SimpleNamespace()
|
||||
|
||||
def initialize(target):
|
||||
attempts.append(target)
|
||||
if len(attempts) == 1:
|
||||
raise RuntimeError("fixture failure")
|
||||
target.settings = "complete"
|
||||
|
||||
loader = SimpleNamespace(name="fixture", exec_module=initialize)
|
||||
namespace = definitions("runtime.py", {"_legacy"}, dict(
|
||||
_LEGACY=None, _LEGACY_LOCK=threading.RLock(), PREFIX=Path("/fixture"),
|
||||
importlib=SimpleNamespace(
|
||||
machinery=SimpleNamespace(SourceFileLoader=lambda *args: loader),
|
||||
util=SimpleNamespace(spec_from_loader=lambda *args: object(),
|
||||
module_from_spec=lambda spec: module))))
|
||||
with self.assertRaisesRegex(RuntimeError, "fixture failure"):
|
||||
namespace["_legacy"]()
|
||||
self.assertIsNone(namespace["_LEGACY"])
|
||||
self.assertIs(namespace["_legacy"](), module)
|
||||
self.assertEqual(len(attempts), 2)
|
||||
self.assertEqual(module.settings, "complete")
|
||||
|
||||
def test_fleet_bootstrap_serializes_shared_path_check_remove_and_append(self):
|
||||
started, release = threading.Event(), threading.Event()
|
||||
lock = ObservedRLock()
|
||||
broker, fleet = Path("/fixture/broker"), Path("/fixture/fleet")
|
||||
|
||||
class Paths(list):
|
||||
def remove(self, value):
|
||||
if not started.is_set():
|
||||
started.set()
|
||||
if not release.wait(2):
|
||||
raise AssertionError("fixture path update was not released")
|
||||
super().remove(value)
|
||||
|
||||
paths = Paths([str(broker), str(fleet)])
|
||||
package = SimpleNamespace(__file__=str(broker / "framebridge_linux_chatgpt/__init__.py"),
|
||||
__path__=[])
|
||||
namespace = definitions("fleet.py", {"bootstrap_original"}, dict(
|
||||
_BOOTSTRAP_LOCK=lock, original_paths=lambda: (broker, fleet), Path=Path,
|
||||
sys=SimpleNamespace(path=paths),
|
||||
_require=lambda valid, code: self.assertTrue(valid, code)))
|
||||
with mock.patch.dict(sys.modules, {"framebridge_linux_chatgpt": package}):
|
||||
with ThreadPoolExecutor(max_workers=2) as pool:
|
||||
first = pool.submit(namespace["bootstrap_original"])
|
||||
try:
|
||||
self.assertTrue(started.wait(1))
|
||||
second = pool.submit(lock.watched, namespace["bootstrap_original"])
|
||||
self.assertTrue(lock.waiting.wait(1))
|
||||
self.assertFalse(second.done())
|
||||
finally:
|
||||
release.set()
|
||||
first.result(1)
|
||||
second.result(1)
|
||||
self.assertEqual(paths.count(str(broker)), 1)
|
||||
self.assertEqual(paths.count(str(fleet)), 1)
|
||||
self.assertCountEqual(package.__path__, [str(root / "framebridge_linux_chatgpt")
|
||||
for root in (broker, fleet)])
|
||||
|
||||
|
||||
def broker_namespace(self, config):
|
||||
def validate(url):
|
||||
if url == "invalid-fixture":
|
||||
raise ValueError("fixture endpoint rejected")
|
||||
return url
|
||||
|
||||
return definitions("broker.py", {"configure", "_settings"}, dict(
|
||||
_state=None, _url=None, _control_path=None,
|
||||
_SETTINGS_LOCK=ObservedRLock(), Path=Path,
|
||||
state_path=lambda state=None: Path(state or "/fixture/default"),
|
||||
_config=config, validate_endpoint=validate))
|
||||
|
||||
def test_broker_initial_settings_are_published_as_one_complete_profile(self):
|
||||
started, release = threading.Event(), threading.Event()
|
||||
loads = []
|
||||
|
||||
def config(state):
|
||||
loads.append(state)
|
||||
started.set()
|
||||
if not release.wait(2):
|
||||
raise AssertionError("fixture config was not released")
|
||||
return {"ports": {"broker": 23456}}
|
||||
|
||||
namespace = self.broker_namespace(config)
|
||||
lock = namespace["_SETTINGS_LOCK"]
|
||||
with ThreadPoolExecutor(max_workers=2) as pool:
|
||||
first = pool.submit(namespace["_settings"])
|
||||
try:
|
||||
self.assertTrue(started.wait(1))
|
||||
self.assertEqual(tuple(namespace[key] for key in ("_state", "_url", "_control_path")),
|
||||
(None, None, None))
|
||||
second = pool.submit(lock.watched, namespace["_settings"])
|
||||
self.assertTrue(lock.waiting.wait(1))
|
||||
self.assertFalse(second.done())
|
||||
finally:
|
||||
release.set()
|
||||
expected = (Path("/fixture/default"), "http://127.0.0.1:23456",
|
||||
Path("/fixture/default/broker/control.key"))
|
||||
self.assertEqual(first.result(1), expected)
|
||||
self.assertEqual(second.result(1), expected)
|
||||
self.assertEqual(loads, [Path("/fixture/default")])
|
||||
|
||||
def test_broker_explicit_profile_change_cannot_expose_mixed_endpoint_or_key_path(self):
|
||||
started, release = threading.Event(), threading.Event()
|
||||
|
||||
def config(state):
|
||||
if state == Path("/fixture/next"):
|
||||
started.set()
|
||||
if not release.wait(2):
|
||||
raise AssertionError("fixture config was not released")
|
||||
return {"ports": {"broker": 23456}}
|
||||
|
||||
namespace = self.broker_namespace(config)
|
||||
previous = namespace["_settings"]()
|
||||
lock = namespace["_SETTINGS_LOCK"]
|
||||
with ThreadPoolExecutor(max_workers=2) as pool:
|
||||
first = pool.submit(namespace["configure"], "https://fixture.example:24443",
|
||||
"/fixture/next", "/fixture/next/explicit-control.key")
|
||||
try:
|
||||
self.assertTrue(started.wait(1))
|
||||
self.assertEqual(tuple(namespace[key] for key in ("_state", "_url", "_control_path")),
|
||||
previous)
|
||||
second = pool.submit(lock.watched, namespace["_settings"])
|
||||
self.assertTrue(lock.waiting.wait(1))
|
||||
self.assertFalse(second.done())
|
||||
finally:
|
||||
release.set()
|
||||
first.result(1)
|
||||
self.assertEqual(second.result(1), (Path("/fixture/next"),
|
||||
"https://fixture.example:24443", Path("/fixture/next/explicit-control.key")))
|
||||
|
||||
def test_broker_failed_validation_publishes_nothing_and_can_retry(self):
|
||||
namespace = self.broker_namespace(lambda state: {})
|
||||
with self.assertRaisesRegex(ValueError, "fixture endpoint rejected"):
|
||||
namespace["configure"]("invalid-fixture", "/fixture/rejected")
|
||||
self.assertEqual(tuple(namespace[key] for key in ("_state", "_url", "_control_path")),
|
||||
(None, None, None))
|
||||
previous = namespace["_settings"]()
|
||||
with self.assertRaisesRegex(ValueError, "fixture endpoint rejected"):
|
||||
namespace["configure"]("invalid-fixture", "/fixture/rejected")
|
||||
self.assertEqual(namespace["_settings"](), previous)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,341 @@
|
||||
"""Socket fault injection against the real FBPAIR broker and enrolled identities."""
|
||||
import errno
|
||||
import hashlib
|
||||
import selectors
|
||||
import shlex
|
||||
import signal
|
||||
import socket
|
||||
import ssl
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
import urllib.error
|
||||
from unittest import mock
|
||||
|
||||
from framebridge_expanded import broker, node
|
||||
import test_fleet_broker_node as fixtures
|
||||
|
||||
|
||||
class StopLoop(BaseException):
|
||||
pass
|
||||
|
||||
|
||||
class NodeRecoveryTests(unittest.TestCase):
|
||||
setUp = fixtures.BrokerNodeTests.setUp
|
||||
tearDown = fixtures.BrokerNodeTests.tearDown
|
||||
make_node = fixtures.BrokerNodeTests.make_node
|
||||
pair = fixtures.BrokerNodeTests.pair
|
||||
|
||||
@staticmethod
|
||||
def disconnect(handler):
|
||||
handler.close_connection = True
|
||||
handler.connection.shutdown(socket.SHUT_RDWR)
|
||||
handler.connection.close()
|
||||
|
||||
def test_transient_poll_disconnects_backoff_cap_and_inventory_recovery(self):
|
||||
state, ident = self.pair("recover-node")
|
||||
node.run(state, once=True)
|
||||
store = self.server.expanded_store
|
||||
store.node_last_seen[ident] = time.time() - 10
|
||||
self.assertFalse(broker.list_confirmed_nodes()[0]["ready"])
|
||||
identity_before = (state / "node/identity.json").read_bytes()
|
||||
original = self.server.RequestHandlerClass.do_POST
|
||||
real_post = broker.post_json
|
||||
attempts, waits, envelopes = [], [], []
|
||||
recovered = []
|
||||
|
||||
def fault(handler):
|
||||
if handler.path.endswith("/node/poll"):
|
||||
attempts.append(1)
|
||||
if len(attempts) <= 8:
|
||||
self.disconnect(handler)
|
||||
return
|
||||
return original(handler)
|
||||
|
||||
def post(url, envelope, **kwargs):
|
||||
if url.endswith("/node/poll"):
|
||||
envelopes.append(envelope)
|
||||
result = real_post(url, envelope, **kwargs)
|
||||
if url.endswith("/node/poll"):
|
||||
recovered.append(True)
|
||||
return result
|
||||
|
||||
def pause(seconds):
|
||||
if recovered:
|
||||
raise StopLoop
|
||||
self.assertFalse(store.node_runtime_status(ident)["ready"])
|
||||
waits.append(seconds)
|
||||
|
||||
with mock.patch.object(self.server.RequestHandlerClass, "do_POST", fault), \
|
||||
mock.patch.object(broker, "post_json", side_effect=post), \
|
||||
mock.patch.object(node, "time", wraps=time) as clock, \
|
||||
self.assertLogs(node._LOG, level="WARNING") as logs:
|
||||
clock.sleep.side_effect = pause
|
||||
with self.assertRaises(StopLoop):
|
||||
node.run(state)
|
||||
self.assertEqual(waits, [1, 2, 4, 8, 16, 30, 30, 30])
|
||||
self.assertEqual(len(attempts), 9)
|
||||
self.assertEqual(len({e["nonce"] for e in envelopes}), 9)
|
||||
self.assertTrue(store.node_runtime_status(ident)["ready"])
|
||||
self.assertEqual((state / "node/identity.json").read_bytes(), identity_before)
|
||||
self.assertTrue(any("NODE_BROKER_RECOVERED" in line for line in logs.output))
|
||||
self.assertTrue(all("signature" not in line and "payload" not in line for line in logs.output))
|
||||
|
||||
def test_result_disconnect_before_and_after_acceptance_never_reexecutes(self):
|
||||
state, ident = self.pair("ack-node")
|
||||
marker = self.root / "executed.txt"
|
||||
real_post = broker.post_json
|
||||
original_post = self.server.RequestHandlerClass.do_POST
|
||||
original_send = self.server.RequestHandlerClass.send_json
|
||||
for after_acceptance in (False, True):
|
||||
with self.subTest(after_acceptance=after_acceptance):
|
||||
marker.write_text("")
|
||||
script = "printf x >> " + shlex.quote(str(marker))
|
||||
payload = {"authorized": True, "job_id": "recovery-command", "timeout_s": 10,
|
||||
"script": script, "sha256": hashlib.sha256(script.encode()).hexdigest()}
|
||||
queued, error = self.server.expanded_store.queue_admin(ident, "admin.exec", payload)
|
||||
self.assertIsNone(error)
|
||||
posted, dispatched, waits, dropped, delivered = [], [], [], [], []
|
||||
|
||||
def execute(value):
|
||||
dispatched.append(value["job_id"])
|
||||
subprocess.run(["/bin/sh", "-c", value["script"]], check=True)
|
||||
return {"ok": True, "exit_code": "0"}
|
||||
|
||||
def fault_post(handler):
|
||||
if not after_acceptance and handler.path.endswith("/node/result") and not dropped:
|
||||
dropped.append(True)
|
||||
self.disconnect(handler)
|
||||
return
|
||||
return original_post(handler)
|
||||
|
||||
def fault_send(handler, code, value):
|
||||
if after_acceptance and handler.path.endswith("/node/result") and not dropped:
|
||||
dropped.append(True)
|
||||
self.disconnect(handler)
|
||||
return
|
||||
return original_send(handler, code, value)
|
||||
|
||||
def post(url, envelope, **kwargs):
|
||||
if url.endswith("/node/result"):
|
||||
posted.append(envelope)
|
||||
result = real_post(url, envelope, **kwargs)
|
||||
if url.endswith("/node/result"):
|
||||
delivered.append(True)
|
||||
return result
|
||||
|
||||
def pause(seconds):
|
||||
if delivered:
|
||||
raise StopLoop
|
||||
waits.append(seconds)
|
||||
|
||||
with mock.patch.object(self.server.RequestHandlerClass, "do_POST", fault_post), \
|
||||
mock.patch.object(self.server.RequestHandlerClass, "send_json", fault_send), \
|
||||
mock.patch.object(broker, "post_json", side_effect=post), \
|
||||
mock.patch("framebridge_expanded.runtime.adapter_request", side_effect=execute), \
|
||||
mock.patch.object(node, "time", wraps=time) as clock, \
|
||||
self.assertLogs(node._LOG, level="WARNING"):
|
||||
clock.sleep.side_effect = pause
|
||||
with self.assertRaises(StopLoop):
|
||||
node.run(state)
|
||||
self.assertEqual(marker.read_text(), "x")
|
||||
self.assertEqual(dispatched, ["recovery-command"])
|
||||
self.assertEqual(waits, [1])
|
||||
self.assertEqual(len(posted), 2)
|
||||
self.assertNotEqual(posted[0]["nonce"], posted[1]["nonce"])
|
||||
self.assertEqual(posted[0]["payload"], posted[1]["payload"])
|
||||
saved = self.server.expanded_store.get_result(queued["request_id"])
|
||||
self.assertEqual(saved["state"], "complete")
|
||||
self.assertTrue(saved["result"]["ok"])
|
||||
|
||||
def test_broker_restart_after_execution_is_terminal_without_replay(self):
|
||||
state, ident = self.pair("orphan-node")
|
||||
marker = self.root / "orphan-effect.txt"
|
||||
script = "printf x >> " + shlex.quote(str(marker))
|
||||
payload = {"authorized": True, "job_id": "orphan-command", "timeout_s": 10,
|
||||
"script": script, "sha256": hashlib.sha256(script.encode()).hexdigest()}
|
||||
queued, error = self.server.expanded_store.queue_admin(ident, "admin.exec", payload)
|
||||
self.assertIsNone(error)
|
||||
real_post = broker.post_json
|
||||
result_attempts = []
|
||||
|
||||
def execute(value):
|
||||
subprocess.run(["/bin/sh", "-c", value["script"]], check=True)
|
||||
return {"ok": True, "exit_code": "0"}
|
||||
|
||||
def post(url, envelope, **kwargs):
|
||||
if url.endswith("/node/result"):
|
||||
result_attempts.append(1)
|
||||
self.assertEqual(marker.read_text(), "x")
|
||||
self.server.shutdown()
|
||||
self.server.server_close()
|
||||
self.thread.join(2)
|
||||
self.server = broker.make_server(port=self.port, state_dir=self.central)
|
||||
self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
|
||||
self.thread.start()
|
||||
self.assertIsNone(self.server.expanded_store.get_result(queued["request_id"]))
|
||||
return real_post(url, envelope, **kwargs)
|
||||
|
||||
with mock.patch.object(broker, "post_json", side_effect=post), \
|
||||
mock.patch("framebridge_expanded.runtime.adapter_request", side_effect=execute) as adapter, \
|
||||
mock.patch.object(node, "time", wraps=time) as clock:
|
||||
with self.assertRaisesRegex(broker.BrokerError, "^BROKER_HTTP_REJECTED_400$"):
|
||||
node.run(state)
|
||||
adapter.assert_called_once()
|
||||
clock.sleep.assert_not_called()
|
||||
self.assertEqual(result_attempts, [1])
|
||||
self.assertEqual(marker.read_text(), "x")
|
||||
# A later operator restart may poll, but must never reconstruct/replay
|
||||
# the orphaned command from its lost acknowledgement.
|
||||
self.assertTrue(node.run(state, once=True)["idle"])
|
||||
self.assertEqual(marker.read_text(), "x")
|
||||
|
||||
def test_gateway_failures_recover_without_replaying_signed_envelope(self):
|
||||
state, _ = self.pair("gateway-node")
|
||||
original = self.server.RequestHandlerClass.do_POST
|
||||
codes, waits, successes = [502, 503, 504], [], []
|
||||
|
||||
def gateway(handler):
|
||||
if codes:
|
||||
handler.send_json(codes.pop(0), {"ok": False})
|
||||
return
|
||||
successes.append(True)
|
||||
return original(handler)
|
||||
|
||||
def pause(seconds):
|
||||
if successes:
|
||||
raise StopLoop
|
||||
waits.append(seconds)
|
||||
|
||||
with mock.patch.object(self.server.RequestHandlerClass, "do_POST", gateway), \
|
||||
mock.patch.object(node, "time", wraps=time) as clock, \
|
||||
self.assertLogs(node._LOG, level="WARNING"):
|
||||
clock.sleep.side_effect = pause
|
||||
with self.assertRaises(StopLoop):
|
||||
node.run(state)
|
||||
self.assertEqual(waits, [1, 2, 4])
|
||||
|
||||
def test_auth_and_invalid_json_fail_without_retry(self):
|
||||
state, _ = self.pair("reject-node")
|
||||
for code, body in [(403, b'{"ok":false}'), (500, b'{"ok":false}'),
|
||||
(200, b'not-json'), (200, b'{"ok":false}')]:
|
||||
with self.subTest(code=code, body=body):
|
||||
attempts = []
|
||||
|
||||
def reject(handler):
|
||||
attempts.append(1)
|
||||
handler.send_response(code)
|
||||
handler.send_header("Content-Length", str(len(body)))
|
||||
handler.end_headers()
|
||||
handler.wfile.write(body)
|
||||
|
||||
with mock.patch.object(self.server.RequestHandlerClass, "do_POST", reject), \
|
||||
mock.patch.object(node, "time", wraps=time) as clock:
|
||||
with self.assertRaises((broker.BrokerError, node.NodeError)):
|
||||
node.run(state)
|
||||
clock.sleep.assert_not_called()
|
||||
self.assertEqual(len(attempts), 1)
|
||||
|
||||
def test_tls_and_unknown_transport_causes_are_terminal(self):
|
||||
state, _ = self.pair("tls-node")
|
||||
for cause in [urllib.error.URLError(ssl.SSLCertVerificationError("untrusted test CA")),
|
||||
urllib.error.URLError(socket.gaierror(socket.EAI_NONAME, "invalid host")),
|
||||
OSError(errno.EACCES, "denied"), ValueError("invalid response")]:
|
||||
with self.subTest(cause=type(cause).__name__), \
|
||||
mock.patch("urllib.request.OpenerDirector.open", side_effect=cause) as opened, \
|
||||
mock.patch.object(node, "time", wraps=time) as clock:
|
||||
with self.assertRaises(broker.BrokerError):
|
||||
node.run(state)
|
||||
self.assertEqual(opened.call_count, 1)
|
||||
clock.sleep.assert_not_called()
|
||||
|
||||
def test_once_is_fail_fast_and_retry_wait_propagates_cancellation(self):
|
||||
state, _ = self.pair("cancel-node")
|
||||
attempts = []
|
||||
|
||||
def fault(handler):
|
||||
attempts.append(1)
|
||||
self.disconnect(handler)
|
||||
|
||||
with mock.patch.object(self.server.RequestHandlerClass, "do_POST", fault), \
|
||||
mock.patch.object(node, "time", wraps=time) as clock:
|
||||
with self.assertRaises(broker.BrokerError):
|
||||
node.run(state, once=True)
|
||||
clock.sleep.assert_not_called()
|
||||
clock.sleep.side_effect = KeyboardInterrupt
|
||||
with self.assertLogs(node._LOG, level="WARNING"), self.assertRaises(KeyboardInterrupt):
|
||||
node.run(state)
|
||||
self.assertEqual(len(attempts), 2)
|
||||
clock.sleep.assert_called_once_with(1)
|
||||
|
||||
def test_real_broker_restart_recovers_same_node_process_and_identity(self):
|
||||
state, ident = self.pair("restart-node")
|
||||
identity_before = (state / "node/identity.json").read_bytes()
|
||||
code = "from framebridge_expanded.node import run; run(state_dir=" + repr(str(state)) + ")"
|
||||
child = subprocess.Popen([sys.executable, "-u", "-c", code],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
|
||||
|
||||
def wait_ready():
|
||||
deadline = time.monotonic() + 8
|
||||
while time.monotonic() < deadline:
|
||||
self.assertIsNone(child.poll(), "node exited during recovery")
|
||||
if self.server.expanded_store.node_runtime_status(ident)["ready"]:
|
||||
return
|
||||
threading.Event().wait(0.02)
|
||||
self.fail("node heartbeat did not recover")
|
||||
|
||||
try:
|
||||
wait_ready()
|
||||
self.server.shutdown()
|
||||
self.server.server_close()
|
||||
self.thread.join(2)
|
||||
with selectors.DefaultSelector() as selector:
|
||||
selector.register(child.stderr, selectors.EVENT_READ)
|
||||
self.assertTrue(selector.select(10), "child did not report broker interruption")
|
||||
self.assertIn("NODE_BROKER_RETRY", child.stderr.readline())
|
||||
self.assertIsNone(child.poll())
|
||||
self.server = broker.make_server(port=self.port, state_dir=self.central)
|
||||
self.assertFalse(self.server.expanded_store.node_runtime_status(ident)["ready"])
|
||||
self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
|
||||
self.thread.start()
|
||||
wait_ready()
|
||||
self.assertEqual((state / "node/identity.json").read_bytes(), identity_before)
|
||||
self.assertEqual(broker.list_confirmed_nodes()[0]["node_id"], ident)
|
||||
finally:
|
||||
child.terminate()
|
||||
try:
|
||||
child.wait(timeout=3)
|
||||
finally:
|
||||
if child.poll() is None:
|
||||
child.kill()
|
||||
child.wait(timeout=3)
|
||||
child.stderr.close()
|
||||
|
||||
def test_sigterm_stops_real_process_during_retry_wait(self):
|
||||
state, _ = self.pair("stop-node")
|
||||
self.server.shutdown()
|
||||
self.server.server_close()
|
||||
self.thread.join(2)
|
||||
code = "from framebridge_expanded.node import run; run(state_dir=" + repr(str(state)) + ")"
|
||||
child = subprocess.Popen([sys.executable, "-u", "-c", code],
|
||||
stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
|
||||
try:
|
||||
with selectors.DefaultSelector() as selector:
|
||||
selector.register(child.stderr, selectors.EVENT_READ)
|
||||
self.assertTrue(selector.select(10), "child did not report retry")
|
||||
self.assertIn("NODE_BROKER_RETRY", child.stderr.readline())
|
||||
started = time.monotonic()
|
||||
child.terminate()
|
||||
self.assertEqual(child.wait(timeout=3), -signal.SIGTERM)
|
||||
self.assertLess(time.monotonic() - started, 3)
|
||||
finally:
|
||||
if child.poll() is None:
|
||||
child.kill()
|
||||
child.wait(timeout=3)
|
||||
child.stderr.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,91 @@
|
||||
"""Private package versions retain the fixed privileged-helper rollback floor.
|
||||
|
||||
Only a temporary maintainer script is run. Cleanup Python is replaced with a
|
||||
recording stub; no installed helper, profile or package transaction is touched.
|
||||
"""
|
||||
import ast
|
||||
import json
|
||||
from pathlib import Path
|
||||
import shlex
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
WORK = Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0, str(WORK / "packaging"))
|
||||
import build_expanded
|
||||
|
||||
PRIVATE_VERSION = "0.3.2+private.20261028.1"
|
||||
|
||||
|
||||
class PrivateCandidatePrermTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.temp.cleanup)
|
||||
self.root = Path(self.temp.name)
|
||||
self.control = self.root / "DEBIAN"
|
||||
self.control.mkdir()
|
||||
self.log = self.root / "cleanup-calls.jsonl"
|
||||
self.python_stub = self.root / "record-cleanup"
|
||||
self.python_stub.write_text(
|
||||
"#!/usr/bin/python3\nimport json,sys\nfrom pathlib import Path\n"
|
||||
+ "with Path(" + repr(str(self.log)) + ").open('a') as stream:\n"
|
||||
+ " stream.write(json.dumps(sys.argv[1:])+'\\n')\n")
|
||||
self.python_stub.chmod(0o755)
|
||||
|
||||
def candidate_script(self, package_version=PRIVATE_VERSION):
|
||||
# Execute the real production write expression with an arbitrary package
|
||||
# VERSION, rather than reconstructing its floor selection in the test.
|
||||
tree = ast.parse(Path(build_expanded.__file__).read_text())
|
||||
main = next(node for node in tree.body
|
||||
if isinstance(node, ast.FunctionDef) and node.name == "main")
|
||||
writers = [node for node in ast.walk(main)
|
||||
if isinstance(node, ast.Expr) and isinstance(node.value, ast.Call)
|
||||
and isinstance(node.value.func, ast.Attribute)
|
||||
and node.value.func.attr == "write_text"
|
||||
and node.value.args and isinstance(node.value.args[0], ast.Call)
|
||||
and isinstance(node.value.args[0].func, ast.Name)
|
||||
and node.value.args[0].func.id == "maintainer_prerm"]
|
||||
self.assertEqual(len(writers), 1)
|
||||
namespace = dict(control=self.control, VERSION=package_version,
|
||||
PRIVILEGED_HELPER_FLOOR=build_expanded.PRIVILEGED_HELPER_FLOOR,
|
||||
maintainer_prerm=build_expanded.maintainer_prerm)
|
||||
exec(compile(ast.Module(body=writers, type_ignores=[]), build_expanded.__file__, "exec"), namespace)
|
||||
script = self.control / "prerm"
|
||||
text = script.read_text()
|
||||
self.assertIn('"$2" lt 0.3.2;', text)
|
||||
self.assertNotIn(package_version, text)
|
||||
# The real dpkg is used only for its read-only version validator and
|
||||
# comparator. Every privileged cleanup command is a recording stub.
|
||||
script.write_text(text.replace("/usr/bin/python3", shlex.quote(str(self.python_stub))))
|
||||
return script
|
||||
|
||||
def upgrade_target(self, target, package_version=PRIVATE_VERSION):
|
||||
script = self.candidate_script(package_version)
|
||||
result = subprocess.run(["/bin/sh", str(script), "upgrade", target],
|
||||
capture_output=True, text=True, timeout=5)
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
return [json.loads(line) for line in self.log.read_text().splitlines()] if self.log.exists() else []
|
||||
|
||||
def test_private_to_official_032_preserves_helper_grants(self):
|
||||
self.assertEqual(self.upgrade_target("0.3.2"), [])
|
||||
|
||||
def test_target_below_032_revokes_helper_grants(self):
|
||||
calls = self.upgrade_target("0.3.1")
|
||||
self.assertEqual(len(calls), 1)
|
||||
self.assertEqual(calls[0][:3], ["-I", "-S", "-c"])
|
||||
self.assertIn('["purge_credential_helpers"]()', calls[0][3])
|
||||
|
||||
def test_same_private_version_preserves_helper_grants(self):
|
||||
self.assertEqual(self.upgrade_target(PRIVATE_VERSION), [])
|
||||
|
||||
def test_forward_upgrade_preserves_helper_grants(self):
|
||||
self.assertEqual(self.upgrade_target("0.3.3"), [])
|
||||
|
||||
def test_arbitrary_future_package_version_keeps_the_032_floor(self):
|
||||
self.assertEqual(self.upgrade_target("0.3.2", package_version="9.9.9+private.1"), [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -3,6 +3,10 @@
|
||||
The bounded executor is a fixture. This is not VM, GUI, TTY or input acceptance.
|
||||
"""
|
||||
import ast
|
||||
import base64
|
||||
import json
|
||||
import math
|
||||
import shlex
|
||||
from pathlib import Path
|
||||
import unittest
|
||||
|
||||
@@ -13,9 +17,26 @@ from test_screen_frontdoor import Executor, Image, MCP, NODE_ID
|
||||
PORTABLE = Path(__file__).resolve().parents[1] / "portable/front_door_mcp.py"
|
||||
|
||||
|
||||
class TypedChunkExecutor(Executor):
|
||||
"""Model the explicitly negotiated typed route without changing legacy fixtures."""
|
||||
def __call__(self, **kwargs):
|
||||
outer = super().__call__(**kwargs)
|
||||
request = json.loads(base64.b64decode(shlex.split(kwargs["script"])[-1]))
|
||||
if request["op"] == "read_chunk" and request.get("chunk_bytes"):
|
||||
size, index = request["chunk_bytes"], request["index"]
|
||||
envelope = json.loads(outer["stdout"])
|
||||
envelope["result"].update(
|
||||
total_chunks=math.ceil(len(self.data) / size),
|
||||
data=base64.b64encode(self.data[index * size:(index + 1) * size]).decode(),
|
||||
more=(index + 1) * size < len(self.data),
|
||||
)
|
||||
outer["stdout"] = json.dumps(envelope)
|
||||
return outer
|
||||
|
||||
|
||||
class PortableScreenRegistrationTests(unittest.TestCase):
|
||||
def registration(self, target="LOCAL", corrupt=None):
|
||||
fixture = Executor(hostname="fb-central-lab" if target == "LOCAL" else "fb-node-lab")
|
||||
fixture = TypedChunkExecutor(hostname="fb-central-lab" if target == "LOCAL" else "fb-node-lab")
|
||||
calls, legacy_calls = [], []
|
||||
|
||||
def expanded(alias, script, authorized, timeout_s=120, job_id="", sha256="", include_stdout=False):
|
||||
@@ -51,6 +72,7 @@ class PortableScreenRegistrationTests(unittest.TestCase):
|
||||
framebridge_exec=legacy,
|
||||
register_framebridge_screen_tools=register_framebridge_screen_tools,
|
||||
resolve_profile_node=lambda node: node)
|
||||
scope["_blocking_tools"] = scope["mcp"] # This fixture calls registrations synchronously.
|
||||
# Load the definition and registration from production source rather
|
||||
# than recreating the adapter inside this test.
|
||||
module = ast.Module(body=definitions + registrations, type_ignores=[])
|
||||
|
||||
@@ -0,0 +1,431 @@
|
||||
"""Deterministic protocol fixtures only: never observe or control a desktop."""
|
||||
import base64
|
||||
import concurrent.futures
|
||||
import contextlib
|
||||
import hashlib
|
||||
import gzip
|
||||
import io
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
from pathlib import Path
|
||||
import shlex
|
||||
import sys
|
||||
import types
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from framebridge_expanded.real_screen import screen_rpc as rpc
|
||||
from framebridge_expanded.real_screen import screen_transport as transport
|
||||
from framebridge_expanded.real_screen.screen_frontdoor import (
|
||||
CHUNK_BYTES, EXPANDED_CHUNK_BYTES, EXPANDED_STDOUT_BYTES,
|
||||
MAX_STDOUT_BYTES, RPC_COMMAND, ScreenClient, ScreenError,
|
||||
)
|
||||
from test_screen_frontdoor import CAPTURE_ID, NODE_ID, Image, png, Executor
|
||||
|
||||
|
||||
class Clock:
|
||||
now = 100.0
|
||||
|
||||
def __call__(self):
|
||||
return self.now
|
||||
|
||||
|
||||
class ProtocolFixture:
|
||||
"""Saved generated PNG fixture through actual producer, RPC and typed transport.
|
||||
|
||||
The private records model a previously captured artifact for protocol tests;
|
||||
they are generated test pixels and do not certify any real screen observation.
|
||||
"""
|
||||
def __init__(self, directory, width=384, height=256):
|
||||
self.directory = directory
|
||||
self.data = png(width, height)
|
||||
self.manifest = {
|
||||
'capture_id': CAPTURE_ID, 'node': 'test-fixture-host', 'uid': os.getuid(),
|
||||
'source': 'linux-framebuffer', 'screen_type': 'tty', 'active_vt': 'tty1',
|
||||
'bytes': len(self.data), 'sha256': hashlib.sha256(self.data).hexdigest(),
|
||||
'width': width, 'height': height, 'synthetic': False,
|
||||
'pixel_source_verified': True, 'full_screen': True, 'authentic_pixels': True,
|
||||
'session_id': 'screen_fixture', 'requested_session_visible': False,
|
||||
'control_available': True, '_action_consumed': False,
|
||||
'binding': {'window_id': 'fixture-window', 'geometry': [0, 0, width, height]},
|
||||
}
|
||||
path = directory / (CAPTURE_ID + '.png')
|
||||
path.write_bytes(self.data)
|
||||
path.chmod(0o600)
|
||||
for suffix in ('.json', '.bound.json'):
|
||||
rpc.write_record(directory / (CAPTURE_ID + suffix), self.manifest)
|
||||
self.calls, self.reply_sizes = [], []
|
||||
self.rpc, self.transport = rpc, transport
|
||||
self.variant = None
|
||||
self.clock, self.latency = None, 0
|
||||
|
||||
def __call__(self, **kwargs):
|
||||
request = json.loads(base64.b64decode(shlex.split(kwargs['script'])[-1]))
|
||||
self.calls.append((request, kwargs['timeout_s']))
|
||||
if self.clock is not None:
|
||||
self.clock.now += self.latency
|
||||
payload = {key: kwargs[key] for key in (
|
||||
'script', 'authorized', 'timeout_s', 'job_id', 'sha256', 'include_stdout')}
|
||||
payload.update(version=1, mode='EXEC', request_type='screen_rpc')
|
||||
|
||||
def runner(argument, timeout_s):
|
||||
# No real observe/input functions are invoked by this fixture executor.
|
||||
if request['op'] in ('key', 'text'):
|
||||
manifest = self.rpc.public_manifest(self.manifest)
|
||||
manifest['action'] = {'kind': request['op'], 'previous_capture_id': request['capture_id'],
|
||||
'result': {'performed': True}}
|
||||
envelope = {'ok': True, 'result': manifest}
|
||||
else:
|
||||
envelope = self.rpc.dispatch(request, self.directory)
|
||||
if request['op'] == 'read_chunk' and self.variant == 'corrupt':
|
||||
encoded = envelope['result']['data']
|
||||
data = base64.b64decode(encoded)
|
||||
envelope['result']['data'] = base64.b64encode(bytes([data[0] ^ 1]) + data[1:]).decode()
|
||||
return 0, self.rpc.encode(envelope) + b'\n', b''
|
||||
|
||||
outer = self.transport.handle_request(payload, runner=runner)
|
||||
outer.update(target=kwargs['target'], node_id=NODE_ID,
|
||||
mcp_job_id=kwargs['job_id'], mcp_script_sha256=kwargs['sha256'])
|
||||
self.reply_sizes.append(len(outer['stdout'].encode()))
|
||||
if self.variant == 'truncated' and request['op'] == 'read_chunk':
|
||||
outer['stdout_truncated'] = True
|
||||
return outer
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def frozen_receiver(testcase):
|
||||
"""Load the hashed base receiver fixture, including in exported source archives.
|
||||
|
||||
The capsule contains only the five original real-screen modules from the
|
||||
stated immutable revision. Its source hashes and outer digest are verified;
|
||||
no Git checkout, installed receiver, or live desktop is used by this test.
|
||||
"""
|
||||
revision = 'f6b14c52109e0d55b919f179a38c8a4c8f8868f9'
|
||||
path = Path(__file__).resolve().parent / 'fixtures/frozen_screen_receiver_f6b14c52.json.gz'
|
||||
raw = path.read_bytes()
|
||||
testcase.assertEqual(hashlib.sha256(raw).hexdigest(), '47670e1325c0f77fabeea9c69c09923c3fd0ed2463a8300a4940bb5e4bafffc5')
|
||||
capsule = json.loads(gzip.decompress(raw))
|
||||
testcase.assertEqual(capsule['revision'], revision)
|
||||
names = ('screen_frontdoor', 'screen_capture', 'screen_control', 'screen_rpc', 'screen_transport')
|
||||
testcase.assertEqual(set(capsule['files']), set(names))
|
||||
package_name = '_framebridge_frozen_screen_fixture'
|
||||
package = types.ModuleType(package_name)
|
||||
package.__path__ = []
|
||||
with patch.dict(sys.modules, {package_name: package}):
|
||||
modules = {}
|
||||
for name in names:
|
||||
entry = capsule['files'][name]
|
||||
source = entry['source']
|
||||
testcase.assertEqual(hashlib.sha256(source.encode()).hexdigest(), entry['sha256'])
|
||||
full_name = package_name + '.' + name
|
||||
module = types.ModuleType(full_name)
|
||||
module.__package__ = package_name
|
||||
sys.modules[full_name] = module
|
||||
exec(compile(source, revision + ':' + name + '.py', 'exec'), module.__dict__)
|
||||
modules[name] = module
|
||||
yield modules['screen_rpc'], modules['screen_transport']
|
||||
|
||||
|
||||
class ScreenTransferBudgetTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.temp.cleanup)
|
||||
self.directory = Path(self.temp.name)
|
||||
self.fixture = ProtocolFixture(self.directory)
|
||||
|
||||
def client(self, **kwargs):
|
||||
return ScreenClient(self.fixture, Image, **kwargs)
|
||||
|
||||
def read(self, client):
|
||||
return client.call('read_capture', 'GERAMEG', True, capture_id=CAPTURE_ID)
|
||||
|
||||
def test_expanded_complete_original_through_producer_rpc_and_transport(self):
|
||||
result = self.read(self.client(stdout_limit_bytes=EXPANDED_STDOUT_BYTES))
|
||||
self.assertEqual(result[1].data, self.fixture.data)
|
||||
self.assertEqual(hashlib.sha256(result[1].data).hexdigest(), self.fixture.manifest['sha256'])
|
||||
self.assertEqual(len(self.fixture.calls), 1 + math.ceil(len(self.fixture.data) / EXPANDED_CHUNK_BYTES))
|
||||
self.assertTrue(all(r['chunk_bytes'] == EXPANDED_CHUNK_BYTES for r, _ in self.fixture.calls[1:]))
|
||||
self.assertLessEqual(max(self.fixture.reply_sizes), EXPANDED_STDOUT_BYTES)
|
||||
self.assertGreater(max(self.fixture.reply_sizes), MAX_STDOUT_BYTES)
|
||||
# A full-sized serialized reply is measured, not estimated from base64 alone.
|
||||
self.assertGreater(max(self.fixture.reply_sizes), 60000)
|
||||
self.assertGreater(EXPANDED_STDOUT_BYTES - max(self.fixture.reply_sizes), 4096)
|
||||
for key in ('capture_id', 'node', 'session_id', 'width', 'height', 'binding', 'requested_session_visible'):
|
||||
self.assertEqual(result[0]['result'][key], self.fixture.manifest[key])
|
||||
self.assertEqual(result[0]['node_id'], NODE_ID)
|
||||
legacy_calls = 1 + math.ceil(len(self.fixture.data) / CHUNK_BYTES)
|
||||
self.assertGreater(legacy_calls, 10 * len(self.fixture.calls))
|
||||
|
||||
def test_legacy_omits_extension_and_stays_within_4096_bytes(self):
|
||||
result = self.read(self.client())
|
||||
self.assertEqual(result[1].data, self.fixture.data)
|
||||
self.assertEqual(len(self.fixture.calls), 1 + math.ceil(len(self.fixture.data) / CHUNK_BYTES))
|
||||
self.assertTrue(all('chunk_bytes' not in r for r, _ in self.fixture.calls))
|
||||
self.assertLessEqual(max(self.fixture.reply_sizes), MAX_STDOUT_BYTES)
|
||||
|
||||
def test_new_central_with_actual_frozen_receiver_uses_legacy_schema(self):
|
||||
with frozen_receiver(self) as (frozen_rpc, frozen_transport):
|
||||
self.fixture.rpc, self.fixture.transport = frozen_rpc, frozen_transport
|
||||
self.assertNotIn('max_chunk_bytes', frozen_rpc.public_manifest(self.fixture.manifest))
|
||||
# Only the capture itself is mocked: all frozen protocol and PNG reads
|
||||
# execute the actual base code. Never access a live desktop in tests.
|
||||
with patch.object(frozen_rpc, 'observe', return_value=self.fixture.manifest):
|
||||
for op, arguments in [('observe', {}), ('read_capture', {'capture_id': CAPTURE_ID})]:
|
||||
with self.subTest(op=op):
|
||||
self.fixture.calls.clear()
|
||||
result = self.client(stdout_limit_bytes=EXPANDED_STDOUT_BYTES).call(
|
||||
op, 'GERAMEG', True, **arguments)
|
||||
self.assertEqual(result[1].data, self.fixture.data)
|
||||
self.assertEqual(result[0]['result']['sha256'], self.fixture.manifest['sha256'])
|
||||
self.assertNotIn('max_chunk_bytes', result[0]['result'])
|
||||
self.assertTrue(all('chunk_bytes' not in request for request, _ in self.fixture.calls))
|
||||
self.assertEqual(len(self.fixture.calls), 1 + math.ceil(len(self.fixture.data) / CHUNK_BYTES))
|
||||
self.assertLessEqual(max(self.fixture.reply_sizes), MAX_STDOUT_BYTES)
|
||||
|
||||
def test_capability_is_exact_bounded_int_and_cannot_expand_transport_budget(self):
|
||||
original = self.fixture.rpc.public_manifest
|
||||
for capability in (True, 0, -1, EXPANDED_CHUNK_BYTES + 1, 1.5, '45000'):
|
||||
with self.subTest(capability=capability), patch.object(
|
||||
self.fixture.rpc, 'public_manifest',
|
||||
side_effect=lambda value: dict(original(value), max_chunk_bytes=capability)):
|
||||
self.fixture.calls.clear()
|
||||
result = self.read(self.client(stdout_limit_bytes=EXPANDED_STDOUT_BYTES))
|
||||
self.assertEqual(result['error'], 'SCREEN_TRANSFER_CAPABILITY_INVALID')
|
||||
self.assertEqual(len(self.fixture.calls), 1)
|
||||
self.fixture.calls.clear()
|
||||
result = self.read(self.client())
|
||||
self.assertEqual(result[1].data, self.fixture.data)
|
||||
self.assertTrue(all('chunk_bytes' not in request for request, _ in self.fixture.calls))
|
||||
self.assertLessEqual(max(self.fixture.reply_sizes), MAX_STDOUT_BYTES)
|
||||
|
||||
def test_ambiguous_submitted_input_failures_require_observe_without_claiming_performed(self):
|
||||
variants = ('non_ok', 'truncated', 'malformed', 'binding')
|
||||
for op, arguments in [('key', {'key': 'ENTER'}), ('text', {'text': 'fixture-only'})]:
|
||||
for variant in variants:
|
||||
with self.subTest(op=op, variant=variant):
|
||||
calls = []
|
||||
fixture = Executor()
|
||||
def execute(**kwargs):
|
||||
calls.append(kwargs)
|
||||
outer = fixture(**kwargs)
|
||||
if variant == 'non_ok':
|
||||
return {'ok': False, 'error': 'SCREEN_RPC_FAILED'}
|
||||
if variant == 'truncated':
|
||||
outer['stdout_truncated'] = True
|
||||
elif variant == 'malformed':
|
||||
outer['stdout'] = '{invalid fixture'
|
||||
elif variant == 'binding':
|
||||
outer['mcp_script_sha256'] = '0' * 64
|
||||
return outer
|
||||
result = ScreenClient(execute, Image).call(
|
||||
op, 'GERAMEG', True, capture_id=CAPTURE_ID, **arguments)
|
||||
self.assertFalse(result['ok'])
|
||||
self.assertTrue(result['observe_again'])
|
||||
self.assertEqual(result['capture_id'], CAPTURE_ID)
|
||||
self.assertNotIn('performed', result)
|
||||
self.assertNotIn('capture_consumed', result)
|
||||
self.assertEqual(len(calls), 1)
|
||||
|
||||
def test_unconfirmed_input_image_integrity_failure_does_not_assert_performed(self):
|
||||
for variant in ('dimensions', 'chunk_hash', 'image_hash'):
|
||||
with self.subTest(variant=variant):
|
||||
fixture = Executor(variant)
|
||||
result = ScreenClient(fixture, Image).call(
|
||||
'key', 'GERAMEG', True, capture_id=CAPTURE_ID, key='ENTER')
|
||||
self.assertFalse(result['ok'])
|
||||
self.assertTrue(result['observe_again'])
|
||||
self.assertNotIn('performed', result)
|
||||
self.assertNotIn('capture_consumed', result)
|
||||
self.assertEqual(sum(request['op'] == 'key' for request in fixture.requests), 1)
|
||||
|
||||
def test_generic_rpc_exception_after_consumption_remains_uncertain(self):
|
||||
attempts = []
|
||||
def execute(**kwargs):
|
||||
request = shlex.split(kwargs['script'])[-1]
|
||||
output = io.StringIO()
|
||||
def interrupted(*args, **arguments):
|
||||
attempts.append('fixture-input-attempt')
|
||||
raise OSError('fixture interruption after consuming capture')
|
||||
with patch.object(rpc, 'state_directory', return_value=self.directory), \
|
||||
patch.object(rpc, 'act', side_effect=interrupted), contextlib.redirect_stdout(output):
|
||||
code = rpc.main([request])
|
||||
return {'ok': False, 'exit_code': str(code), 'target': kwargs['target'], 'node_id': NODE_ID,
|
||||
'mcp_job_id': kwargs['job_id'], 'mcp_script_sha256': kwargs['sha256'],
|
||||
'stdout': output.getvalue(), 'stdout_truncated': False}
|
||||
result = ScreenClient(execute, Image).call(
|
||||
'key', 'GERAMEG', True, capture_id=CAPTURE_ID, key='ENTER')
|
||||
self.assertEqual(result['error'], 'SCREEN_REQUEST_FAILED')
|
||||
self.assertTrue(result['observe_again'])
|
||||
self.assertNotIn('performed', result)
|
||||
self.assertNotIn('capture_consumed', result)
|
||||
self.assertEqual(attempts, ['fixture-input-attempt'])
|
||||
self.assertTrue(rpc.load_capture(self.directory, CAPTURE_ID)['_action_consumed'])
|
||||
|
||||
def test_chunk_rejection_after_unconfirmed_input_still_requires_observation(self):
|
||||
fixture = Executor()
|
||||
def execute(**kwargs):
|
||||
result = fixture(**kwargs)
|
||||
if fixture.requests[-1]['op'] == 'read_chunk':
|
||||
result.update(ok=False, exit_code='1', stdout=json.dumps(
|
||||
{'ok': False, 'error': {'code': 'UNKNOWN_REQUEST_FIELDS'}}))
|
||||
return result
|
||||
result = ScreenClient(execute, Image).call(
|
||||
'key', 'GERAMEG', True, capture_id=CAPTURE_ID, key='ENTER')
|
||||
self.assertEqual(result['error'], 'UNKNOWN_REQUEST_FIELDS')
|
||||
self.assertTrue(result['observe_again'])
|
||||
self.assertNotIn('performed', result)
|
||||
self.assertNotIn('capture_consumed', result)
|
||||
self.assertEqual([request['op'] for request in fixture.requests], ['key', 'read_chunk'])
|
||||
|
||||
def test_known_backend_rejection_keeps_original_meaning_without_uncertainty(self):
|
||||
for op, arguments in [('key', {'key': 'ENTER'}), ('text', {'text': 'fixture-only'})]:
|
||||
with self.subTest(op=op):
|
||||
fixture = Executor('error_exit1')
|
||||
result = ScreenClient(fixture, Image).call(
|
||||
op, 'GERAMEG', True, capture_id=CAPTURE_ID, **arguments)
|
||||
self.assertEqual(result['error'], 'SESSION_NOT_VISIBLE')
|
||||
self.assertNotIn('observe_again', result)
|
||||
self.assertNotIn('performed', result)
|
||||
self.assertEqual(len(fixture.requests), 1)
|
||||
|
||||
def test_truncation_and_corruption_never_return_partial_success(self):
|
||||
for variant, error in [('truncated', 'SCREEN_STDOUT_INVALID'), ('corrupt', 'SCREEN_IMAGE_HASH_MISMATCH')]:
|
||||
with self.subTest(variant=variant):
|
||||
self.fixture.variant = variant
|
||||
result = self.read(self.client(stdout_limit_bytes=EXPANDED_STDOUT_BYTES))
|
||||
self.assertIs(result['ok'], False)
|
||||
self.assertEqual(result['error'], error)
|
||||
self.assertNotIn('result', result)
|
||||
|
||||
def test_monotonic_total_deadline_propagates_remaining_whole_seconds(self):
|
||||
clock = Clock()
|
||||
self.fixture.clock, self.fixture.latency = clock, 1.2
|
||||
with patch('framebridge_expanded.real_screen.screen_frontdoor.time.monotonic', clock):
|
||||
result = self.read(self.client(stdout_limit_bytes=EXPANDED_STDOUT_BYTES, transfer_timeout_s=4))
|
||||
self.assertEqual(result['error'], 'SCREEN_TRANSFER_TIMEOUT')
|
||||
self.assertEqual([t for _, t in self.fixture.calls], [4, 2, 1])
|
||||
self.assertLess(clock.now, 104)
|
||||
self.assertNotIn('result', result)
|
||||
|
||||
def test_late_executor_response_cannot_publish_stale_success(self):
|
||||
clock = Clock()
|
||||
self.fixture.clock, self.fixture.latency = clock, 6
|
||||
with patch('framebridge_expanded.real_screen.screen_frontdoor.time.monotonic', clock):
|
||||
result = self.read(self.client(transfer_timeout_s=5))
|
||||
self.assertEqual(result['error'], 'SCREEN_TRANSFER_TIMEOUT')
|
||||
self.assertEqual(len(self.fixture.calls), 1)
|
||||
|
||||
def test_deadline_after_confirmed_input_preserves_no_replay_warning(self):
|
||||
clock = Clock()
|
||||
self.fixture.clock, self.fixture.latency = clock, 1.2
|
||||
with patch('framebridge_expanded.real_screen.screen_frontdoor.time.monotonic', clock):
|
||||
result = self.client(transfer_timeout_s=3).call(
|
||||
'key', 'GERAMEG', True, capture_id=CAPTURE_ID, key='ENTER')
|
||||
self.assertEqual(result['error'], 'SCREEN_TRANSFER_TIMEOUT')
|
||||
self.assertTrue(result['performed'])
|
||||
self.assertTrue(result['capture_consumed'])
|
||||
self.assertTrue(result['observe_again'])
|
||||
self.assertEqual(result['capture_id'], CAPTURE_ID)
|
||||
|
||||
def test_executor_exception_during_input_requires_fresh_observation(self):
|
||||
def failed(**kwargs):
|
||||
raise TimeoutError('test fixture timeout')
|
||||
result = ScreenClient(failed, Image).call(
|
||||
'key', 'GERAMEG', True, capture_id=CAPTURE_ID, key='ENTER')
|
||||
self.assertFalse(result['ok'])
|
||||
self.assertTrue(result['observe_again'])
|
||||
self.assertNotIn('performed', result)
|
||||
self.assertEqual(result['capture_id'], CAPTURE_ID)
|
||||
|
||||
def test_chunk_extension_is_bounded_strict_and_legacy_default_remains(self):
|
||||
request = {'op': 'read_chunk', 'capture_id': CAPTURE_ID, 'index': 0,
|
||||
'expected_sha256': self.fixture.manifest['sha256']}
|
||||
self.assertEqual(rpc.dispatch(request, self.directory)['result']['chunk_bytes'], CHUNK_BYTES)
|
||||
for value in (True, 0, -1, EXPANDED_CHUNK_BYTES + 1, 1.5, '45000'):
|
||||
with self.subTest(value=value), self.assertRaisesRegex(rpc.ScreenRpcError, 'CHUNK_BOUNDS_INVALID'):
|
||||
rpc.validate_request(dict(request, chunk_bytes=value))
|
||||
with self.assertRaisesRegex(rpc.ScreenRpcError, 'UNKNOWN_REQUEST_FIELDS'):
|
||||
rpc.validate_request(dict(request, extra='invalid'))
|
||||
|
||||
def test_client_stdout_budgets_reject_one_byte_over_exact_boundary(self):
|
||||
for budget in (MAX_STDOUT_BYTES, EXPANDED_STDOUT_BYTES):
|
||||
for over in (0, 1):
|
||||
with self.subTest(budget=budget, over=over):
|
||||
def execute(**kwargs):
|
||||
envelope = {'ok': True, 'result': {'padding': ''}}
|
||||
envelope['result']['padding'] = 'x' * (budget + over - len(rpc.encode(envelope)))
|
||||
return {'ok': True, 'exit_code': '0', 'target': kwargs['target'], 'node_id': NODE_ID,
|
||||
'mcp_job_id': kwargs['job_id'], 'mcp_script_sha256': kwargs['sha256'],
|
||||
'stdout': rpc.encode(envelope).decode(), 'stdout_truncated': False}
|
||||
client = ScreenClient(execute, Image, stdout_limit_bytes=budget)
|
||||
if over:
|
||||
with self.assertRaisesRegex(ScreenError, 'SCREEN_STDOUT_INVALID'):
|
||||
client.remote('GERAMEG', {'op': 'status'})
|
||||
else:
|
||||
self.assertIn('padding', client.remote('GERAMEG', {'op': 'status'}))
|
||||
|
||||
def test_typed_transport_exact_serialized_boundary(self):
|
||||
script = RPC_COMMAND + shlex.quote(base64.b64encode(rpc.encode({'op': 'status'})).decode()) + '\n'
|
||||
payload = {'version': 1, 'mode': 'EXEC', 'authorized': True, 'script': script,
|
||||
'timeout_s': 30, 'job_id': 'test-screen-boundary', 'sha256': hashlib.sha256(script.encode()).hexdigest(),
|
||||
'include_stdout': True, 'request_type': 'screen_rpc'}
|
||||
for over in (0, 1):
|
||||
envelope = {'ok': True, 'result': {'padding': ''}}
|
||||
envelope['result']['padding'] = 'x' * (EXPANDED_STDOUT_BYTES + over - len(rpc.encode(envelope)) - 1)
|
||||
stdout = rpc.encode(envelope) + b'\n'
|
||||
self.assertEqual(len(stdout), EXPANDED_STDOUT_BYTES + over)
|
||||
result = transport.handle_request(payload, runner=lambda *args: (0, stdout, b''))
|
||||
self.assertEqual(result['ok'], not over)
|
||||
if over:
|
||||
self.assertIn('SCREEN_RPC_OUTPUT_LIMIT', result['stdout'])
|
||||
|
||||
def test_main_serializes_competing_inputs_and_consumes_capture_once(self):
|
||||
request = base64.b64encode(rpc.encode({'op': 'key', 'capture_id': CAPTURE_ID, 'key': 'ENTER'})).decode()
|
||||
barrier = threading.Barrier(2)
|
||||
def invoke():
|
||||
barrier.wait(timeout=2)
|
||||
return rpc.main([request])
|
||||
with patch.object(rpc, 'state_directory', return_value=self.directory), \
|
||||
patch.object(rpc, 'act', return_value={'ok': True, 'performed': True}) as act, \
|
||||
patch.object(rpc, 'observe', return_value=dict(self.fixture.manifest)), \
|
||||
patch('builtins.print'), concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
|
||||
outcomes = list(pool.map(lambda _: invoke(), range(2)))
|
||||
self.assertEqual(sorted(outcomes), [0, 1])
|
||||
self.assertEqual(act.call_count, 1)
|
||||
self.assertTrue(rpc.load_capture(self.directory, CAPTURE_ID)['_action_consumed'])
|
||||
|
||||
def test_main_lock_covers_dispatch_in_competing_rpc_entries(self):
|
||||
first_entered, second_entered, release = threading.Event(), threading.Event(), threading.Event()
|
||||
count = 0
|
||||
mutex = threading.Lock()
|
||||
def dispatch(*args):
|
||||
nonlocal count
|
||||
with mutex:
|
||||
count += 1
|
||||
current = count
|
||||
if current == 1:
|
||||
first_entered.set()
|
||||
self.assertTrue(release.wait(2))
|
||||
else:
|
||||
second_entered.set()
|
||||
return {'ok': True, 'result': {}}
|
||||
request = base64.b64encode(rpc.encode({'op': 'status'})).decode()
|
||||
with patch.object(rpc, 'state_directory', return_value=self.directory), \
|
||||
patch.object(rpc, 'dispatch', side_effect=dispatch), patch('builtins.print'), \
|
||||
concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
|
||||
first = pool.submit(rpc.main, [request])
|
||||
try:
|
||||
self.assertTrue(first_entered.wait(2))
|
||||
second = pool.submit(rpc.main, [request])
|
||||
self.assertFalse(second_entered.wait(0.1))
|
||||
finally:
|
||||
release.set()
|
||||
self.assertEqual((first.result(timeout=2), second.result(timeout=2)), (0, 0))
|
||||
self.assertTrue(second_entered.is_set())
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user