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