"""Interactive pseudo-terminal lifecycle and WebSocket streaming.""" import base64 import hashlib import json import secrets import socket import struct import time import unittest import urllib.parse import urllib.request from harness import Albedo, exclusive from inspect_support import Inspection class TerminalWebSocket: def __init__(self, app, terminal_id, timeout=20, **params): address = urllib.parse.urlsplit(app.base) self.socket = socket.create_connection( (address.hostname, address.port), timeout ) self.reader = self.socket.makefile("rb") key = base64.b64encode(secrets.token_bytes(16)).decode() query = urllib.parse.urlencode(params) path = f"/terminals/{terminal_id}" + (f"?{query}" if query else "") self.socket.sendall( ( f"GET {path} HTTP/1.1\r\nHost: {address.netloc}\r\n" "Connection: Upgrade\r\nUpgrade: websocket\r\n" f"Sec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n" f"Authorization: Bearer {app.connection['token']}\r\n\r\n" ).encode() ) try: status = self.reader.readline() if b" 101 " not in status: raise AssertionError(f"WebSocket upgrade failed: {status!r}") headers = {} while (line := self.reader.readline()) not in (b"\r\n", b""): name, value = line.decode().split(":", 1) headers[name.lower()] = value.strip() expected = base64.b64encode( hashlib.sha1( (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode() ).digest() ).decode() if headers.get("sec-websocket-accept") != expected: raise AssertionError(f"bad WebSocket accept: {headers}") except BaseException: self.close() raise def close(self): try: self.reader.close() self.socket.close() except OSError: pass def frame(self, opcode, payload): mask = secrets.token_bytes(4) size = len(payload) header = bytes([0x80 | opcode]) if size < 126: header += bytes([0x80 | size]) elif size < 65536: header += bytes([0x80 | 126]) + struct.pack("!H", size) else: header += bytes([0x80 | 127]) + struct.pack("!Q", size) self.socket.sendall( header + mask + bytes(value ^ mask[i % 4] for i, value in enumerate(payload)) ) def send_text(self, **message): self.frame(1, json.dumps(message).encode("utf-8")) def send_bytes(self, data): self.frame(2, data) def read_exact(self, size): value = self.reader.read(size) if len(value) != size: raise AssertionError("WebSocket closed before a complete frame") return value def receive_frame(self, timeout=10): self.socket.settimeout(timeout) flags, length = self.read_exact(2) opcode = flags & 15 if length & 128: raise AssertionError("server sent a masked frame") if length == 126: length = struct.unpack("!H", self.read_exact(2))[0] elif length == 127: length = struct.unpack("!Q", self.read_exact(8))[0] payload = self.read_exact(length) if opcode == 9: self.frame(10, payload) return self.receive_frame(timeout) if opcode == 10: return self.receive_frame(timeout) return opcode, payload class TerminalsTest(unittest.TestCase): def setUp(self): self.app = Albedo(providers={}).__enter__() self.addCleanup(self.app.__exit__, None, None, None) def test_terminal_session_io_and_lifecycle(self): tid = "term-test-1" ws = TerminalWebSocket(self.app, tid, cols=80, rows=24, shell="/bin/sh") self.addCleanup(ws.close) # 1. Expect initial ready frame (text) opcode, payload = ws.receive_frame() self.assertEqual(opcode, 1) ready = json.loads(payload.decode("utf-8")) self.assertEqual(ready["type"], "ready", f"terminal returned error: {ready}") self.assertEqual(ready["id"], tid) self.assertEqual(ready["cols"], 80) self.assertEqual(ready["rows"], 24) self.assertEqual(ready["state"], "running") # 2. Send command to PTY ws.send_bytes(b"echo ALBEDO_TERM_OK\r\n") # 3. Read output frames until our string appears output = bytearray() deadline = time.monotonic() + 10.0 while time.monotonic() < deadline: opcode, chunk = ws.receive_frame(timeout=5) if opcode == 2: output.extend(chunk) if b"ALBEDO_TERM_OK" in output: break self.assertIn(b"ALBEDO_TERM_OK", output) # 4. Test resize command ws.send_text(type="resize", cols=100, rows=30) ws.send_bytes(b"stty size\r\n") deadline = time.monotonic() + 10.0 size_output = bytearray() while time.monotonic() < deadline: opcode, chunk = ws.receive_frame(timeout=5) if opcode == 2: size_output.extend(chunk) if b"30 100" in size_output: break self.assertIn(b"30 100", size_output) # 5. Test signal (SIGINT stops running cat) ws.send_bytes(b"cat\r\n") time.sleep(0.05) ws.send_text(type="signal", signal="SIGINT") ws.send_bytes(b"echo AFTER_SIGINT\r\n") deadline = time.monotonic() + 10.0 sig_output = bytearray() while time.monotonic() < deadline: opcode, chunk = ws.receive_frame(timeout=5) if opcode == 2: sig_output.extend(chunk) if b"AFTER_SIGINT" in sig_output: break self.assertIn(b"AFTER_SIGINT", sig_output) ws.close() # 6. Reconnect to the same terminal: scrollback history should be preserved ws2 = TerminalWebSocket(self.app, tid, shell="/bin/sh") self.addCleanup(ws2.close) opcode, payload = ws2.receive_frame() self.assertEqual(opcode, 1) ready2 = json.loads(payload.decode("utf-8")) self.assertEqual(ready2["id"], tid) # Following frame is binary history opcode, history = ws2.receive_frame() self.assertEqual(opcode, 2) self.assertIn(b"ALBEDO_TERM_OK", history) # 7. Test in-stream restart observer = TerminalWebSocket(self.app, tid, shell="/bin/sh") self.addCleanup(observer.close) self.assertEqual(observer.receive_frame()[0], 1) self.assertEqual(observer.receive_frame()[0], 2) ws2.send_text(type="restart", cols=80, rows=24) while True: opcode, payload = ws2.receive_frame(timeout=5) if opcode == 1: break restarted = json.loads(payload.decode("utf-8")) self.assertEqual(restarted["type"], "ready") self.assertEqual(restarted["id"], tid) while True: opcode, payload = observer.receive_frame(timeout=5) if opcode == 1: observer_event = json.loads(payload.decode("utf-8")) break self.assertEqual(observer_event["type"], "closed") ws2.close() # 8. Explicit delete endpoint req = urllib.request.Request( f"{self.app.base}/terminals/{tid}", headers={"Authorization": f"Bearer {self.app.connection['token']}"}, method="DELETE", ) with urllib.request.urlopen(req, timeout=5) as resp: self.assertEqual(resp.status, 200) res = json.load(resp) self.assertEqual(res["closed"], tid) # 9. Connecting after deletion gives fresh terminal with no old history ws3 = TerminalWebSocket(self.app, tid, shell="/bin/sh") self.addCleanup(ws3.close) opcode, payload = ws3.receive_frame() self.assertEqual(opcode, 1) ready3 = json.loads(payload.decode("utf-8")) self.assertEqual(ready3["type"], "ready") ws3.send_bytes(b"echo FRESH_AFTER_DELETE\r\n") fresh_out = bytearray() deadline = time.monotonic() + 10.0 while time.monotonic() < deadline: opcode, chunk = ws3.receive_frame(timeout=5) if opcode == 2: fresh_out.extend(chunk) if b"FRESH_AFTER_DELETE" in fresh_out: break self.assertIn(b"FRESH_AFTER_DELETE", fresh_out) self.assertNotIn(b"AFTER_SIGINT", fresh_out) ws3.close() def test_reconnect_sends_ready_before_live_output(self): tid = "term-reconnect-order" ws = TerminalWebSocket(self.app, tid, shell="/bin/sh") self.addCleanup(ws.close) opcode, _ = ws.receive_frame() self.assertEqual(opcode, 1) for attempt in range(16): ws.send_bytes( b"i=0; while [ $i -lt 10000 ]; do echo RECONNECT_OUTPUT; " b"i=$((i+1)); done\r\n" ) output = bytearray() while output.count(b"RECONNECT_OUTPUT") < 2: opcode, payload = ws.receive_frame(timeout=1) if opcode == 2: output.extend(payload) ws.close() ws = TerminalWebSocket(self.app, tid, shell="/bin/sh") self.addCleanup(ws.close) opcode, payload = ws.receive_frame() self.assertEqual( opcode, 1, f"output arrived before ready on attach {attempt}: {payload[:80]!r}", ) ready = json.loads(payload.decode("utf-8")) self.assertEqual(ready["type"], "ready") req = urllib.request.Request( f"{self.app.base}/terminals/{tid}", headers={"Authorization": f"Bearer {self.app.connection['token']}"}, method="DELETE", ) with urllib.request.urlopen(req, timeout=5) as resp: self.assertEqual(resp.status, 200) # exclusive: ALBEDO_INSPECT must be set before the daemon boots so the test can # read terminal mailbox lengths over erlang distribution (erlang:process_info) @exclusive class TerminalPressureTest(unittest.TestCase): def setUp(self): self.app = Albedo( providers={}, prepare=lambda app: app.daemon.env.update(ALBEDO_INSPECT="1"), ).__enter__() self.addCleanup(self.app.__exit__, None, None, None) def mailbox_lengths(self): """Queue lengths of every terminal worker and the sockets watching it.""" return json.loads( Inspection(self.app).evaluate( timeout=10, expression="Length = fun(P) -> case Call(erlang, process_info, " "[P, message_queue_len]) of {message_queue_len, N} -> N; " "_ -> 0 end end, " "Monitored = fun(P) -> case Call(erlang, process_info, " "[P, monitors]) of {monitors, Ms} -> " "[Q || {process, Q} <- Ms, is_pid(Q)]; _ -> [] end end, " "Manager = Call(erlang, whereis, [albedo_terminals_manager]), " "Workers = Monitored(Manager), " 'io:format("~w", [[[Length(W) | [Length(P) || P <- Monitored(W)]]' " || W <- Workers]])", ) ) def test_a_flooding_terminal_stays_bounded_and_accepts_new_sockets(self): tid = "term-pressure" unread = TerminalWebSocket(self.app, tid, timeout=5, shell="/bin/sh") self.addCleanup(unread.close) self.assertEqual(unread.receive_frame(timeout=5)[0], 1) unread.send_bytes(b"yes PRESSURE_OUTPUT\r\n") self.assertEqual(unread.receive_frame(timeout=5)[0], 2) # From here the unread socket stops reading. Its send buffer fills # within a moment, after which an unbounded worker or socket mailbox # only grows; the first sample already sees it. for _ in range(3): time.sleep(1) lengths = self.mailbox_lengths() self.assertTrue(lengths, "no terminal worker found") self.assertLessEqual( max(max(row) for row in lengths), 100, f"terminal mailboxes [worker, watchers...] grew: {lengths}", ) # Each new socket must get through Mist's 500 ms init window while # the producer keeps writing. deadline = time.monotonic() + 15 for attempt in range(16): self.assertLess(time.monotonic(), deadline, f"attach {attempt} too slow") ws = TerminalWebSocket(self.app, tid, timeout=3, shell="/bin/sh") self.addCleanup(ws.close) opcode, payload = ws.receive_frame(timeout=3) self.assertEqual(opcode, 1, f"attach {attempt}: {payload[:80]!r}") self.assertEqual(json.loads(payload)["type"], "ready") ws.close() # Resize and signal reach the PTY ahead of the flood: Ctrl-C stops # the producer and the shell reports the new size. control = TerminalWebSocket(self.app, tid, timeout=3, shell="/bin/sh") self.addCleanup(control.close) self.assertEqual(control.receive_frame(timeout=3)[0], 1) control.send_text(type="resize", cols=100, rows=30) control.send_text(type="signal", signal="SIGINT") control.send_bytes(b"stty size\r\n") self.assertIn(b"30 100", self.read_until(control, b"30 100")) # Once the unread socket reads again it is told what it missed. self.assertIn(b"[albedo: skipped", self.read_until(unread, b"[albedo: skipped")) req = urllib.request.Request( f"{self.app.base}/terminals/{tid}", headers={"Authorization": f"Bearer {self.app.connection['token']}"}, method="DELETE", ) with urllib.request.urlopen(req, timeout=5) as resp: self.assertEqual(resp.status, 200) def read_until(self, ws, marker, seconds=10): """Read output frames until marker appears; return the last frames read.""" window = b"" deadline = time.monotonic() + seconds while marker not in window and time.monotonic() < deadline: opcode, payload = ws.receive_frame(timeout=3) if opcode == 2: window = window[-len(marker) :] + payload return window if __name__ == "__main__": unittest.main()