Skip to main content

agentmux_remote/
frame.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! The wire protocol between srv and a durable session (spec §7.2), carried
5//! end to end over `ssh -T host -- agentmux-remote attach ...`: `attach` only
6//! relays bytes between its stdio and the daemon's socket, so these frames go
7//! straight from the daemon to srv and back.
8//!
9//! A frame is `[kind: u8][len: u32 big-endian][payload: len bytes]`. Output
10//! carries the session's byte offset, so srv appends exactly what it has not
11//! seen and a reattach asks for exactly what it missed.
12
13/// Largest payload accepted; a longer length means a broken stream.
14pub const MAX_PAYLOAD: usize = 1 << 20;
15
16/// The protocol version `attach` and srv must share (the handshake line).
17pub const PROTOCOL: u32 = 1;
18
19/// One message.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum Frame {
22    // srv -> session
23    /// Keystrokes for the session's terminal.
24    Input(Vec<u8>),
25    /// The pane's size.
26    Resize {
27        cols: u16,
28        rows: u16,
29    },
30    /// Are you there? Answered with `Pong`; silence means a stalled link (§7.5).
31    Ping,
32    /// Leave the session running and drop this connection.
33    Detach,
34    /// End the session: its shell and everything it started.
35    End,
36
37    // session -> srv
38    /// Terminal output starting at byte `offset` of the session's stream.
39    Output {
40        offset: u64,
41        data: Vec<u8>,
42    },
43    /// Bytes `from..to` fell out of the session's ring before they could be
44    /// replayed: the pane marks a gap there.
45    Lost {
46        from: u64,
47        to: u64,
48    },
49    /// The session's shell exited with `code`; the session is over.
50    Exited {
51        code: i32,
52    },
53    Pong,
54    /// Sent first: the session id, whether it was just created, and the next
55    /// offset its output will have.
56    Hello {
57        session: String,
58        created: bool,
59        end: u64,
60    },
61    /// Something went wrong, in words for the pane.
62    Error(String),
63}
64
65mod kind {
66    pub const INPUT: u8 = 1;
67    pub const RESIZE: u8 = 2;
68    pub const PING: u8 = 3;
69    pub const DETACH: u8 = 4;
70    pub const END: u8 = 5;
71    pub const OUTPUT: u8 = 10;
72    pub const LOST: u8 = 11;
73    pub const EXITED: u8 = 12;
74    pub const PONG: u8 = 13;
75    pub const HELLO: u8 = 14;
76    pub const ERROR: u8 = 15;
77}
78
79impl Frame {
80    /// The frame's bytes on the wire.
81    pub fn encode(&self) -> Vec<u8> {
82        let (k, payload): (u8, Vec<u8>) = match self {
83            Frame::Input(d) => (kind::INPUT, d.clone()),
84            Frame::Resize { cols, rows } => {
85                let mut p = cols.to_be_bytes().to_vec();
86                p.extend_from_slice(&rows.to_be_bytes());
87                (kind::RESIZE, p)
88            }
89            Frame::Ping => (kind::PING, Vec::new()),
90            Frame::Detach => (kind::DETACH, Vec::new()),
91            Frame::End => (kind::END, Vec::new()),
92            Frame::Output { offset, data } => {
93                let mut p = offset.to_be_bytes().to_vec();
94                p.extend_from_slice(data);
95                (kind::OUTPUT, p)
96            }
97            Frame::Lost { from, to } => {
98                let mut p = from.to_be_bytes().to_vec();
99                p.extend_from_slice(&to.to_be_bytes());
100                (kind::LOST, p)
101            }
102            Frame::Exited { code } => (kind::EXITED, code.to_be_bytes().to_vec()),
103            Frame::Pong => (kind::PONG, Vec::new()),
104            Frame::Hello {
105                session,
106                created,
107                end,
108            } => {
109                let mut p = end.to_be_bytes().to_vec();
110                p.push(u8::from(*created));
111                p.extend_from_slice(session.as_bytes());
112                (kind::HELLO, p)
113            }
114            Frame::Error(msg) => (kind::ERROR, msg.as_bytes().to_vec()),
115        };
116        let mut out = Vec::with_capacity(5 + payload.len());
117        out.push(k);
118        out.extend_from_slice(&(payload.len() as u32).to_be_bytes());
119        out.extend_from_slice(&payload);
120        out
121    }
122
123    fn decode(k: u8, p: &[u8]) -> Result<Frame, String> {
124        let u64_at = |i: usize| -> Result<u64, String> {
125            p.get(i..i + 8)
126                .map(|b| u64::from_be_bytes(b.try_into().unwrap()))
127                .ok_or_else(|| format!("frame {k}: too short"))
128        };
129        Ok(match k {
130            kind::INPUT => Frame::Input(p.to_vec()),
131            kind::RESIZE => {
132                if p.len() != 4 {
133                    return Err("resize: wrong length".into());
134                }
135                Frame::Resize {
136                    cols: u16::from_be_bytes([p[0], p[1]]),
137                    rows: u16::from_be_bytes([p[2], p[3]]),
138                }
139            }
140            kind::PING => Frame::Ping,
141            kind::DETACH => Frame::Detach,
142            kind::END => Frame::End,
143            kind::OUTPUT => Frame::Output {
144                offset: u64_at(0)?,
145                data: p[8..].to_vec(),
146            },
147            kind::LOST => Frame::Lost {
148                from: u64_at(0)?,
149                to: u64_at(8)?,
150            },
151            kind::EXITED => {
152                let b: [u8; 4] = p
153                    .try_into()
154                    .map_err(|_| "exited: wrong length".to_string())?;
155                Frame::Exited {
156                    code: i32::from_be_bytes(b),
157                }
158            }
159            kind::PONG => Frame::Pong,
160            kind::HELLO => {
161                let end = u64_at(0)?;
162                let created = *p.get(8).ok_or("hello: too short")? != 0;
163                let session = String::from_utf8(p[9..].to_vec())
164                    .map_err(|_| "hello: session id is not UTF-8")?;
165                Frame::Hello {
166                    session,
167                    created,
168                    end,
169                }
170            }
171            kind::ERROR => Frame::Error(String::from_utf8_lossy(p).into_owned()),
172            other => return Err(format!("unknown frame kind {other}")),
173        })
174    }
175}
176
177/// Reassembles frames from a byte stream that arrives in arbitrary pieces.
178#[derive(Debug, Default)]
179pub struct Decoder {
180    buf: Vec<u8>,
181}
182
183impl Decoder {
184    pub fn new() -> Self {
185        Self::default()
186    }
187
188    /// Add bytes; return every frame now complete. `Err` for a stream that is
189    /// not frames (an unknown kind or an impossible length): the connection
190    /// should be dropped.
191    pub fn push(&mut self, bytes: &[u8]) -> Result<Vec<Frame>, String> {
192        self.buf.extend_from_slice(bytes);
193        let mut frames = Vec::new();
194        loop {
195            if self.buf.len() < 5 {
196                break;
197            }
198            let len =
199                u32::from_be_bytes([self.buf[1], self.buf[2], self.buf[3], self.buf[4]]) as usize;
200            if len > MAX_PAYLOAD {
201                return Err(format!("frame of {len} bytes is over the limit"));
202            }
203            if self.buf.len() < 5 + len {
204                break;
205            }
206            let frame = Frame::decode(self.buf[0], &self.buf[5..5 + len])?;
207            self.buf.drain(..5 + len);
208            frames.push(frame);
209        }
210        Ok(frames)
211    }
212}
213
214#[cfg(test)]
215mod tests {
216    use super::*;
217
218    fn all() -> Vec<Frame> {
219        vec![
220            Frame::Input(b"ls -la\r".to_vec()),
221            Frame::Resize {
222                cols: 200,
223                rows: 50,
224            },
225            Frame::Ping,
226            Frame::Detach,
227            Frame::End,
228            Frame::Output {
229                offset: 1 << 40,
230                data: b"\x1b[31mred\x1b[0m".to_vec(),
231            },
232            Frame::Lost {
233                from: 10,
234                to: 8_000_000,
235            },
236            Frame::Exited { code: -1 },
237            Frame::Pong,
238            Frame::Hello {
239                session: "s-1".into(),
240                created: true,
241                end: 42,
242            },
243            Frame::Error("no such session".into()),
244        ]
245    }
246
247    #[test]
248    fn every_frame_round_trips() {
249        for f in all() {
250            let mut d = Decoder::new();
251            assert_eq!(d.push(&f.encode()).unwrap(), vec![f.clone()], "{f:?}");
252        }
253    }
254
255    #[test]
256    fn frames_split_anywhere_and_joined_reassemble() {
257        let stream: Vec<u8> = all().iter().flat_map(Frame::encode).collect();
258        for chunk in [1, 2, 3, 7, 64] {
259            let mut d = Decoder::new();
260            let mut got = Vec::new();
261            for piece in stream.chunks(chunk) {
262                got.extend(d.push(piece).unwrap());
263            }
264            assert_eq!(got, all(), "chunk {chunk}");
265        }
266    }
267
268    #[test]
269    fn a_stream_that_is_not_frames_is_an_error() {
270        assert!(
271            Decoder::new().push(&[99, 0, 0, 0, 0]).is_err(),
272            "unknown kind"
273        );
274        assert!(
275            Decoder::new().push(&[1, 0xff, 0xff, 0xff, 0xff]).is_err(),
276            "impossible length"
277        );
278        assert!(
279            Decoder::new().push(&[2, 0, 0, 0, 1, 7]).is_err(),
280            "bad resize"
281        );
282    }
283}