1"""Local side of a portal session.
2
3Runs on the developer's machine. `connect` establishes the encrypted
4tunnel through the relay and listens on a Unix socket. Subsequent
5commands (exec, pull, push) talk to the connect process over the socket.
6"""
7
8import asyncio
9import fcntl
10import functools
11import json
12import os
13import signal
14import struct
15import sys
16from collections.abc import Callable
17
18import websockets.exceptions
19from websockets.asyncio.client import connect as ws_connect
20
21from .codegen import validate_code
22from .crypto import channel_id, perform_key_exchange
23from .protocol import (
24 DEFAULT_RELAY_HOST,
25 make_ping,
26 make_relay_url,
27)
28
29
30@functools.lru_cache
31def _portal_dir() -> str:
32 """Return .plain/portal/ in the project root, creating it if needed."""
33 from plain.runtime import PLAIN_TEMP_PATH
34
35 d = os.path.join(PLAIN_TEMP_PATH, "portal")
36 os.makedirs(d, exist_ok=True)
37 return d
38
39
40def _socket_path() -> str:
41 return os.path.join(_portal_dir(), "portal.sock")
42
43
44def _lock_path() -> str:
45 return os.path.join(_portal_dir(), "portal.lock")
46
47
48_lock_fd = None
49
50
51async def _send_framed(writer: asyncio.StreamWriter, data: bytes) -> None:
52 """Write a length-prefixed message to a stream."""
53 writer.write(struct.pack("!I", len(data)))
54 writer.write(data)
55 await writer.drain()
56
57
58# 75MB — large enough for 50MB files base64-encoded (~67MB), prevents unbounded allocation
59_MAX_FRAME_SIZE = 75 * 1024 * 1024
60
61
62async def _recv_framed(reader: asyncio.StreamReader) -> bytes:
63 """Read a length-prefixed message from a stream."""
64 length_bytes = await reader.readexactly(4)
65 length = struct.unpack("!I", length_bytes)[0]
66 if length > _MAX_FRAME_SIZE:
67 raise ValueError(f"Frame too large: {length} bytes (max {_MAX_FRAME_SIZE})")
68 return await reader.readexactly(length)
69
70
71async def connect(
72 code: str,
73 *,
74 relay_host: str = DEFAULT_RELAY_HOST,
75) -> None:
76 """Connect to a remote portal session and run the daemon."""
77
78 if not validate_code(code):
79 print(f"Invalid portal code: {code}", file=sys.stderr)
80 sys.exit(1)
81
82 # Acquire an exclusive file lock before anything else. Holds for the
83 # lifetime of the process — released automatically on exit/crash.
84 # Stored at module level to prevent GC from closing the fd.
85 global _lock_fd
86 _lock_fd = open(_lock_path(), "w") # noqa: SIM115, ASYNC230 — held for the process lifetime
87 try:
88 fcntl.flock(_lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
89 except OSError:
90 _lock_fd.close()
91 _lock_fd = None
92 print("A portal session is already active.", file=sys.stderr)
93 sys.exit(1)
94
95 # Clean up any stale socket from a previous crash.
96 _cleanup()
97
98 cid = channel_id(code)
99 relay_url = make_relay_url(relay_host, cid, "connect")
100
101 try:
102 ws = await ws_connect(relay_url)
103 except Exception as e:
104 print(f"Failed to connect to relay: {e}", file=sys.stderr)
105 sys.exit(1)
106
107 encryptor = await perform_key_exchange(ws, code, side="connect")
108
109 print("Connected to remote. Session active.")
110
111 # Exec requests use queues (for streaming exec_stdout + exec_result).
112 # All other request types use single-shot futures.
113 pending_responses: dict[int, asyncio.Future] = {}
114 pending_queues: dict[int, asyncio.Queue] = {}
115 file_data_accumulators: dict[int, dict] = {}
116 request_counter = 0
117
118 async def handle_local_client(
119 reader: asyncio.StreamReader, writer: asyncio.StreamWriter
120 ) -> None:
121 """Handle a command from a local CLI invocation (exec/pull/push)."""
122 nonlocal request_counter
123 req_id = None
124 is_exec = False
125
126 try:
127 data = await _recv_framed(reader)
128 request = json.loads(data.decode("utf-8"))
129
130 request_counter += 1
131 req_id = request_counter
132 request["_req_id"] = req_id
133 is_exec = request.get("type") == "exec"
134
135 if is_exec:
136 # Exec uses a queue so we can stream exec_stdout messages
137 queue: asyncio.Queue = asyncio.Queue()
138 pending_queues[req_id] = queue
139 await ws.send(encryptor.encrypt_message(request))
140
141 # Read from the queue until we get the final exec_result
142 exec_timeout = request.get("timeout", 120) + 30 # extra margin
143 while True:
144 msg = await asyncio.wait_for(queue.get(), timeout=exec_timeout)
145 await _send_framed(writer, json.dumps(msg).encode("utf-8"))
146 if msg.get("type") != "exec_stdout":
147 break
148 else:
149 # Non-exec: single request/response via future
150 future: asyncio.Future = asyncio.get_running_loop().create_future()
151 pending_responses[req_id] = future
152 await ws.send(encryptor.encrypt_message(request))
153 response = await asyncio.wait_for(future, timeout=300)
154 await _send_framed(writer, json.dumps(response).encode("utf-8"))
155
156 except TimeoutError:
157 await _send_framed(
158 writer,
159 json.dumps({"error": "Request timed out"}).encode("utf-8"),
160 )
161 except Exception as e:
162 await _send_framed(writer, json.dumps({"error": str(e)}).encode("utf-8"))
163 finally:
164 if req_id is not None:
165 pending_responses.pop(req_id, None)
166 pending_queues.pop(req_id, None)
167 file_data_accumulators.pop(req_id, None)
168 writer.close()
169 await writer.wait_closed()
170
171 async def relay_listener() -> None:
172 """Listen for messages from the remote side via WebSocket."""
173 try:
174 async for raw in ws:
175 if isinstance(raw, str):
176 continue
177
178 msg = encryptor.decrypt_message(raw)
179 msg_type = msg.get("type")
180
181 if msg_type == "ping":
182 await ws.send(encryptor.encrypt_message({"type": "pong"}))
183 continue
184
185 if msg_type == "pong":
186 continue
187
188 req_id = msg.pop("_req_id", None)
189 if not req_id:
190 continue
191
192 # Streaming exec messages go through the queue
193 if msg_type in ("exec_stdout", "exec_result"):
194 if req_id in pending_queues:
195 await pending_queues[req_id].put(msg)
196 continue
197
198 # File data accumulation (multiple chunks → single response)
199 if msg_type == "file_data":
200 if req_id not in pending_responses:
201 continue
202 if req_id not in file_data_accumulators:
203 file_data_accumulators[req_id] = {
204 "name": msg["name"],
205 "chunks": msg["chunks"],
206 "received": {},
207 }
208 acc = file_data_accumulators[req_id]
209 acc["received"][msg["chunk"]] = msg["data"]
210 if len(acc["received"]) == acc["chunks"]:
211 all_data = "".join(
212 acc["received"][i] for i in range(acc["chunks"])
213 )
214 del file_data_accumulators[req_id]
215 pending_responses[req_id].set_result(
216 {
217 "type": "file_data",
218 "name": acc["name"],
219 "data": all_data,
220 }
221 )
222 continue
223
224 # Everything else resolves the future directly
225 if req_id in pending_responses:
226 pending_responses[req_id].set_result(msg)
227
228 except websockets.exceptions.ConnectionClosed:
229 pass
230 finally:
231 for future in pending_responses.values():
232 if not future.done():
233 future.set_result({"error": "Remote disconnected"})
234 for queue in pending_queues.values():
235 await queue.put({"type": "error", "error": "Remote disconnected"})
236 _cleanup()
237
238 # Set restrictive umask so the socket is created owner-only (no TOCTOU window)
239 old_umask = os.umask(0o177)
240 try:
241 server = await asyncio.start_unix_server(
242 handle_local_client, path=_socket_path()
243 )
244 finally:
245 os.umask(old_umask)
246
247 loop = asyncio.get_running_loop()
248
249 def _handle_signal() -> None:
250 _cleanup()
251 loop.stop()
252
253 loop.add_signal_handler(signal.SIGTERM, _handle_signal)
254 loop.add_signal_handler(signal.SIGINT, _handle_signal)
255
256 async def send_keepalive_pings() -> None:
257 while True:
258 await asyncio.sleep(30)
259 await ws.send(encryptor.encrypt_message(make_ping()))
260
261 keepalive_task = asyncio.create_task(send_keepalive_pings())
262
263 try:
264 await relay_listener()
265 finally:
266 keepalive_task.cancel()
267 server.close()
268 await server.wait_closed()
269 _cleanup()
270
271
272def _cleanup() -> None:
273 """Remove the socket file. The lock file is left in place — the flock
274 is on the inode, so unlinking it would let a new process acquire a
275 lock on a different inode."""
276 try:
277 os.unlink(_socket_path())
278 except FileNotFoundError:
279 pass
280
281
282async def send_command(request: dict) -> dict:
283 """Send a command to the connect process via Unix socket.
284
285 Returns a single response. For streaming exec, use send_exec_streaming instead.
286 """
287 try:
288 reader, writer = await asyncio.open_unix_connection(_socket_path())
289 except FileNotFoundError, ConnectionRefusedError:
290 print(
291 "No active portal session. Run 'plain portal connect <code>' first.",
292 file=sys.stderr,
293 )
294 sys.exit(1)
295
296 try:
297 await _send_framed(writer, json.dumps(request).encode("utf-8"))
298 response_data = await _recv_framed(reader)
299 return json.loads(response_data.decode("utf-8"))
300 finally:
301 writer.close()
302 await writer.wait_closed()
303
304
305async def send_exec_streaming(
306 request: dict,
307 on_stdout: Callable[[str], None],
308) -> dict:
309 """Send an exec request and stream stdout chunks as they arrive.
310
311 Calls on_stdout(data) for each exec_stdout chunk.
312 Returns the final exec_result response.
313 """
314 try:
315 reader, writer = await asyncio.open_unix_connection(_socket_path())
316 except FileNotFoundError, ConnectionRefusedError:
317 print(
318 "No active portal session. Run 'plain portal connect <code>' first.",
319 file=sys.stderr,
320 )
321 sys.exit(1)
322
323 try:
324 await _send_framed(writer, json.dumps(request).encode("utf-8"))
325 while True:
326 response_data = await _recv_framed(reader)
327 msg = json.loads(response_data.decode("utf-8"))
328 if msg.get("type") == "exec_stdout":
329 on_stdout(msg["data"])
330 else:
331 return msg
332 finally:
333 writer.close()
334 await writer.wait_closed()