@neoline/hostbridge 2.0.4 → 2.0.5
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +11 -9
- package/docs/ARCHITECTURE.md +62 -0
- package/package.json +2 -1
- package/pyproject.toml +1 -1
- package/src/server_control_mcp/__init__.py +4 -3
- package/src/server_control_mcp/async_lifecycle.py +41 -0
- package/src/server_control_mcp/cli.py +1 -1
- package/src/server_control_mcp/client.py +302 -287
- package/src/server_control_mcp/config.py +2 -1
- package/src/server_control_mcp/daemon.py +233 -526
- package/src/server_control_mcp/doctor.py +34 -2
- package/src/server_control_mcp/hosts.py +136 -2
- package/src/server_control_mcp/mock_ssh.py +52 -14
- package/src/server_control_mcp/mock_ssh_exec.py +147 -58
- package/src/server_control_mcp/mock_ssh_session.py +3 -2
- package/src/server_control_mcp/mock_ssh_sftp.py +126 -28
- package/src/server_control_mcp/mock_ssh_tunnel.py +141 -444
- package/src/server_control_mcp/mux_connection.py +326 -0
- package/src/server_control_mcp/mux_daemon.py +186 -0
- package/src/server_control_mcp/mux_protocol.py +279 -0
- package/src/server_control_mcp/mux_records.py +71 -0
- package/src/server_control_mcp/mux_rpc.py +1779 -0
- package/src/server_control_mcp/mux_service.py +506 -0
- package/src/server_control_mcp/mux_stream.py +283 -0
- package/src/server_control_mcp/mux_sync_client.py +224 -0
- package/src/server_control_mcp/remote_agent_bundle.py +56 -0
- package/src/server_control_mcp/remote_mux_agent.py +769 -0
- package/src/server_control_mcp/runtime_services.py +862 -0
- package/src/server_control_mcp/server.py +33 -18
- package/src/server_control_mcp/services.py +2 -666
- package/src/server_control_mcp/transports/__init__.py +4 -8
- package/src/server_control_mcp/transports/base.py +60 -38
- package/src/server_control_mcp/transports/ssh.py +206 -259
- package/src/server_control_mcp/tunnel_manager.py +1245 -0
- package/src/server_control_mcp/tunnel_native_ssh.py +454 -0
- package/src/server_control_mcp/tunnel_providers.py +20 -0
- package/src/server_control_mcp/tunnel_pty_agent.py +909 -0
- package/src/server_control_mcp/protocol.py +0 -363
- package/src/server_control_mcp/remote_tunnel_agent.py +0 -143
- package/src/server_control_mcp/transports/pty.py +0 -515
|
@@ -1,37 +1,24 @@
|
|
|
1
|
-
"""Typed synchronous client for the HostBridge daemon protocol."""
|
|
1
|
+
"""Typed synchronous client for the HostBridge multiplexed daemon protocol."""
|
|
2
2
|
|
|
3
3
|
from __future__ import annotations
|
|
4
4
|
|
|
5
|
-
import base64
|
|
6
|
-
import binascii
|
|
7
5
|
import hashlib
|
|
8
|
-
import json
|
|
9
|
-
import socket
|
|
10
6
|
import threading
|
|
7
|
+
import time
|
|
11
8
|
import uuid
|
|
12
9
|
from collections.abc import Iterable, Mapping
|
|
13
10
|
from contextlib import suppress
|
|
14
11
|
from dataclasses import dataclass
|
|
15
12
|
from pathlib import Path
|
|
13
|
+
from types import TracebackType
|
|
16
14
|
|
|
17
15
|
from . import __version__
|
|
18
|
-
from .
|
|
19
|
-
|
|
20
|
-
|
|
21
|
-
|
|
22
|
-
ErrorInfo,
|
|
23
|
-
FrameFlags,
|
|
24
|
-
FrameSequenceValidator,
|
|
25
|
-
ProtocolError,
|
|
26
|
-
Request,
|
|
27
|
-
Response,
|
|
28
|
-
encode_frame,
|
|
29
|
-
read_frame,
|
|
30
|
-
)
|
|
16
|
+
from .mux_protocol import MUX_PROTOCOL_VERSION
|
|
17
|
+
from .mux_records import MuxRecordChannel, MuxRecordDecoder, MuxRecordError
|
|
18
|
+
from .mux_rpc import MuxCapacityError, MuxConnectionClosed, MuxRemoteError
|
|
19
|
+
from .mux_sync_client import MuxSyncClient, MuxSyncStream
|
|
31
20
|
from .runtime import runtime_paths
|
|
32
|
-
from .transports.base import ShellReadResult, TransferResult
|
|
33
|
-
|
|
34
|
-
MAX_CONTROL_LINE_BYTES = 4 * 1024 * 1024
|
|
21
|
+
from .transports.base import ShellReadResult, ShellTerminalResult, TransferResult
|
|
35
22
|
|
|
36
23
|
|
|
37
24
|
def _required_text(value: object, field: str) -> str:
|
|
@@ -40,23 +27,6 @@ def _required_text(value: object, field: str) -> str:
|
|
|
40
27
|
return value
|
|
41
28
|
|
|
42
29
|
|
|
43
|
-
def _encode_inline_bytes(value: bytes | None) -> dict[str, str] | None:
|
|
44
|
-
if value is None:
|
|
45
|
-
return None
|
|
46
|
-
if not isinstance(value, bytes):
|
|
47
|
-
raise TypeError("stdin must be bytes or None")
|
|
48
|
-
return {"encoding": "base64", "data": base64.b64encode(value).decode("ascii")}
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
def _decode_inline_bytes(value: object, field: str) -> bytes:
|
|
52
|
-
if not isinstance(value, dict) or value.get("encoding") != "base64" or not isinstance(value.get("data"), str):
|
|
53
|
-
raise HostBridgeError("protocol_mismatch", f"response field {field} must contain base64 data")
|
|
54
|
-
try:
|
|
55
|
-
return base64.b64decode(value["data"], validate=True)
|
|
56
|
-
except (binascii.Error, ValueError) as exc:
|
|
57
|
-
raise HostBridgeError("protocol_mismatch", f"response field {field} contains invalid base64") from exc
|
|
58
|
-
|
|
59
|
-
|
|
60
30
|
@dataclass(frozen=True, slots=True)
|
|
61
31
|
class SessionInfo:
|
|
62
32
|
session_id: str
|
|
@@ -83,21 +53,70 @@ class ExecResult:
|
|
|
83
53
|
stderr: bytes
|
|
84
54
|
timed_out: bool = False
|
|
85
55
|
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
|
|
91
|
-
|
|
92
|
-
|
|
93
|
-
|
|
94
|
-
|
|
95
|
-
|
|
96
|
-
|
|
97
|
-
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
|
|
56
|
+
|
|
57
|
+
class HostBridgeShell:
|
|
58
|
+
def __init__(self, stream: MuxSyncStream) -> None:
|
|
59
|
+
self._stream = stream
|
|
60
|
+
self._state_lock = threading.Lock()
|
|
61
|
+
self._closed = False
|
|
62
|
+
self._input_closed = False
|
|
63
|
+
self._terminal: ShellTerminalResult | None = None
|
|
64
|
+
|
|
65
|
+
@property
|
|
66
|
+
def closed(self) -> bool:
|
|
67
|
+
with self._state_lock:
|
|
68
|
+
return self._closed
|
|
69
|
+
|
|
70
|
+
@property
|
|
71
|
+
def exit_code(self) -> int | None:
|
|
72
|
+
with self._state_lock:
|
|
73
|
+
return None if self._terminal is None else self._terminal.exit_code
|
|
74
|
+
|
|
75
|
+
@property
|
|
76
|
+
def exit_signal(self) -> str | None:
|
|
77
|
+
with self._state_lock:
|
|
78
|
+
return None if self._terminal is None else self._terminal.signal
|
|
79
|
+
|
|
80
|
+
def receive(self, max_bytes: int = 64 * 1024) -> bytes | None:
|
|
81
|
+
if max_bytes <= 0:
|
|
82
|
+
raise ValueError("max_bytes must be positive")
|
|
83
|
+
with self._state_lock:
|
|
84
|
+
if self._closed:
|
|
85
|
+
return None
|
|
86
|
+
chunk = self._stream.receive_unbounded(max_bytes)
|
|
87
|
+
if chunk is not None:
|
|
88
|
+
return chunk
|
|
89
|
+
result = self._stream.finish(timeout=1)
|
|
90
|
+
try:
|
|
91
|
+
terminal = ShellTerminalResult.from_mapping(result)
|
|
92
|
+
except ValueError as exc:
|
|
93
|
+
raise HostBridgeError("protocol_mismatch", f"invalid shell terminal result: {exc}") from exc
|
|
94
|
+
with self._state_lock:
|
|
95
|
+
self._terminal = terminal
|
|
96
|
+
self._closed = True
|
|
97
|
+
return None
|
|
98
|
+
|
|
99
|
+
def send(self, data: bytes) -> None:
|
|
100
|
+
if not isinstance(data, bytes):
|
|
101
|
+
raise TypeError("shell data must be bytes")
|
|
102
|
+
with self._state_lock:
|
|
103
|
+
if self._closed or self._input_closed:
|
|
104
|
+
raise HostBridgeError("connection_lost", "shell input is closed")
|
|
105
|
+
self._stream.send(data)
|
|
106
|
+
|
|
107
|
+
def send_eof(self) -> None:
|
|
108
|
+
with self._state_lock:
|
|
109
|
+
if self._closed or self._input_closed:
|
|
110
|
+
return
|
|
111
|
+
self._input_closed = True
|
|
112
|
+
self._stream.send_eof()
|
|
113
|
+
|
|
114
|
+
def close(self) -> None:
|
|
115
|
+
with self._state_lock:
|
|
116
|
+
if self._closed:
|
|
117
|
+
return
|
|
118
|
+
self._closed = True
|
|
119
|
+
self._stream.reset("shell output consumer closed", timeout=1)
|
|
101
120
|
|
|
102
121
|
|
|
103
122
|
class HostBridgeError(RuntimeError):
|
|
@@ -114,10 +133,6 @@ class HostBridgeError(RuntimeError):
|
|
|
114
133
|
self.retryable = retryable
|
|
115
134
|
self.details = details or {}
|
|
116
135
|
|
|
117
|
-
@classmethod
|
|
118
|
-
def from_info(cls, error: ErrorInfo) -> HostBridgeError:
|
|
119
|
-
return cls(error.code, error.message, retryable=error.retryable, details=error.details)
|
|
120
|
-
|
|
121
136
|
|
|
122
137
|
class HostBridgeClient:
|
|
123
138
|
def __init__(
|
|
@@ -130,6 +145,25 @@ class HostBridgeClient:
|
|
|
130
145
|
raise ValueError("default_timeout must be positive")
|
|
131
146
|
self.socket_path = runtime_paths(socket_path).socket
|
|
132
147
|
self.default_timeout = float(default_timeout)
|
|
148
|
+
self._mux = MuxSyncClient(self.socket_path, default_timeout=self.default_timeout)
|
|
149
|
+
|
|
150
|
+
@property
|
|
151
|
+
def closed(self) -> bool:
|
|
152
|
+
return self._mux.closed
|
|
153
|
+
|
|
154
|
+
def close(self) -> None:
|
|
155
|
+
self._mux.close()
|
|
156
|
+
|
|
157
|
+
def __enter__(self) -> HostBridgeClient:
|
|
158
|
+
return self
|
|
159
|
+
|
|
160
|
+
def __exit__(
|
|
161
|
+
self,
|
|
162
|
+
exc_type: type[BaseException] | None,
|
|
163
|
+
exc: BaseException | None,
|
|
164
|
+
traceback: TracebackType | None,
|
|
165
|
+
) -> None:
|
|
166
|
+
self.close()
|
|
133
167
|
|
|
134
168
|
def request(
|
|
135
169
|
self,
|
|
@@ -138,19 +172,16 @@ class HostBridgeClient:
|
|
|
138
172
|
*,
|
|
139
173
|
timeout: float | None = None,
|
|
140
174
|
) -> dict[str, object]:
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
|
|
144
|
-
|
|
145
|
-
raise HostBridgeError("protocol_mismatch", "failed response omitted error details")
|
|
146
|
-
raise HostBridgeError.from_info(response.error)
|
|
147
|
-
return response.result
|
|
175
|
+
try:
|
|
176
|
+
return self._mux.request(method, params, timeout=timeout)
|
|
177
|
+
except BaseException as exc:
|
|
178
|
+
raise _translate_error(exc) from exc
|
|
148
179
|
|
|
149
180
|
def hello(self) -> dict[str, object]:
|
|
150
181
|
result = self.request("system.hello", {})
|
|
151
182
|
if (
|
|
152
183
|
result.get("service") != "hostbridge"
|
|
153
|
-
or result.get("protocol") !=
|
|
184
|
+
or result.get("protocol") != MUX_PROTOCOL_VERSION
|
|
154
185
|
or result.get("version") != __version__
|
|
155
186
|
):
|
|
156
187
|
raise HostBridgeError(
|
|
@@ -159,6 +190,25 @@ class HostBridgeClient:
|
|
|
159
190
|
)
|
|
160
191
|
return result
|
|
161
192
|
|
|
193
|
+
def list_hosts(self) -> list[dict[str, object]]:
|
|
194
|
+
hosts = self.request("host.list", {}).get("hosts")
|
|
195
|
+
if not isinstance(hosts, list) or any(not isinstance(host, dict) for host in hosts):
|
|
196
|
+
raise HostBridgeError("protocol_mismatch", "host list response must contain objects")
|
|
197
|
+
return hosts
|
|
198
|
+
|
|
199
|
+
def tunnel_health(self) -> dict[str, object]:
|
|
200
|
+
result = self.request("tunnel.health", {})
|
|
201
|
+
hosts = result.get("hosts")
|
|
202
|
+
if not isinstance(hosts, list) or any(not isinstance(host, dict) for host in hosts):
|
|
203
|
+
raise HostBridgeError("protocol_mismatch", "tunnel health response must contain host objects")
|
|
204
|
+
return result
|
|
205
|
+
|
|
206
|
+
def list_sessions(self) -> list[dict[str, object]]:
|
|
207
|
+
sessions = self.request("session.list", {}).get("sessions")
|
|
208
|
+
if not isinstance(sessions, list) or any(not isinstance(session, dict) for session in sessions):
|
|
209
|
+
raise HostBridgeError("protocol_mismatch", "session list response must contain objects")
|
|
210
|
+
return sessions
|
|
211
|
+
|
|
162
212
|
def open_session(
|
|
163
213
|
self,
|
|
164
214
|
host_id: str,
|
|
@@ -185,19 +235,17 @@ class HostBridgeClient:
|
|
|
185
235
|
output_limit: int | None = None,
|
|
186
236
|
owner: str | None = None,
|
|
187
237
|
) -> ExecResult:
|
|
188
|
-
|
|
189
|
-
"
|
|
190
|
-
|
|
191
|
-
|
|
192
|
-
|
|
193
|
-
|
|
194
|
-
|
|
195
|
-
|
|
196
|
-
|
|
197
|
-
|
|
198
|
-
|
|
199
|
-
result = self.request("exec.start", params, timeout=timeout)
|
|
200
|
-
return ExecResult.from_dict(result)
|
|
238
|
+
if stdin is not None and not isinstance(stdin, bytes):
|
|
239
|
+
raise TypeError("stdin must be bytes or None")
|
|
240
|
+
chunks: Iterable[bytes] = () if stdin is None else (stdin,)
|
|
241
|
+
return self.exec_stream(
|
|
242
|
+
session_id,
|
|
243
|
+
command,
|
|
244
|
+
chunks,
|
|
245
|
+
timeout=timeout,
|
|
246
|
+
output_limit=output_limit,
|
|
247
|
+
owner=owner,
|
|
248
|
+
)
|
|
201
249
|
|
|
202
250
|
def exec_stream(
|
|
203
251
|
self,
|
|
@@ -210,58 +258,43 @@ class HostBridgeClient:
|
|
|
210
258
|
owner: str | None = None,
|
|
211
259
|
cancel_event: threading.Event | None = None,
|
|
212
260
|
) -> ExecResult:
|
|
213
|
-
|
|
214
|
-
params: dict[str, object] = {
|
|
215
|
-
"session_id": session_id,
|
|
216
|
-
"command": command,
|
|
217
|
-
"stream_id": stream_id,
|
|
218
|
-
}
|
|
261
|
+
params: dict[str, object] = {"session_id": session_id, "command": command}
|
|
219
262
|
if timeout is not None:
|
|
220
263
|
params["timeout"] = timeout
|
|
221
264
|
if output_limit is not None:
|
|
222
265
|
params["output_limit"] = output_limit
|
|
223
266
|
if owner is not None:
|
|
224
267
|
params["owner"] = owner
|
|
225
|
-
|
|
226
|
-
|
|
227
|
-
|
|
268
|
+
deadline = _deadline(self.default_timeout if timeout is None else timeout)
|
|
269
|
+
stream = self._open_stream("exec.start", params, deadline)
|
|
270
|
+
decoder = MuxRecordDecoder()
|
|
271
|
+
stdout = bytearray()
|
|
272
|
+
stderr = bytearray()
|
|
228
273
|
try:
|
|
229
|
-
|
|
230
|
-
|
|
231
|
-
|
|
232
|
-
|
|
233
|
-
|
|
234
|
-
|
|
235
|
-
|
|
236
|
-
|
|
237
|
-
|
|
238
|
-
|
|
239
|
-
|
|
240
|
-
|
|
241
|
-
|
|
242
|
-
|
|
243
|
-
|
|
244
|
-
|
|
245
|
-
|
|
246
|
-
|
|
247
|
-
|
|
248
|
-
|
|
249
|
-
|
|
250
|
-
|
|
251
|
-
|
|
252
|
-
|
|
253
|
-
raise
|
|
254
|
-
with connection.makefile("rb") as reader:
|
|
255
|
-
response = self._read_response(reader, request)
|
|
256
|
-
except HostBridgeError:
|
|
257
|
-
raise
|
|
258
|
-
except TimeoutError as exc:
|
|
259
|
-
raise HostBridgeError("timeout", "HostBridge exec stream timed out", retryable=True) from exc
|
|
260
|
-
except OSError as exc:
|
|
261
|
-
raise HostBridgeError(
|
|
262
|
-
"connection_lost", "HostBridge exec stream connection failed", retryable=True
|
|
263
|
-
) from exc
|
|
264
|
-
return ExecResult.from_dict(response.result)
|
|
274
|
+
self._send_input(stream, stdin, deadline, cancel_event=cancel_event, label="exec stdin")
|
|
275
|
+
while chunk := stream.receive(timeout=_remaining(deadline)):
|
|
276
|
+
for record in decoder.feed(chunk):
|
|
277
|
+
target = stdout if record.channel is MuxRecordChannel.STDOUT else stderr
|
|
278
|
+
target.extend(record.payload)
|
|
279
|
+
if decoder.buffered_bytes:
|
|
280
|
+
raise HostBridgeError("protocol_mismatch", "exec output ended with an incomplete binary record")
|
|
281
|
+
result = stream.finish(timeout=_remaining(deadline))
|
|
282
|
+
except BaseException as exc:
|
|
283
|
+
self._reset_stream(stream, str(exc))
|
|
284
|
+
raise _translate_error(exc) from exc
|
|
285
|
+
exit_code = result.get("exit_code")
|
|
286
|
+
if exit_code is not None and (not isinstance(exit_code, int) or isinstance(exit_code, bool)):
|
|
287
|
+
raise HostBridgeError("protocol_mismatch", "response field exit_code must be an integer or null")
|
|
288
|
+
timed_out = result.get("timed_out", False)
|
|
289
|
+
if not isinstance(timed_out, bool):
|
|
290
|
+
raise HostBridgeError("protocol_mismatch", "response field timed_out must be a boolean")
|
|
291
|
+
return ExecResult(
|
|
292
|
+
_required_text(result.get("session_id"), "session_id"),
|
|
293
|
+
exit_code,
|
|
294
|
+
bytes(stdout),
|
|
295
|
+
bytes(stderr),
|
|
296
|
+
timed_out,
|
|
297
|
+
)
|
|
265
298
|
|
|
266
299
|
def close_session(self, session_id: str, *, owner: str | None = None) -> None:
|
|
267
300
|
params: dict[str, object] = {"session_id": session_id}
|
|
@@ -269,6 +302,9 @@ class HostBridgeClient:
|
|
|
269
302
|
params["owner"] = owner
|
|
270
303
|
self.request("session.close", params)
|
|
271
304
|
|
|
305
|
+
def release_mock_ssh_session(self, session_id: str) -> None:
|
|
306
|
+
self.request("session.release_mock_ssh", {"session_id": session_id})
|
|
307
|
+
|
|
272
308
|
def write_text(
|
|
273
309
|
self,
|
|
274
310
|
session_id: str,
|
|
@@ -288,8 +324,7 @@ class HostBridgeClient:
|
|
|
288
324
|
}
|
|
289
325
|
if owner is not None:
|
|
290
326
|
params["owner"] = owner
|
|
291
|
-
|
|
292
|
-
return self._transfer_result(result)
|
|
327
|
+
return self._transfer_result(self.request("file.write_text", params))
|
|
293
328
|
|
|
294
329
|
def read_text(
|
|
295
330
|
self,
|
|
@@ -299,28 +334,55 @@ class HostBridgeClient:
|
|
|
299
334
|
max_bytes: int = 4 * 1024 * 1024,
|
|
300
335
|
owner: str | None = None,
|
|
301
336
|
) -> str:
|
|
302
|
-
params: dict[str, object] = {
|
|
337
|
+
params: dict[str, object] = {
|
|
338
|
+
"session_id": session_id,
|
|
339
|
+
"remote_path": remote_path,
|
|
340
|
+
"max_bytes": max_bytes,
|
|
341
|
+
}
|
|
303
342
|
if owner is not None:
|
|
304
343
|
params["owner"] = owner
|
|
305
|
-
|
|
306
|
-
content = result.get("content")
|
|
344
|
+
content = self.request("file.read_text", params).get("content")
|
|
307
345
|
if not isinstance(content, str):
|
|
308
346
|
raise HostBridgeError("protocol_mismatch", "text response content must be a string")
|
|
309
347
|
return content
|
|
310
348
|
|
|
311
349
|
def shell_write(self, session_id: str, data: bytes, *, owner: str | None = None) -> int:
|
|
312
|
-
|
|
313
|
-
"
|
|
314
|
-
|
|
315
|
-
}
|
|
350
|
+
if not isinstance(data, bytes):
|
|
351
|
+
raise TypeError("shell data must be bytes")
|
|
352
|
+
params: dict[str, object] = {"session_id": session_id, "max_bytes": max(1, len(data))}
|
|
316
353
|
if owner is not None:
|
|
317
354
|
params["owner"] = owner
|
|
318
|
-
|
|
355
|
+
deadline = _deadline(self.default_timeout)
|
|
356
|
+
stream = self._open_stream("shell.write", params, deadline)
|
|
357
|
+
try:
|
|
358
|
+
self._send_input(stream, (data,), deadline, cancel_event=None, label="shell input")
|
|
359
|
+
while stream.receive(timeout=_remaining(deadline)) is not None:
|
|
360
|
+
pass
|
|
361
|
+
result = stream.finish(timeout=_remaining(deadline))
|
|
362
|
+
except BaseException as exc:
|
|
363
|
+
self._reset_stream(stream, str(exc))
|
|
364
|
+
raise _translate_error(exc) from exc
|
|
319
365
|
written = result.get("written")
|
|
320
366
|
if not isinstance(written, int) or isinstance(written, bool) or written < 0:
|
|
321
367
|
raise HostBridgeError("protocol_mismatch", "shell write response contains invalid byte count")
|
|
322
368
|
return written
|
|
323
369
|
|
|
370
|
+
def open_shell(
|
|
371
|
+
self,
|
|
372
|
+
session_id: str,
|
|
373
|
+
*,
|
|
374
|
+
owner: str | None = None,
|
|
375
|
+
max_bytes: int = 64 * 1024,
|
|
376
|
+
) -> HostBridgeShell:
|
|
377
|
+
if max_bytes <= 0:
|
|
378
|
+
raise ValueError("max_bytes must be positive")
|
|
379
|
+
params: dict[str, object] = {"session_id": session_id, "max_bytes": max_bytes}
|
|
380
|
+
if owner is not None:
|
|
381
|
+
params["owner"] = owner
|
|
382
|
+
deadline = _deadline(self.default_timeout)
|
|
383
|
+
stream = self._open_stream("shell.attach", params, deadline)
|
|
384
|
+
return HostBridgeShell(stream)
|
|
385
|
+
|
|
324
386
|
def shell_read(
|
|
325
387
|
self,
|
|
326
388
|
session_id: str,
|
|
@@ -329,18 +391,24 @@ class HostBridgeClient:
|
|
|
329
391
|
max_bytes: int = 4096,
|
|
330
392
|
owner: str | None = None,
|
|
331
393
|
) -> ShellReadResult:
|
|
332
|
-
params: dict[str, object] = {
|
|
333
|
-
"session_id": session_id,
|
|
334
|
-
"timeout": timeout,
|
|
335
|
-
"max_bytes": max_bytes,
|
|
336
|
-
}
|
|
394
|
+
params: dict[str, object] = {"session_id": session_id, "timeout": timeout, "max_bytes": max_bytes}
|
|
337
395
|
if owner is not None:
|
|
338
396
|
params["owner"] = owner
|
|
339
|
-
|
|
397
|
+
deadline = _deadline(max(self.default_timeout, timeout + 1))
|
|
398
|
+
stream = self._open_stream("shell.read", params, deadline)
|
|
399
|
+
output = bytearray()
|
|
400
|
+
try:
|
|
401
|
+
stream.send_eof(timeout=_remaining(deadline))
|
|
402
|
+
while chunk := stream.receive(timeout=_remaining(deadline)):
|
|
403
|
+
output.extend(chunk)
|
|
404
|
+
result = stream.finish(timeout=_remaining(deadline))
|
|
405
|
+
except BaseException as exc:
|
|
406
|
+
self._reset_stream(stream, str(exc))
|
|
407
|
+
raise _translate_error(exc) from exc
|
|
340
408
|
alive = result.get("alive")
|
|
341
409
|
if not isinstance(alive, bool):
|
|
342
410
|
raise HostBridgeError("protocol_mismatch", "shell read response contains invalid alive flag")
|
|
343
|
-
return ShellReadResult(
|
|
411
|
+
return ShellReadResult(bytes(output), alive)
|
|
344
412
|
|
|
345
413
|
def upload_file(
|
|
346
414
|
self,
|
|
@@ -356,55 +424,37 @@ class HostBridgeClient:
|
|
|
356
424
|
max_bytes: int | None = None,
|
|
357
425
|
) -> TransferResult:
|
|
358
426
|
self._validate_chunk_size(chunk_size)
|
|
427
|
+
self._validate_optional_byte_limit(max_bytes)
|
|
359
428
|
source = Path(local_path)
|
|
360
429
|
expected_size = source.stat().st_size
|
|
361
|
-
self._validate_optional_byte_limit(max_bytes)
|
|
362
430
|
if max_bytes is not None and expected_size > max_bytes:
|
|
363
431
|
raise HostBridgeError("resource_limit", f"upload exceeds {max_bytes} bytes")
|
|
364
|
-
stream_id = uuid.uuid4().hex
|
|
365
432
|
params: dict[str, object] = {
|
|
366
433
|
"session_id": session_id,
|
|
367
434
|
"remote_path": remote_path,
|
|
368
|
-
"stream_id": stream_id,
|
|
369
435
|
"mode": mode,
|
|
370
436
|
"expected_size": expected_size,
|
|
371
437
|
}
|
|
372
438
|
if owner is not None:
|
|
373
439
|
params["owner"] = owner
|
|
374
|
-
|
|
440
|
+
deadline = _deadline(self.default_timeout if timeout is None else timeout)
|
|
441
|
+
stream = self._open_stream("transfer.upload", params, deadline)
|
|
375
442
|
digest = hashlib.sha256()
|
|
376
|
-
|
|
377
|
-
|
|
443
|
+
|
|
444
|
+
def chunks() -> Iterable[bytes]:
|
|
445
|
+
with source.open("rb") as handle:
|
|
446
|
+
while chunk := handle.read(chunk_size):
|
|
447
|
+
digest.update(chunk)
|
|
448
|
+
yield chunk
|
|
449
|
+
|
|
378
450
|
try:
|
|
379
|
-
|
|
380
|
-
|
|
381
|
-
|
|
382
|
-
|
|
383
|
-
|
|
384
|
-
|
|
385
|
-
|
|
386
|
-
payload = json.dumps(
|
|
387
|
-
{"code": "cancelled", "message": "upload cancelled by client"},
|
|
388
|
-
separators=(",", ":"),
|
|
389
|
-
).encode("utf-8")
|
|
390
|
-
connection.sendall(
|
|
391
|
-
encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.CANCEL, payload))
|
|
392
|
-
)
|
|
393
|
-
break
|
|
394
|
-
chunk = local_file.read(chunk_size)
|
|
395
|
-
if not chunk:
|
|
396
|
-
connection.sendall(encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.EOF)))
|
|
397
|
-
break
|
|
398
|
-
digest.update(chunk)
|
|
399
|
-
connection.sendall(encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.DATA, chunk)))
|
|
400
|
-
sequence += 1
|
|
401
|
-
with connection.makefile("rb") as reader:
|
|
402
|
-
response = self._read_response(reader, request)
|
|
403
|
-
except TimeoutError as exc:
|
|
404
|
-
raise HostBridgeError("timeout", "HostBridge upload timed out", retryable=True) from exc
|
|
405
|
-
except OSError as exc:
|
|
406
|
-
raise HostBridgeError("connection_lost", "HostBridge upload connection failed", retryable=True) from exc
|
|
407
|
-
result = self._transfer_result(response.result)
|
|
451
|
+
self._send_input(stream, chunks(), deadline, cancel_event=cancel_event, label="upload")
|
|
452
|
+
while stream.receive(timeout=_remaining(deadline)) is not None:
|
|
453
|
+
pass
|
|
454
|
+
result = self._transfer_result(stream.finish(timeout=_remaining(deadline)))
|
|
455
|
+
except BaseException as exc:
|
|
456
|
+
self._reset_stream(stream, str(exc))
|
|
457
|
+
raise _translate_error(exc) from exc
|
|
408
458
|
if result.bytes_transferred != expected_size or result.sha256 != digest.hexdigest():
|
|
409
459
|
raise HostBridgeError("integrity_error", "HostBridge upload integrity check failed")
|
|
410
460
|
return result
|
|
@@ -425,66 +475,70 @@ class HostBridgeClient:
|
|
|
425
475
|
destination = Path(local_path)
|
|
426
476
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
427
477
|
temporary = destination.with_name(f".{destination.name}.hostbridge-{uuid.uuid4().hex}.tmp")
|
|
428
|
-
stream_id = uuid.uuid4().hex
|
|
429
478
|
params: dict[str, object] = {
|
|
430
479
|
"session_id": session_id,
|
|
431
480
|
"remote_path": remote_path,
|
|
432
|
-
"stream_id": stream_id,
|
|
433
481
|
"chunk_size": chunk_size,
|
|
434
482
|
}
|
|
435
483
|
if owner is not None:
|
|
436
484
|
params["owner"] = owner
|
|
437
|
-
|
|
485
|
+
deadline = _deadline(self.default_timeout if timeout is None else timeout)
|
|
486
|
+
stream = self._open_stream("transfer.download", params, deadline)
|
|
438
487
|
digest = hashlib.sha256()
|
|
439
488
|
transferred = 0
|
|
440
|
-
validator = FrameSequenceValidator()
|
|
441
|
-
request_timeout = self.default_timeout if timeout is None else timeout
|
|
442
489
|
try:
|
|
443
|
-
|
|
444
|
-
|
|
445
|
-
|
|
446
|
-
|
|
447
|
-
|
|
448
|
-
|
|
449
|
-
|
|
450
|
-
|
|
451
|
-
|
|
452
|
-
|
|
453
|
-
|
|
454
|
-
|
|
455
|
-
|
|
456
|
-
|
|
457
|
-
if frame.flags & FrameFlags.CANCEL:
|
|
458
|
-
self._raise_cancel_frame(frame)
|
|
459
|
-
if frame.flags & FrameFlags.DATA:
|
|
460
|
-
local_file.write(frame.payload)
|
|
461
|
-
digest.update(frame.payload)
|
|
462
|
-
transferred += len(frame.payload)
|
|
463
|
-
if max_bytes is not None and transferred > max_bytes:
|
|
464
|
-
raise HostBridgeError("resource_limit", f"download exceeds {max_bytes} bytes")
|
|
465
|
-
if frame.flags & FrameFlags.EOF:
|
|
466
|
-
metadata = self._frame_metadata(frame.payload)
|
|
467
|
-
break
|
|
468
|
-
except HostBridgeError:
|
|
469
|
-
temporary.unlink(missing_ok=True)
|
|
470
|
-
raise
|
|
471
|
-
except TimeoutError as exc:
|
|
490
|
+
stream.send_eof(timeout=_remaining(deadline))
|
|
491
|
+
with temporary.open("wb") as handle:
|
|
492
|
+
while chunk := stream.receive(timeout=_remaining(deadline)):
|
|
493
|
+
transferred += len(chunk)
|
|
494
|
+
if max_bytes is not None and transferred > max_bytes:
|
|
495
|
+
raise HostBridgeError("resource_limit", f"download exceeds {max_bytes} bytes")
|
|
496
|
+
handle.write(chunk)
|
|
497
|
+
digest.update(chunk)
|
|
498
|
+
result = self._transfer_result(stream.finish(timeout=_remaining(deadline)))
|
|
499
|
+
if result.bytes_transferred != transferred or result.sha256 != digest.hexdigest():
|
|
500
|
+
raise HostBridgeError("integrity_error", "HostBridge download integrity check failed")
|
|
501
|
+
temporary.replace(destination)
|
|
502
|
+
return result
|
|
503
|
+
except BaseException as exc:
|
|
472
504
|
temporary.unlink(missing_ok=True)
|
|
473
|
-
|
|
474
|
-
|
|
475
|
-
|
|
476
|
-
|
|
477
|
-
|
|
478
|
-
|
|
479
|
-
|
|
480
|
-
raise
|
|
481
|
-
|
|
482
|
-
|
|
505
|
+
self._reset_stream(stream, str(exc))
|
|
506
|
+
raise _translate_error(exc) from exc
|
|
507
|
+
|
|
508
|
+
def _open_stream(self, method: str, params: Mapping[str, object], deadline: float) -> MuxSyncStream:
|
|
509
|
+
try:
|
|
510
|
+
return self._mux.open_stream(method, params, timeout=_remaining(deadline))
|
|
511
|
+
except BaseException as exc:
|
|
512
|
+
raise _translate_error(exc) from exc
|
|
513
|
+
|
|
514
|
+
@staticmethod
|
|
515
|
+
def _send_input(
|
|
516
|
+
stream: MuxSyncStream,
|
|
517
|
+
chunks: Iterable[bytes],
|
|
518
|
+
deadline: float,
|
|
519
|
+
*,
|
|
520
|
+
cancel_event: threading.Event | None,
|
|
521
|
+
label: str,
|
|
522
|
+
) -> None:
|
|
523
|
+
for chunk in chunks:
|
|
524
|
+
if cancel_event is not None and cancel_event.is_set():
|
|
525
|
+
raise HostBridgeError("cancelled", f"{label} cancelled")
|
|
526
|
+
if not isinstance(chunk, bytes):
|
|
527
|
+
raise TypeError(f"{label} chunks must be bytes")
|
|
528
|
+
stream.send(chunk, timeout=_remaining(deadline))
|
|
529
|
+
if cancel_event is not None and cancel_event.is_set():
|
|
530
|
+
raise HostBridgeError("cancelled", f"{label} cancelled")
|
|
531
|
+
stream.send_eof(timeout=_remaining(deadline))
|
|
532
|
+
|
|
533
|
+
@staticmethod
|
|
534
|
+
def _reset_stream(stream: MuxSyncStream, message: str) -> None:
|
|
535
|
+
with suppress(BaseException):
|
|
536
|
+
stream.reset(message or "client operation failed", timeout=1)
|
|
483
537
|
|
|
484
538
|
@staticmethod
|
|
485
539
|
def _validate_chunk_size(chunk_size: int) -> None:
|
|
486
|
-
if not isinstance(chunk_size, int) or isinstance(chunk_size, bool) or not 4096 <= chunk_size <=
|
|
487
|
-
raise ValueError(
|
|
540
|
+
if not isinstance(chunk_size, int) or isinstance(chunk_size, bool) or not 4096 <= chunk_size <= 1024 * 1024:
|
|
541
|
+
raise ValueError("chunk_size must be between 4096 and 1048576 bytes")
|
|
488
542
|
|
|
489
543
|
@staticmethod
|
|
490
544
|
def _validate_optional_byte_limit(max_bytes: int | None) -> None:
|
|
@@ -501,70 +555,31 @@ class HostBridgeClient:
|
|
|
501
555
|
raise HostBridgeError("protocol_mismatch", "transfer response contains invalid SHA-256")
|
|
502
556
|
return TransferResult(byte_count, sha256.lower())
|
|
503
557
|
|
|
504
|
-
def _read_response(self, reader, request: Request) -> Response:
|
|
505
|
-
raw_response = reader.readline(MAX_CONTROL_LINE_BYTES + 1)
|
|
506
|
-
if not raw_response:
|
|
507
|
-
raise HostBridgeError("connection_lost", "HostBridge daemon closed before responding", retryable=True)
|
|
508
|
-
if len(raw_response) > MAX_CONTROL_LINE_BYTES:
|
|
509
|
-
raise HostBridgeError("protocol_mismatch", "HostBridge daemon response exceeds the control limit")
|
|
510
|
-
try:
|
|
511
|
-
response = Response.from_json(raw_response.decode("utf-8"))
|
|
512
|
-
except (ProtocolError, UnicodeDecodeError) as exc:
|
|
513
|
-
raise HostBridgeError("protocol_mismatch", "HostBridge daemon returned an invalid response") from exc
|
|
514
|
-
if response.request_id != request.request_id:
|
|
515
|
-
raise HostBridgeError("protocol_mismatch", "HostBridge daemon response request_id does not match")
|
|
516
|
-
if not response.ok:
|
|
517
|
-
if response.error is None:
|
|
518
|
-
raise HostBridgeError("protocol_mismatch", "failed response omitted error details")
|
|
519
|
-
raise HostBridgeError.from_info(response.error)
|
|
520
|
-
return response
|
|
521
|
-
|
|
522
|
-
@staticmethod
|
|
523
|
-
def _frame_metadata(payload: bytes) -> dict[str, object]:
|
|
524
|
-
try:
|
|
525
|
-
value = json.loads(payload.decode("utf-8"))
|
|
526
|
-
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
|
527
|
-
raise HostBridgeError("protocol_mismatch", "download EOF contains invalid metadata") from exc
|
|
528
|
-
if not isinstance(value, dict):
|
|
529
|
-
raise HostBridgeError("protocol_mismatch", "download EOF metadata must be an object")
|
|
530
|
-
return value
|
|
531
|
-
|
|
532
|
-
@staticmethod
|
|
533
|
-
def _raise_cancel_frame(frame: BinaryFrame) -> None:
|
|
534
|
-
metadata = HostBridgeClient._frame_metadata(frame.payload)
|
|
535
|
-
code = metadata.get("code")
|
|
536
|
-
message = metadata.get("message")
|
|
537
|
-
raise HostBridgeError(
|
|
538
|
-
code if isinstance(code, str) else "cancelled",
|
|
539
|
-
message if isinstance(message, str) else "HostBridge transfer cancelled",
|
|
540
|
-
)
|
|
541
|
-
|
|
542
|
-
def _exchange(self, request: Request, *, timeout: float) -> Response:
|
|
543
|
-
if timeout <= 0:
|
|
544
|
-
raise ValueError("timeout must be positive")
|
|
545
|
-
try:
|
|
546
|
-
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
|
|
547
|
-
connection.settimeout(timeout)
|
|
548
|
-
connection.connect(str(self.socket_path))
|
|
549
|
-
connection.sendall(request.to_json().encode("utf-8") + b"\n")
|
|
550
|
-
with connection.makefile("rb") as reader:
|
|
551
|
-
raw_response = reader.readline(MAX_CONTROL_LINE_BYTES + 1)
|
|
552
|
-
except TimeoutError as exc:
|
|
553
|
-
raise HostBridgeError("timeout", "HostBridge daemon request timed out", retryable=True) from exc
|
|
554
|
-
except OSError as exc:
|
|
555
|
-
raise HostBridgeError("connection_lost", "HostBridge daemon connection failed", retryable=True) from exc
|
|
556
|
-
if not raw_response:
|
|
557
|
-
raise HostBridgeError("connection_lost", "HostBridge daemon closed before responding", retryable=True)
|
|
558
|
-
if len(raw_response) > MAX_CONTROL_LINE_BYTES:
|
|
559
|
-
raise HostBridgeError("protocol_mismatch", "HostBridge daemon response exceeds the control limit")
|
|
560
|
-
return self._read_response_from_bytes(raw_response, request)
|
|
561
558
|
|
|
562
|
-
|
|
563
|
-
|
|
564
|
-
|
|
565
|
-
|
|
566
|
-
|
|
567
|
-
|
|
568
|
-
|
|
569
|
-
|
|
570
|
-
|
|
559
|
+
def _deadline(timeout: float) -> float:
|
|
560
|
+
if timeout <= 0:
|
|
561
|
+
raise ValueError("timeout must be positive")
|
|
562
|
+
return time.monotonic() + timeout
|
|
563
|
+
|
|
564
|
+
|
|
565
|
+
def _remaining(deadline: float) -> float:
|
|
566
|
+
remaining = deadline - time.monotonic()
|
|
567
|
+
if remaining <= 0:
|
|
568
|
+
raise HostBridgeError("timeout", "HostBridge operation timed out", retryable=True)
|
|
569
|
+
return remaining
|
|
570
|
+
|
|
571
|
+
|
|
572
|
+
def _translate_error(error: BaseException) -> BaseException:
|
|
573
|
+
if isinstance(error, MuxRecordError):
|
|
574
|
+
return HostBridgeError("protocol_mismatch", str(error))
|
|
575
|
+
if isinstance(error, HostBridgeError | TypeError | ValueError):
|
|
576
|
+
return error
|
|
577
|
+
if isinstance(error, MuxRemoteError):
|
|
578
|
+
return HostBridgeError(error.code, str(error), retryable=error.retryable, details=error.details)
|
|
579
|
+
if isinstance(error, MuxCapacityError):
|
|
580
|
+
return HostBridgeError("busy", str(error), retryable=True)
|
|
581
|
+
if isinstance(error, MuxConnectionClosed | RuntimeError):
|
|
582
|
+
return HostBridgeError("connection_lost", str(error), retryable=True)
|
|
583
|
+
if isinstance(error, TimeoutError):
|
|
584
|
+
return HostBridgeError("timeout", str(error), retryable=True)
|
|
585
|
+
return error
|