diff --git a/pyisolate/migration.py b/pyisolate/migration.py index e90a648..4e6c78d 100644 --- a/pyisolate/migration.py +++ b/pyisolate/migration.py @@ -16,6 +16,11 @@ _MAGIC = b"PYISOMIG1" _HEADER_LIMIT = 16 * 1024 _DEFAULT_PORT = 8765 +# Checkpoints are JSON sandbox state plus AEAD overhead, so even generous +# deployments stay far below this. The bound exists because the peer-declared +# ``blob_len`` is read into memory *before* the HMAC can be verified; without +# it an unauthenticated peer could make the endpoint buffer arbitrary amounts. +DEFAULT_MAX_BLOB_LEN = 64 * 1024 * 1024 @dataclass(frozen=True) @@ -126,6 +131,7 @@ def handle_migration_connection( key: bytes, *, restore_fn: Callable[[bytes, bytes], Sandbox] = restore, + max_blob_len: int = DEFAULT_MAX_BLOB_LEN, ) -> MigrationResponse: """Restore one checkpoint received from *conn* and send a response.""" if len(key) != 32: @@ -138,6 +144,8 @@ def handle_migration_connection( auth = header.get("auth") if isinstance(blob_len, bool) or not isinstance(blob_len, int) or blob_len < 0: raise MigrationProtocolError("invalid checkpoint length") + if blob_len > max_blob_len: + raise MigrationProtocolError("checkpoint exceeds maximum size") if not isinstance(auth, str): raise MigrationProtocolError("missing migration authenticator") blob = _read_exact(conn, blob_len) @@ -154,14 +162,20 @@ def handle_migration_connection( def serve_migration_once( - host: str, key: bytes, *, restore_fn: Callable[[bytes, bytes], Sandbox] = restore + host: str, + key: bytes, + *, + restore_fn: Callable[[bytes, bytes], Sandbox] = restore, + max_blob_len: int = DEFAULT_MAX_BLOB_LEN, ) -> MigrationResponse: """Listen on *host* and handle a single migration request.""" address = _parse_host(host) with socket.create_server(address) as server: conn, _peer = server.accept() with conn: - return handle_migration_connection(conn, key, restore_fn=restore_fn) + return handle_migration_connection( + conn, key, restore_fn=restore_fn, max_blob_len=max_blob_len + ) def migrate(sandbox: Sandbox, host: str, key: bytes) -> MigrationResponse: diff --git a/tests/test_migration.py b/tests/test_migration.py index e585e38..f278a1d 100644 --- a/tests/test_migration.py +++ b/tests/test_migration.py @@ -150,3 +150,48 @@ def restore_fn(seen_blob, seen_key): # pragma: no cover - should never run assert response["ok"] is False assert "authenticator" in response["error"] assert restored == [] + + +def test_endpoint_rejects_oversized_checkpoint_before_reading_it(): + # blob_len is peer-controlled and read into memory before the HMAC can be + # checked, so the endpoint must reject an oversized declaration up front + # instead of buffering whatever the peer streams. + server, client = socket.socketpair() + restored = [] + result = [] + + def restore_fn(seen_blob, seen_key): # pragma: no cover - should never run + restored.append((seen_blob, seen_key)) + return object() + + def run_server(): + with server: + result.append( + migration.handle_migration_connection( + server, KEY, restore_fn=restore_fn, max_blob_len=1024 + ) + ) + + thread = threading.Thread(target=run_server) + thread.start() + try: + migration._send_json( + client, + { + "magic": "PYISOMIG1", + "version": 1, + "blob_len": 1025, + "auth": "0" * 64, + }, + ) + response = migration._recv_json(client) + finally: + client.close() + thread.join(timeout=1) + + assert response["ok"] is False + assert "maximum size" in response["error"] + assert result == [ + migration.MigrationResponse(ok=False, error="checkpoint exceeds maximum size") + ] + assert restored == []