Merge remote-tracking branch 'origin/main' into stlc/promote-next

Author: kernel-internal[bot]

src/hypeman/lib/_ws.py

@@ -16,6 +16,9 @@
 ]
 
 
+MAX_INBOUND_MESSAGE_SIZE = 2**20
+
+
 class ClientConfig(Protocol):
     """The generated client settings used by the custom WebSocket APIs."""
 
@@ -46,7 +49,7 @@ def __exit__(
 
 
 class SyncWebSocketConnector(Protocol):
-    def __call__(self, url: str, *, additional_headers: dict[str, str], max_size: int | None) -> SyncWebSocket: ...
+    def __call__(self, url: str, *, additional_headers: dict[str, str], max_size: int) -> SyncWebSocket: ...
 
 
 class AsyncWebSocket(Protocol):
@@ -81,16 +84,24 @@ async def __aexit__(
 
 
 class AsyncWebSocketConnector(Protocol):
-    def __call__(
-        self, url: str, *, additional_headers: dict[str, str], max_size: int | None
-    ) -> AsyncWebSocketContext: ...
+    def __call__(self, url: str, *, additional_headers: dict[str, str], max_size: int) -> AsyncWebSocketContext: ...
 
 
-def sync_connect(url: str, *, additional_headers: dict[str, str], max_size: int | None) -> SyncWebSocket:
+def sync_connect(
+    url: str,
+    *,
+    additional_headers: dict[str, str],
+    max_size: int = MAX_INBOUND_MESSAGE_SIZE,
+) -> SyncWebSocket:
     return cast(SyncWebSocket, websocket_connect(url, additional_headers=additional_headers, max_size=max_size))
 
 
-def async_connect(url: str, *, additional_headers: dict[str, str], max_size: int | None) -> AsyncWebSocketContext:
+def async_connect(
+    url: str,
+    *,
+    additional_headers: dict[str, str],
+    max_size: int = MAX_INBOUND_MESSAGE_SIZE,
+) -> AsyncWebSocketContext:
     return cast(
         AsyncWebSocketContext,
         async_websocket_connect(url, additional_headers=additional_headers, max_size=max_size),

src/hypeman/lib/cp.py

@@ -12,6 +12,7 @@
 from collections.abc import Callable
 
 from ._ws import (
+    MAX_INBOUND_MESSAGE_SIZE,
     ClientConfig,
     SyncWebSocket,
     AsyncWebSocket,
@@ -220,7 +221,7 @@ def cp_to_instance(
     )
     url, headers = connection_settings(client, instance_id, "cp")
     for entry in entries:
-        with connector(url, additional_headers=headers, max_size=None) as websocket:
+        with connector(url, additional_headers=headers, max_size=MAX_INBOUND_MESSAGE_SIZE) as websocket:
             websocket.send(_request(entry))
             if entry.is_dir:
                 websocket.send('{"type":"end"}')
@@ -280,7 +281,7 @@ async def cp_to_instance_async(
     )
     url, headers = connection_settings(client, instance_id, "cp")
     for entry in entries:
-        async with connector(url, additional_headers=headers, max_size=None) as websocket:
+        async with connector(url, additional_headers=headers, max_size=MAX_INBOUND_MESSAGE_SIZE) as websocket:
             await websocket.send(_request(entry))
             if entry.is_dir:
                 await websocket.send('{"type":"end"}')
@@ -507,7 +508,7 @@ def cp_from_instance(
     url, headers = connection_settings(client, instance_id, "cp")
     state = _DownloadState(Path(dst_path), archive, callbacks)
     try:
-        with connector(url, additional_headers=headers, max_size=None) as websocket:
+        with connector(url, additional_headers=headers, max_size=MAX_INBOUND_MESSAGE_SIZE) as websocket:
             websocket.send(_download_request(src_path, follow_symlinks))
             while not state.complete:
                 try:
@@ -536,7 +537,7 @@ async def cp_from_instance_async(
     url, headers = connection_settings(client, instance_id, "cp")
     state = _DownloadState(Path(dst_path), archive, callbacks)
     try:
-        async with connector(url, additional_headers=headers, max_size=None) as websocket:
+        async with connector(url, additional_headers=headers, max_size=MAX_INBOUND_MESSAGE_SIZE) as websocket:
             await websocket.send(_download_request(src_path, follow_symlinks))
             while not state.complete:
                 try:

src/hypeman/lib/exec.py

@@ -7,6 +7,7 @@
 from typing_extensions import TypeAlias
 
 from ._ws import (
+    MAX_INBOUND_MESSAGE_SIZE,
     ClientConfig,
     SyncWebSocketConnector,
     AsyncWebSocketConnector,
@@ -53,6 +54,7 @@ class _ExecRequest:
     wait_for_agent: int | None
     rows: int | None
     cols: int | None
+    resize: tuple[tuple[int, int], ...]
 
     def encode(self) -> str:
         payload: dict[str, object] = {"command": self.command, "tty": self.tty}
@@ -79,6 +81,7 @@ def _request(
     tty: bool,
     rows: int | None,
     cols: int | None,
+    resize: Iterable[tuple[int, int]],
 ) -> _ExecRequest:
     if isinstance(command, str):
         raise ValueError("command must be an argument sequence, not a string")
@@ -93,7 +96,13 @@ def _request(
             raise ValueError(f"{name} must be positive")
     if (rows is not None or cols is not None) and not tty:
         raise ValueError("rows and cols require tty=True")
-    return _ExecRequest(argv, tty, env, cwd, timeout, wait_for_agent, rows, cols)
+    resize_events = tuple(resize)
+    for resize_rows, resize_cols in resize_events:
+        if any(isinstance(value, bool) or value <= 0 for value in (resize_rows, resize_cols)):
+            raise ValueError("resize dimensions must be positive integers")
+    if resize_events and not tty:
+        raise ValueError("resize requires tty=True")
+    return _ExecRequest(argv, tty, env, cwd, timeout, wait_for_agent, rows, cols, resize_events)
 
 
 def _stdin_chunks(stdin: Stdin | None) -> Iterable[bytes]:
@@ -161,17 +170,16 @@ def exec(
         tty=tty,
         rows=rows,
         cols=cols,
+        resize=resize,
     )
     url, headers = connection_settings(client, instance_id, "exec")
     output = bytearray()
-    with connector(url, additional_headers=headers, max_size=None) as websocket:
+    with connector(url, additional_headers=headers, max_size=MAX_INBOUND_MESSAGE_SIZE) as websocket:
         websocket.send(request.encode())
         for chunk in _stdin_chunks(stdin):
             if chunk:
                 websocket.send(chunk)
-        for resize_rows, resize_cols in resize:
-            if not tty or resize_rows <= 0 or resize_cols <= 0:
-                raise ValueError("resize dimensions must be positive and require tty=True")
+        for resize_rows, resize_cols in request.resize:
             websocket.send(json.dumps({"resize": {"rows": resize_rows, "cols": resize_cols}}, separators=(",", ":")))
 
         while True:
@@ -212,10 +220,11 @@ async def exec_async(
         tty=tty,
         rows=rows,
         cols=cols,
+        resize=resize,
     )
     url, headers = connection_settings(client, instance_id, "exec")
     output = bytearray()
-    async with connector(url, additional_headers=headers, max_size=None) as websocket:
+    async with connector(url, additional_headers=headers, max_size=MAX_INBOUND_MESSAGE_SIZE) as websocket:
         await websocket.send(request.encode())
         if isinstance(stdin, AsyncIterable):
             async for chunk in stdin:
@@ -225,9 +234,7 @@ async def exec_async(
             for chunk in _stdin_chunks(stdin):
                 if chunk:
                     await websocket.send(chunk)
-        for resize_rows, resize_cols in resize:
-            if not tty or resize_rows <= 0 or resize_cols <= 0:
-                raise ValueError("resize dimensions must be positive and require tty=True")
+        for resize_rows, resize_cols in request.resize:
             await websocket.send(
                 json.dumps({"resize": {"rows": resize_rows, "cols": resize_cols}}, separators=(",", ":"))
             )

tests/lib/test_websocket_lib.py

@@ -9,6 +9,7 @@
 from collections.abc import AsyncIterator
 
 import pytest
+from websockets.exceptions import PayloadTooBig
 
 from hypeman.lib import (
     CopyCallbacks,
@@ -62,9 +63,9 @@ def __exit__(
 @dataclass
 class FakeConnector:
     connections: deque[FakeWebSocket]
-    calls: list[tuple[str, dict[str, str], int | None]] = field(default_factory=lambda: [])
+    calls: list[tuple[str, dict[str, str], int]] = field(default_factory=lambda: [])
 
-    def __call__(self, url: str, *, additional_headers: dict[str, str], max_size: int | None) -> FakeWebSocket:
+    def __call__(self, url: str, *, additional_headers: dict[str, str], max_size: int) -> FakeWebSocket:
         self.calls.append((url, additional_headers, max_size))
         return self.connections.popleft()
 
@@ -102,9 +103,9 @@ async def __aexit__(
 @dataclass
 class FakeAsyncConnector:
     connections: deque[FakeAsyncWebSocket]
-    calls: list[tuple[str, dict[str, str], int | None]] = field(default_factory=lambda: [])
+    calls: list[tuple[str, dict[str, str], int]] = field(default_factory=lambda: [])
 
-    def __call__(self, url: str, *, additional_headers: dict[str, str], max_size: int | None) -> FakeAsyncWebSocket:
+    def __call__(self, url: str, *, additional_headers: dict[str, str], max_size: int) -> FakeAsyncWebSocket:
         self.calls.append((url, additional_headers, max_size))
         return self.connections.popleft()
 
@@ -155,7 +156,7 @@ def test_exec_uses_client_auth_url_and_all_request_dimensions() -> None:
     assert result.output == b"outerr"
     assert result.exit_code == 0
     assert connector.calls == [
-        ("wss://example.test/api/instances/inst_123/exec", {"Authorization": "Bearer secret"}, None)
+        ("wss://example.test/api/instances/inst_123/exec", {"Authorization": "Bearer secret"}, 2**20)
     ]
     assert json.loads(str(websocket.sent[0])) == {
         "command": ["sh", "-lc", "echo hi"],
@@ -196,6 +197,31 @@ def test_exec_never_retries_after_dispatch() -> None:
     assert len(websocket.sent) == 1
 
 
+@pytest.mark.parametrize(("tty", "resize"), [(False, [(24, 80)]), (True, [(0, 80)]), (True, [(24, -1)])])
+def test_exec_rejects_invalid_resize_before_connect(tty: bool, resize: list[tuple[int, int]]) -> None:
+    connector = FakeConnector(deque())
+    with pytest.raises(ValueError, match="resize"):
+        exec(FakeClient(), "inst", ["true"], tty=tty, resize=resize, connector=connector)
+    assert not connector.calls
+
+
+def test_exec_rejects_oversized_inbound_message() -> None:
+    oversized = PayloadTooBig(2**20 + 1, 2**20)
+    connector = FakeConnector(deque([FakeWebSocket(deque([oversized]))]))
+    with pytest.raises(ExecProtocolError, match="before an exitCode") as exc_info:
+        exec(FakeClient(), "inst", ["true"], connector=connector)
+    assert isinstance(exc_info.value.__cause__, PayloadTooBig)
+    assert connector.calls[0][2] == 2**20
+
+
+@pytest.mark.asyncio
+async def test_exec_async_rejects_invalid_resize_before_connect() -> None:
+    connector = FakeAsyncConnector(deque())
+    with pytest.raises(ValueError, match="resize"):
+        await exec_async(FakeClient(), "inst", ["true"], tty=False, resize=[(24, 80)], connector=connector)
+    assert not connector.calls
+
+
 @pytest.mark.asyncio
 async def test_exec_async_supports_streaming_stdin() -> None:
     websocket = FakeAsyncWebSocket(deque([b"done", '{"exitCode":7}']))
@@ -248,6 +274,7 @@ def test_cp_upload_file_preserves_mode_and_reports_progress(tmp_path: Path) -> N
         expected_request["gid"] = source_stat.st_gid
     assert request == expected_request
     assert websocket.sent[1:] == [b"payload", '{"type":"end"}']
+    assert connector.calls[0][2] == 2**20
     assert events == [
         ("start", (str(source), 7)),
         ("progress", 7),
@@ -420,27 +447,31 @@ async def test_cp_async_upload_and_download(tmp_path: Path) -> None:
     source = tmp_path / "source"
     source.write_bytes(b"async")
     upload_socket = FakeAsyncWebSocket(deque([upload_result(5)]))
+    upload_connector = FakeAsyncConnector(deque([upload_socket]))
     await cp_to_instance_async(
         FakeClient(),
         "inst",
         source,
         "/guest/source",
-        connector=FakeAsyncConnector(deque([upload_socket])),
+        connector=upload_connector,
     )
     assert upload_socket.sent[1] == b"async"
+    assert upload_connector.calls[0][2] == 2**20
 
     download_socket = FakeAsyncWebSocket(
         deque([file_header("result", size=5), b"async", text_frame({"type": "end", "final": True})])
     )
     destination = tmp_path / "dest"
+    download_connector = FakeAsyncConnector(deque([download_socket]))
     await cp_from_instance_async(
         FakeClient(),
         "inst",
         "/guest/result",
         destination,
-        connector=FakeAsyncConnector(deque([download_socket])),
+        connector=download_connector,
     )
     assert (destination / "result").read_bytes() == b"async"
+    assert download_connector.calls[0][2] == 2**20
 
 
 def test_invalid_instance_id_is_rejected_before_connect() -> None: