///|
fn is_auth_handshake_rejection(code : Int) -> Bool {
code == 401 || code == 403
}
///|
/// Raised by `connect_and_register` when the relay rejects the
/// Register with `proto_version_too_old`. The payload is the relay's
/// `min_version`. Tagged as fatal by the retry loop — no amount of
/// retrying will help; the operator must upgrade this tun-server
/// binary to at least that version.
priv suberror ProtoVersionTooOldRejection {
ProtoVersionTooOldRejection(Int)
} derive(Debug)
///|
async fn send_active_message(
bridge : BridgeState,
msg : String,
phase : String,
send : async () -> Unit,
) -> Unit {
send() catch {
err =>
if bridge.shell_exited.val {
@log.warn(
"[relay.send.fail] phase=\"\{phase}\" suppressed=\"shell_exited\" bytes=\{msg.length()} error=\{@debug.to_string(err)}",
)
} else {
@log.error(
"[relay.send.fail] phase=\"\{phase}\" bytes=\{msg.length()} error=\{@debug.to_string(err)}",
)
raise err
}
}
}
///|
async fn send_shutdown_message(
msg : String,
phase : String,
send : async () -> Unit,
) -> Unit {
send() catch {
err =>
@log.warn(
"[relay.send.fail] phase=\"\{phase}\" suppressed=\"shutdown\" bytes=\{msg.length()} error=\{@debug.to_string(err)}",
)
}
}
///|
/// Connect to the relay and send the initial Register message. Returns
/// an open WebSocket only after the relay has accepted the
/// registration with `Registered(...)`. Any other response is treated
/// as an error so the caller can retry or fail before touching tty/pty
/// state.
async fn connect_and_register(config : ConnectConfig) -> @ws.Conn {
@log.info("Connecting to relay: \{config.url}")
let ws = @ws.Conn::connect(config.url.to_string(), headers=config.headers)
let register_msg : @proto.Message = Sys(
Register(
task_uuid=config.task_uuid,
hostname=config.hostname,
cwd=config.cwd,
proto_version=@proto.ProtoVersion,
),
)
ws.send_text(register_msg.to_json().stringify()) catch {
err => {
ws.close()
raise err
}
}
@log.info("Sent: register")
let response = ws.recv().read_all().text() catch {
err => {
ws.close()
raise err
}
}
@log.info("Received: \{response}")
let msg : @proto.Message = @json.from_json(@json.parse(response)) catch {
err => {
ws.close()
raise err
}
}
match msg {
Sys(Registered(task_uuid=echoed_uuid)) => {
if echoed_uuid != config.task_uuid {
ws.close()
fail(
"Registered echoed wrong uuid: got \{echoed_uuid}, expected \{config.task_uuid}",
)
}
@log.info("Registered task_uuid=\{config.task_uuid}")
ws
}
Sys(RegisterRejected(reason=ProtoVersionTooOld(min_version~))) => {
ws.close()
raise ProtoVersionTooOldRejection(min_version)
}
_ => {
ws.close()
fail("unexpected register response: \{@debug.to_string(msg)}")
}
}
}
///|
/// Retry relay connection / registration until a session is
/// established or a fatal error occurs. Runs before tty/pty setup so
/// early failures stay visible on the terminal.
async fn connect_retry_loop(
config : ConnectConfig,
credential : CredentialStore,
) -> @ws.Conn? {
Some(
@async.retry(
ExponentialDelay(initial=1000, factor=2.0, maximum=5000),
fatal_error=fn(err) {
match err {
@ws.HandshakeRejected(_, response) =>
is_auth_handshake_rejection(response.code)
ProtoVersionTooOldRejection(_) => true
_ => false
}
},
() => connect_and_register(config),
),
) catch {
err if @async.is_being_cancelled() => raise err
@ws.HandshakeRejected(_, response) if is_auth_handshake_rejection(
response.code,
) => {
credential.remove_stale()
@log.error(
"relay rejected credential (this server may have been revoked); " +
"removed stale credential, re-run tun-server to re-pair",
)
None
}
@ws.HandshakeRejected(msg, response) => {
@log.error(
"Fatal: relay handshake rejected: \{response.code} \{response.reason}: \{msg}",
)
None
}
@ws.InvalidHandshake(msg) => {
@log.error("Fatal: relay handshake failed: \{msg}")
None
}
ProtoVersionTooOldRejection(min_version) => {
@log.error(
"Fatal: relay requires proto_version >= \{min_version}, this tun-server speaks \{@proto.ProtoVersion}; upgrade tun-server",
)
None
}
err => {
@log.error("Fatal: \{err}")
None
}
}
}
///|
/// Run the relay retry loop. `first_ws` is an already-accepted
/// WebSocket from the startup connect — this function runs that
/// session first, then loops on reconnect.
async fn run_retry_loop(
config : ConnectConfig,
credential : CredentialStore,
bridge : BridgeState,
first_ws? : @ws.Conn,
) -> Unit {
match first_ws {
Some(ws) =>
run_session(bridge, ws) catch {
err if @async.is_being_cancelled() => raise err
err => @log.info("relay session lost, will reconnect: \{err}")
}
None => ()
}
@async.retry(
ExponentialDelay(initial=1000, factor=2.0, maximum=5000),
fatal_error=fn(err) {
match err {
@ws.HandshakeRejected(_, response) =>
is_auth_handshake_rejection(response.code)
ProtoVersionTooOldRejection(_) => true
_ => false
}
},
() => {
let ws = connect_and_register(config)
run_session(bridge, ws)
},
) catch {
err if @async.is_being_cancelled() => raise err
@ws.HandshakeRejected(_, response) if is_auth_handshake_rejection(
response.code,
) => {
credential.remove_stale()
@log.error(
"relay rejected credential (this server may have been revoked); " +
"removed stale credential, re-run tun-server to re-pair",
)
}
@ws.HandshakeRejected(msg, response) =>
@log.error(
"Fatal: relay handshake rejected: \{response.code} \{response.reason}: \{msg}",
)
@ws.InvalidHandshake(msg) =>
@log.error("Fatal: relay handshake failed: \{msg}")
ProtoVersionTooOldRejection(min_version) =>
@log.error(
"Fatal: relay requires proto_version >= \{min_version}, this tun-server speaks \{@proto.ProtoVersion}; upgrade tun-server",
)
err => @log.error("Fatal: \{err}")
}
}
///|
/// A single registered relay session. Constructs a `RelaySession`,
/// publishes it in `bridge.current` so `pty_reader` can broadcast,
/// runs the send+recv loops, and clears `bridge.current` on
/// disconnect.
async fn run_session(bridge : BridgeState, ws : @ws.Conn) -> Unit {
defer ws.close()
@log.info("Entering message loop...")
@async.with_task_group(async fn(group) {
let session = RelaySession::new()
bridge.current.val = Some(session)
group.add_defer(fn() {
bridge.current.val = None
session.clients.clear()
})
// Single writer for ws. Two-phase loop:
// Phase 1: blocking get + send. Cancellation breaks the loop
// without re-raising so we can drain instead of bailing
// mid-shutdown.
// Phase 2: protect_from_cancel drain — non-blocking try_get for
// any messages already enqueued by handle_message before
// the recv loop saw its cancel, then the graceful
// exit_msg.
// Why one writer: the websocket lib aborts (panic, not raise) in
// start_message if write_opcode is still Some from a previously
// aborted send. If a partial send_text leaks state, no other path
// may call start_message on the same Conn. Keeping all writes in
// this task makes that condition impossible.
group.spawn_bg(async fn() {
for ;; {
let msg = session.send_queue.get() catch {
_ if @async.is_being_cancelled() => break
@aqueue.QueueAlreadyClosed => break
err => {
@log.error(
"[relay.send_queue.get.fail] error=\{@debug.to_string(err)}",
)
raise err
}
}
@async.protect_from_cancel(async fn() {
send_active_message(bridge, msg, "active", async fn() {
ws.send_text(msg)
})
})
}
@async.protect_from_cancel(async fn() {
for ;; {
let next = session.send_queue.try_get() catch {
@aqueue.QueueAlreadyClosed => None
err => {
@log.error(
"[relay.send_queue.try_get.fail] error=\{@debug.to_string(err)}",
)
raise err
}
}
match next {
Some(msg) =>
send_shutdown_message(msg, "drain", async fn() {
ws.send_text(msg)
})
None => break
}
}
if bridge.shell_exited.val {
let exit_msg : @proto.Message = Sys(TaskExited(exit_code=0))
let exit_payload = exit_msg.to_json().stringify()
send_shutdown_message(exit_payload, "task_exited", async fn() {
ws.send_text(exit_payload)
})
}
})
})
for ;; {
let text = ws.recv().read_all().text()
let incoming : @proto.Message = @json.from_json(@json.parse(text)) catch {
err => {
@log.error(
"Failed to decode incoming message: raw=\{@debug.to_string(text)}, error=\{@debug.to_string(err)}",
)
continue
}
}
handle_message(incoming, bridge, session) catch {
err if @async.is_being_cancelled() => raise err
err =>
@log.error(
"Error processing message: message=\{incoming.to_log_string()}, raw=\{@debug.to_string(text)}, error=\{@debug.to_string(err)}",
)
}
}
})
}
///|
async fn handle_message(
incoming : @proto.Message,
bridge : BridgeState,
session : RelaySession,
) -> Unit {
match incoming {
Sys(NewClient(client_id~, last_seen_id~)) => {
@log.info(
"NewClient: \{client_id} last_seen_id=\{@debug.to_string(last_seen_id)}",
)
// Ordering invariant — pty_reader runs in a separate task and
// (a) appends (id, bytes) to `bridge.replay`, (b) broadcasts
// Term(Data) for those bytes to every entry in `session.clients`
// via send_queue. A newly-registered client must:
// 1. See every log entry with id > last_seen_id replayed in
// id-order, then
// 2. See every live Term(Data) broadcast that pty_reader
// produces after this point.
// This ordering guarantee is about the retained replay window
// only. `bridge.replay` is a bounded byte tail, so if the client
// cursor falls before the retained head we still replay the
// retained suffix in order, but that suffix is not guaranteed to
// be a self-contained terminal bootstrap.
// The invariant holds iff we (i) register the client in
// `session.clients`, (ii) enqueue the replay range, all inside
// one synchronous block with no awaits. `@aqueue.Queue::put` is
// synchronous, and pty_reader cannot interleave inside a sync
// block — so its first Term(Data) for this client is ordered
// after our replay puts.
let threshold = last_seen_id.unwrap_or(0L)
session.clients[client_id] = ()
for entry in bridge.replay.entries {
if entry.id > threshold {
let data = @base64.encode(entry.bytes[:])
let msg : @proto.Message = Term(
Data(client_id~, id=entry.id, created=Some(entry.created), data~),
)
session.send_queue.put(msg.to_json().stringify())
}
}
}
Term(Data(data~, ..)) =>
if bridge.pty is Some(pty) {
pty.write(@utf8.encode(data))
}
Sys(ClientDisconnected(client_id~)) => {
@log.info("ClientDisconnected: \{client_id}")
session.clients.remove(client_id)
}
Codex(..) =>
@log.warn(
"Codex message received but codex app-server spawning is disabled",
)
Fs(fs_msg) => {
let response = @fs.handle_fs(fs_msg)
session.send_queue.put(response.to_json().stringify())
}
other => @log.warn("Unhandled message: \{@debug.to_string(other)}")
}
}