1pub const MAX_PAYLOAD: usize = 1 << 20;
15
16pub const PROTOCOL: u32 = 1;
18
19#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum Frame {
22 Input(Vec<u8>),
25 Resize {
27 cols: u16,
28 rows: u16,
29 },
30 Ping,
32 Detach,
34 End,
36
37 Output {
40 offset: u64,
41 data: Vec<u8>,
42 },
43 Lost {
46 from: u64,
47 to: u64,
48 },
49 Exited {
51 code: i32,
52 },
53 Pong,
54 Hello {
57 session: String,
58 created: bool,
59 end: u64,
60 },
61 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 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#[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 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}