diff --git a/dataplicity_cli/cli.py b/dataplicity_cli/cli.py index b759c41..3526a7c 100644 --- a/dataplicity_cli/cli.py +++ b/dataplicity_cli/cli.py @@ -605,6 +605,31 @@ def _friendly_response_message(default_message: str, response_data: Any, respons return message +def _debug_api_request( + state: AppContext, + *, + method: str, + path: str, + params: Optional[Dict[str, Any]] = None, + json_data: Optional[Dict[str, Any]] = None, + data: Optional[Dict[str, Any]] = None, +) -> None: + if not state.debug: + return + if path.startswith("http://") or path.startswith("https://"): + url = path + else: + url = f"{state.config.base_url.rstrip('/')}/{path.lstrip('/')}" + payload: Dict[str, Any] = {"method": method.upper(), "url": url} + if params: + payload["params"] = _sanitize_payload(params) + if json_data: + payload["json"] = _sanitize_payload(json_data) + if data: + payload["data"] = _sanitize_payload(data) + sys.stderr.write(f"[debug] API request {json.dumps(payload, sort_keys=True)}\n") + + def _incident_automation_org_hash(endpoint: str) -> str: match = re.search(r"/api/organisations/([^/]+)/", endpoint) if not match: @@ -919,7 +944,7 @@ class _PortForwardDashboard: active_clients: int = 0 total_connections: int = 0 rejected_connections: int = 0 - current_connection_started_at: Optional[float] = None + connection_started_at: Dict[str, float] = field(default_factory=dict) first_byte_latencies_ms: List[float] = field(default_factory=list) protocols_seen: List[str] = field(default_factory=list) listener_ready: bool = False @@ -940,19 +965,25 @@ def apply_event(self, event: Any) -> None: self.stop_reason = "stopped" return if kind == "connection_opened": - self.active_clients = 1 + self.active_clients += 1 self.total_connections += 1 - self.current_connection_started_at = timestamp + self.connection_started_at[detail] = timestamp return if kind == "connection_closed": - self.active_clients = 0 - self.current_connection_started_at = None + self.active_clients = max(0, self.active_clients - 1) + if detail: + self.connection_started_at.pop(detail, None) return if kind == "connection_rejected": self.rejected_connections += 1 return - if kind == "first_remote_byte" and self.current_connection_started_at is not None: - latency_ms = (timestamp - self.current_connection_started_at) * 1000.0 + if kind == "first_remote_byte": + started_at = self.connection_started_at.get(detail) + if started_at is None and self.connection_started_at: + started_at = next(iter(self.connection_started_at.values())) + if started_at is None: + return + latency_ms = (timestamp - started_at) * 1000.0 self.first_byte_latencies_ms.append(latency_ms) self.first_byte_latencies_ms = self.first_byte_latencies_ms[-64:] return @@ -2206,7 +2237,9 @@ def devices_history_list( if include_deleted: params["include_deleted"] = "1" - response = state.api.get(f"/api/developer/devices/{resolved_hash}/timeline/", params=params) + endpoint = f"/api/developer/devices/{resolved_hash}/timeline/" + _debug_api_request(state, method="GET", path=endpoint, params=params) + response = state.api.get(endpoint, params=params) if not response.ok: message = _friendly_response_message("Unable to load device history.", response.data, response.text) if state.json_output: @@ -2270,7 +2303,9 @@ def devices_history_show( params: Dict[str, Any] = {"limit": 500} if include_deleted: params["include_deleted"] = "1" - response = state.api.get(f"/api/developer/devices/{resolved_hash}/timeline/", params=params) + endpoint = f"/api/developer/devices/{resolved_hash}/timeline/" + _debug_api_request(state, method="GET", path=endpoint, params=params) + response = state.api.get(endpoint, params=params) if not response.ok: message = _friendly_response_message("Unable to load device history.", response.data, response.text) if state.json_output: @@ -2361,7 +2396,9 @@ def devices_history_comment( _show_error(state.console, message) raise typer.Exit(code=2) payload["type"] = normalized_type - response = state.api.post(f"/api/developer/devices/{resolved_hash}/timeline/", json_data=payload) + endpoint = f"/api/developer/devices/{resolved_hash}/timeline/" + _debug_api_request(state, method="POST", path=endpoint, json_data=payload) + response = state.api.post(endpoint, json_data=payload) if not response.ok: message = _friendly_response_message("Unable to add device history comment.", response.data, response.text) if state.json_output: @@ -2399,9 +2436,11 @@ def devices_history_delete( state = _ctx(ctx) _require_auth(state) resolved_hash = _resolve_device_hash_interactive(state, device_hash, action_name="history delete") + endpoint = f"/api/developer/devices/{resolved_hash}/timeline/{event_id}/" + _debug_api_request(state, method="DELETE", path=endpoint) response = state.api.request( "DELETE", - f"/api/developer/devices/{resolved_hash}/timeline/{event_id}/", + endpoint, ) if not response.ok: message = _friendly_response_message("Unable to mark history entry as deleted.", response.data, response.text) @@ -2784,24 +2823,32 @@ async def runner() -> Dict[str, Any]: await m2m.connect() try: identity = await m2m.wait_for_identity() - response = _open_remote_access_port( - state, - resolved_hash, - { - "m2m_identity": identity, - "service": "redirect-port", - "port": remote_port, - }, - ) - if not response.ok: - detail = _friendly_response_message("Unable to open port redirect", response.data, response.text) - raise RuntimeError(detail) - channel_port = await m2m.wait_for_channel_open() + open_channel_lock = asyncio.Lock() + + async def open_redirect_channel() -> int: + async with open_channel_lock: + response = await asyncio.to_thread( + _open_remote_access_port, + state, + resolved_hash, + { + "m2m_identity": identity, + "service": "redirect-port", + "port": remote_port, + }, + ) + if not response.ok: + detail = _friendly_response_message("Unable to open port redirect", response.data, response.text) + raise RuntimeError(detail) + return await m2m.wait_for_channel_open() + + channel_port = await open_redirect_channel() forward_task = asyncio.create_task( run_port_forward( m2m, channel_port, local_port, + channel_factory=open_redirect_channel, event_callback=dashboard.apply_event, ) ) @@ -2912,18 +2959,27 @@ def on_forward_event(event: Any) -> None: try: identity = await m2m.wait_for_identity(timeout=float(connect_timeout)) - response = _open_remote_access_port( - state, - resolved_hash, - { - "m2m_identity": identity, - "service": "redirect-port", - "port": remote_port, - }, - ) - if not response.ok: - raise RuntimeError(_friendly_response_message("Unable to open SSH tunnel.", response.data, response.text)) - channel_port = await m2m.wait_for_channel_open(timeout=float(connect_timeout)) + open_channel_lock = asyncio.Lock() + + async def open_redirect_channel() -> int: + async with open_channel_lock: + response = await asyncio.to_thread( + _open_remote_access_port, + state, + resolved_hash, + { + "m2m_identity": identity, + "service": "redirect-port", + "port": remote_port, + }, + ) + if not response.ok: + raise RuntimeError( + _friendly_response_message("Unable to open SSH tunnel.", response.data, response.text) + ) + return await m2m.wait_for_channel_open(timeout=float(connect_timeout)) + + channel_port = await open_redirect_channel() from .remote_access import run_port_forward forward_task = asyncio.create_task( @@ -2931,6 +2987,7 @@ def on_forward_event(event: Any) -> None: m2m, channel_port, resolved_local_port, + channel_factory=open_redirect_channel, event_callback=on_forward_event, ) ) diff --git a/dataplicity_cli/remote_access.py b/dataplicity_cli/remote_access.py index f86205b..0f99dcf 100644 --- a/dataplicity_cli/remote_access.py +++ b/dataplicity_cli/remote_access.py @@ -9,7 +9,7 @@ import time import tty from dataclasses import dataclass -from typing import Callable, Optional +from typing import Awaitable, Callable, Optional from .m2m import M2MClient @@ -43,6 +43,7 @@ class PortForwardEvent: PortForwardEventCallback = Callable[[PortForwardEvent], None] +PortForwardChannelFactory = Callable[[], Awaitable[int]] def _detect_protocol(sample: bytes) -> Optional[str]: @@ -226,13 +227,14 @@ async def run_remote_file( async def run_port_forward( m2m: M2MClient, - channel_port: int, + channel_port: Optional[int], local_port: int, *, + channel_factory: Optional[PortForwardChannelFactory] = None, event_callback: Optional[PortForwardEventCallback] = None, ) -> None: - queue = m2m.channel_queue(channel_port) - active_client = asyncio.Lock() + initial_channel_claimed = False + initial_channel_lock = asyncio.Lock() def emit(kind: str, *, bytes_count: int = 0, detail: str = "") -> None: if event_callback: @@ -245,29 +247,35 @@ def emit(kind: str, *, bytes_count: int = 0, detail: str = "") -> None: ) ) + async def allocate_channel() -> int: + nonlocal initial_channel_claimed + async with initial_channel_lock: + if channel_port is not None and not initial_channel_claimed: + initial_channel_claimed = True + return channel_port + if channel_factory is None: + raise RuntimeError("No remote channel available for this connection.") + return await channel_factory() + async def handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: peer = writer.get_extra_info("peername") peer_label = str(peer) if peer else "unknown" - if active_client.locked(): - emit("connection_rejected", detail=peer_label) - writer.close() - await writer.wait_closed() - return - - await active_client.acquire() - acquired = True - emit("connection_opened", detail=peer_label) + connection_label = f"{peer_label}/{id(writer):x}" + channel_for_client: Optional[int] = None + emit("connection_opened", detail=connection_label) local_started = asyncio.Event() first_remote_byte_seen = False protocols_seen = set() async def local_to_remote() -> None: + if channel_for_client is None: + return while True: data = await reader.read(65536) if not data: break local_started.set() - await m2m.send_route(channel_port, data) + await m2m.send_route(channel_for_client, data) emit("bytes_up", bytes_count=len(data)) protocol = _detect_protocol(data[:32]) if protocol and protocol not in protocols_seen: @@ -276,13 +284,16 @@ async def local_to_remote() -> None: async def remote_to_local() -> None: nonlocal first_remote_byte_seen + if channel_for_client is None: + return + queue = m2m.channel_queue(channel_for_client) while True: data = await queue.get() if data is None: break if local_started.is_set() and not first_remote_byte_seen: first_remote_byte_seen = True - emit("first_remote_byte") + emit("first_remote_byte", detail=connection_label) writer.write(data) await writer.drain() emit("bytes_down", bytes_count=len(data)) @@ -292,13 +303,29 @@ async def remote_to_local() -> None: emit("protocol_detected", detail=protocol) try: - await asyncio.gather(local_to_remote(), remote_to_local()) + try: + channel_for_client = await allocate_channel() + except Exception as exc: + emit("connection_rejected", detail=f"{connection_label}: {exc}") + return + to_remote_task = asyncio.create_task(local_to_remote()) + to_local_task = asyncio.create_task(remote_to_local()) + done, pending = await asyncio.wait( + {to_remote_task, to_local_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + for task in pending: + task.cancel() + await asyncio.gather(*done, *pending, return_exceptions=True) finally: + if channel_for_client is not None: + try: + await m2m.close_channel(channel_for_client) + except Exception as exc: + emit("channel_close_failed", detail=f"{connection_label}: {exc}") writer.close() await writer.wait_closed() - emit("connection_closed", detail=peer_label) - if acquired and active_client.locked(): - active_client.release() + emit("connection_closed", detail=connection_label) server = await asyncio.start_server(handle_client, host="127.0.0.1", port=local_port) emit("listener_started", detail=f"127.0.0.1:{local_port}") diff --git a/tests/test_devices_history_cli.py b/tests/test_devices_history_cli.py index b1c8c8b..8316590 100644 --- a/tests/test_devices_history_cli.py +++ b/tests/test_devices_history_cli.py @@ -223,6 +223,35 @@ def fake_request( self.assertEqual(payload.get("id"), 777) self.assertEqual(payload.get("is_deleted"), True) + def test_history_list_debug_prints_endpoint(self) -> None: + timeline_payload = {"results": [], "count": 0} + + def fake_get(_self, path: str, *, params=None): # type: ignore[no-untyped-def] + _ = params + self.assertEqual(path, f"/api/developer/devices/{self.device_hash}/timeline/") + return ApiResponse(True, 200, timeline_payload, json.dumps(timeline_payload)) + + with tempfile.TemporaryDirectory() as temp_dir: + config_path = Path(temp_dir) / "cli.json" + self._write_authed_config(config_path) + with patch("dataplicity_cli.cli.ApiClient.get", new=fake_get): + result = self.runner.invoke( + app, + [ + "--debug", + "--config", + str(config_path), + "devices", + "history", + "list", + self.device_hash, + ], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + self.assertIn("/api/developer/devices/", result.output) + self.assertIn("/timeline/", result.output) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_remote_access_helpers.py b/tests/test_remote_access_helpers.py index 574905a..bade0d9 100644 --- a/tests/test_remote_access_helpers.py +++ b/tests/test_remote_access_helpers.py @@ -2,19 +2,22 @@ import asyncio import io +import socket import tempfile import unittest from pathlib import Path from typing import Dict, Optional from unittest.mock import patch -from dataplicity_cli.remote_access import _detect_protocol, run_remote_file, run_single_command +from dataplicity_cli.remote_access import _detect_protocol, run_port_forward, run_remote_file, run_single_command class _FakeM2M: def __init__(self) -> None: self.queues: Dict[int, asyncio.Queue[Optional[bytes]]] = {} self.sent_payloads: list[bytes] = [] + self.sent_routes: list[tuple[int, bytes]] = [] + self.closed_channels: list[int] = [] def channel_queue(self, port: int) -> asyncio.Queue[Optional[bytes]]: if port not in self.queues: @@ -22,11 +25,11 @@ def channel_queue(self, port: int) -> asyncio.Queue[Optional[bytes]]: return self.queues[port] async def send_route(self, port: int, data: bytes) -> None: - _ = port + self.sent_routes.append((port, data)) self.sent_payloads.append(data) async def close_channel(self, port: int) -> None: - _ = port + self.closed_channels.append(port) class _StdoutWithBuffer: @@ -34,6 +37,12 @@ def __init__(self) -> None: self.buffer = io.BytesIO() +def _unused_local_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.bind(("127.0.0.1", 0)) + return int(probe.getsockname()[1]) + + class RemoteAccessHelpersTest(unittest.IsolatedAsyncioTestCase): async def test_run_single_command_rejects_empty_command(self) -> None: fake = _FakeM2M() @@ -80,6 +89,80 @@ async def test_detect_protocol_classifies_known_signatures(self) -> None: self.assertEqual(_detect_protocol(bytes([0x16, 0x03, 0x03, 0x00])), "TLS") self.assertIsNone(_detect_protocol(b"\x01\x02\x03")) + async def test_run_port_forward_allocates_channel_per_local_client(self) -> None: + fake = _FakeM2M() + events = [] + next_channels = iter([102]) + + async def channel_factory() -> int: + return next(next_channels) + + local_port = _unused_local_port() + forward_task = asyncio.create_task( + run_port_forward( + fake, + 101, + local_port, + channel_factory=channel_factory, + event_callback=events.append, + ) + ) + try: + for _ in range(50): + if any(event.kind == "listener_started" for event in events): + break + await asyncio.sleep(0.01) + self.assertTrue(any(event.kind == "listener_started" for event in events)) + + reader_one, writer_one = await asyncio.open_connection("127.0.0.1", local_port) + reader_two, writer_two = await asyncio.open_connection("127.0.0.1", local_port) + try: + writer_one.write(b"GET /one HTTP/1.1\r\n\r\n") + writer_two.write(b"GET /two HTTP/1.1\r\n\r\n") + await writer_one.drain() + await writer_two.drain() + + for _ in range(50): + if len(fake.sent_routes) >= 2: + break + await asyncio.sleep(0.01) + + route_by_payload = {payload: port for port, payload in fake.sent_routes[:2]} + self.assertEqual( + route_by_payload, + { + b"GET /one HTTP/1.1\r\n\r\n": 101, + b"GET /two HTTP/1.1\r\n\r\n": 102, + }, + ) + + await fake.channel_queue(route_by_payload[b"GET /one HTTP/1.1\r\n\r\n"]).put(b"HTTP/1.1 200 OK\r\n\r\none") + await fake.channel_queue(route_by_payload[b"GET /one HTTP/1.1\r\n\r\n"]).put(None) + await fake.channel_queue(route_by_payload[b"GET /two HTTP/1.1\r\n\r\n"]).put(b"HTTP/1.1 200 OK\r\n\r\ntwo") + await fake.channel_queue(route_by_payload[b"GET /two HTTP/1.1\r\n\r\n"]).put(None) + + self.assertEqual(await asyncio.wait_for(reader_one.read(1024), timeout=0.5), b"HTTP/1.1 200 OK\r\n\r\none") + self.assertEqual(await asyncio.wait_for(reader_two.read(1024), timeout=0.5), b"HTTP/1.1 200 OK\r\n\r\ntwo") + for _ in range(50): + if sorted(set(fake.closed_channels)) == [101, 102]: + break + await asyncio.sleep(0.01) + self.assertEqual(sorted(set(fake.closed_channels)), [101, 102]) + self.assertEqual(len([event for event in events if event.kind == "connection_rejected"]), 0) + finally: + writer_one.close() + writer_two.close() + await writer_one.wait_closed() + await writer_two.wait_closed() + finally: + forward_task.cancel() + cancelled = False + try: + await forward_task + except asyncio.CancelledError: + cancelled = True + self.assertTrue(cancelled) + if __name__ == "__main__": unittest.main()