hatchery_core/
multiplex.rs

1//! Binary multiplex framing for the C2 byte restreamer (A2).
2//!
3//! C2 sees only the envelope: `seq`, destination, channel kind and size.
4//! Payload bytes are opaque (inner ciphertext). No g4a / m4a parsing here.
5
6use serde::{Deserialize, Serialize};
7use std::fmt;
8
9/// Wire version for multiplex frames.
10pub const MULTIPLEX_VERSION: u8 = 1;
11
12/// Maximum DATA payload length (1 MiB). Larger payloads are split by the sender.
13pub const MAX_DATA_PAYLOAD: usize = 1024 * 1024;
14
15/// Frame header size before the payload: version, type, channel_id, seq, length.
16pub const FRAME_HEADER_LEN: usize = 1 + 1 + 4 + 8 + 4;
17
18
19/// Grow a stream read buffer when a truncated frame cannot fit.
20/// Caps at `MAX_DATA_PAYLOAD + FRAME_HEADER_LEN`. Returns Err if already at cap
21/// (the peer is sending something larger than the protocol allows).
22pub fn grow_read_buf(buf: &mut Vec<u8>, filled: usize) -> Result<(), FrameError> {
23    const CAP: usize = MAX_DATA_PAYLOAD + FRAME_HEADER_LEN;
24    if filled < buf.len() {
25        return Ok(());
26    }
27    if buf.len() >= CAP {
28        return Err(FrameError::PayloadTooLarge(filled));
29    }
30    let next = (buf.len().saturating_mul(2)).clamp(filled + 64 * 1024, CAP);
31    buf.resize(next, 0);
32    Ok(())
33}
34
35
36#[repr(u8)]
37#[derive(Clone, Copy, Debug, PartialEq, Eq)]
38pub enum FrameType {
39    Open = 1,
40    Data = 2,
41    Close = 3,
42    Window = 4,
43    Ping = 5,
44    Pong = 6,
45    GoAway = 7,
46    Ack = 8,
47    Resume = 9,
48    /// C2 → node: request journal entries with seq > N.
49    JournalReq = 10,
50    /// node → C2: journal entries (payload = JournalReplyBody).
51    JournalReply = 11,
52}
53
54impl FrameType {
55    pub fn from_u8(v: u8) -> Option<Self> {
56        Some(match v {
57            1 => Self::Open,
58            2 => Self::Data,
59            3 => Self::Close,
60            4 => Self::Window,
61            5 => Self::Ping,
62            6 => Self::Pong,
63            7 => Self::GoAway,
64            8 => Self::Ack,
65            9 => Self::Resume,
66            10 => Self::JournalReq,
67            11 => Self::JournalReply,
68            _ => return None,
69        })
70    }
71}
72
73/// Logical channel kinds C2 routes without inspecting contents.
74#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
75#[serde(rename_all = "snake_case")]
76pub enum ChannelKind {
77    G4aControl,
78    G4aPty,
79    G4aEvents,
80    M4a,
81    Vault,
82    Observation,
83}
84
85impl ChannelKind {
86    pub fn as_str(self) -> &'static str {
87        match self {
88            Self::G4aControl => "g4a.control",
89            Self::G4aPty => "g4a.pty",
90            Self::G4aEvents => "g4a.events",
91            Self::M4a => "m4a",
92            Self::Vault => "vault",
93            Self::Observation => "observation",
94        }
95    }
96
97    pub fn parse(s: &str) -> Option<Self> {
98        Some(match s {
99            "g4a.control" => Self::G4aControl,
100            "g4a.pty" => Self::G4aPty,
101            "g4a.events" => Self::G4aEvents,
102            "m4a" => Self::M4a,
103            "vault" => Self::Vault,
104            "observation" => Self::Observation,
105            _ => return None,
106        })
107    }
108}
109
110impl fmt::Display for ChannelKind {
111    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112        f.write_str(self.as_str())
113    }
114}
115
116/// Payload of an OPEN frame. No secrets.
117#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
118pub struct OpenPayload {
119    pub channel_kind: ChannelKind,
120    /// Opaque node identifier assigned by C2; core never sees node addresses.
121    pub node_id: String,
122}
123
124/// One multiplex frame.
125#[derive(Clone, Debug, PartialEq, Eq)]
126pub struct Frame {
127    pub version: u8,
128    pub frame_type: FrameType,
129    pub channel_id: u32,
130    /// For DATA: event sequence. For ACK/RESUME: the acknowledged / resume-after number.
131    pub seq: u64,
132    pub payload: Vec<u8>,
133}
134
135#[derive(Clone, Debug, PartialEq, Eq)]
136pub enum FrameError {
137    Truncated,
138    BadVersion(u8),
139    BadType(u8),
140    PayloadTooLarge(usize),
141    Io(String),
142}
143
144impl fmt::Display for FrameError {
145    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
146        match self {
147            Self::Truncated => write!(f, "truncated frame"),
148            Self::BadVersion(v) => write!(f, "unsupported multiplex version {v}"),
149            Self::BadType(t) => write!(f, "unknown frame type {t}"),
150            Self::PayloadTooLarge(n) => write!(f, "payload length {n} exceeds limit"),
151            Self::Io(msg) => write!(f, "{msg}"),
152        }
153    }
154}
155
156impl std::error::Error for FrameError {}
157
158impl Frame {
159    pub fn open(channel_id: u32, payload: &OpenPayload) -> Result<Self, FrameError> {
160        let body = serde_json::to_vec(payload).map_err(|e| FrameError::Io(e.to_string()))?;
161        Ok(Self {
162            version: MULTIPLEX_VERSION,
163            frame_type: FrameType::Open,
164            channel_id,
165            seq: 0,
166            payload: body,
167        })
168    }
169
170    pub fn data(channel_id: u32, seq: u64, payload: Vec<u8>) -> Result<Self, FrameError> {
171        if payload.len() > MAX_DATA_PAYLOAD {
172            return Err(FrameError::PayloadTooLarge(payload.len()));
173        }
174        Ok(Self {
175            version: MULTIPLEX_VERSION,
176            frame_type: FrameType::Data,
177            channel_id,
178            seq,
179            payload,
180        })
181    }
182
183    pub fn ack(seq: u64) -> Self {
184        Self {
185            version: MULTIPLEX_VERSION,
186            frame_type: FrameType::Ack,
187            channel_id: 0,
188            seq,
189            payload: Vec::new(),
190        }
191    }
192
193    pub fn resume(after: u64) -> Self {
194        Self {
195            version: MULTIPLEX_VERSION,
196            frame_type: FrameType::Resume,
197            channel_id: 0,
198            seq: after,
199            payload: Vec::new(),
200        }
201    }
202
203    pub fn window(channel_id: u32, credit: u32) -> Self {
204        Self {
205            version: MULTIPLEX_VERSION,
206            frame_type: FrameType::Window,
207            channel_id,
208            seq: u64::from(credit),
209            payload: Vec::new(),
210        }
211    }
212
213    pub fn ping(nonce: u64) -> Self {
214        Self {
215            version: MULTIPLEX_VERSION,
216            frame_type: FrameType::Ping,
217            channel_id: 0,
218            seq: nonce,
219            payload: Vec::new(),
220        }
221    }
222
223    pub fn pong(nonce: u64) -> Self {
224        Self {
225            version: MULTIPLEX_VERSION,
226            frame_type: FrameType::Pong,
227            channel_id: 0,
228            seq: nonce,
229            payload: Vec::new(),
230        }
231    }
232
233    pub fn goaway(code: u32) -> Self {
234        Self {
235            version: MULTIPLEX_VERSION,
236            frame_type: FrameType::GoAway,
237            channel_id: 0,
238            seq: u64::from(code),
239            payload: Vec::new(),
240        }
241    }
242
243    pub fn close(channel_id: u32) -> Self {
244        Self {
245            version: MULTIPLEX_VERSION,
246            frame_type: FrameType::Close,
247            channel_id,
248            seq: 0,
249            payload: Vec::new(),
250        }
251    }
252
253    /// Encode big-endian wire bytes.
254    pub fn encode(&self) -> Result<Vec<u8>, FrameError> {
255        if self.payload.len() > MAX_DATA_PAYLOAD {
256            return Err(FrameError::PayloadTooLarge(self.payload.len()));
257        }
258        let mut out = Vec::with_capacity(FRAME_HEADER_LEN + self.payload.len());
259        out.push(self.version);
260        out.push(self.frame_type as u8);
261        out.extend_from_slice(&self.channel_id.to_be_bytes());
262        out.extend_from_slice(&self.seq.to_be_bytes());
263        out.extend_from_slice(&(self.payload.len() as u32).to_be_bytes());
264        out.extend_from_slice(&self.payload);
265        Ok(out)
266    }
267
268    /// Decode one frame from a complete buffer. Returns `(frame, bytes_consumed)`.
269    pub fn decode(buf: &[u8]) -> Result<(Self, usize), FrameError> {
270        if buf.len() < FRAME_HEADER_LEN {
271            return Err(FrameError::Truncated);
272        }
273        let version = buf[0];
274        if version != MULTIPLEX_VERSION {
275            return Err(FrameError::BadVersion(version));
276        }
277        let frame_type = FrameType::from_u8(buf[1]).ok_or(FrameError::BadType(buf[1]))?;
278        let channel_id = u32::from_be_bytes([buf[2], buf[3], buf[4], buf[5]]);
279        let seq = u64::from_be_bytes([
280            buf[6], buf[7], buf[8], buf[9], buf[10], buf[11], buf[12], buf[13],
281        ]);
282        let length = u32::from_be_bytes([buf[14], buf[15], buf[16], buf[17]]) as usize;
283        if length > MAX_DATA_PAYLOAD {
284            return Err(FrameError::PayloadTooLarge(length));
285        }
286        let total = FRAME_HEADER_LEN + length;
287        if buf.len() < total {
288            return Err(FrameError::Truncated);
289        }
290        Ok((
291            Self {
292                version,
293                frame_type,
294                channel_id,
295                seq,
296                payload: buf[FRAME_HEADER_LEN..total].to_vec(),
297            },
298            total,
299        ))
300    }
301
302    pub fn parse_open(&self) -> Result<OpenPayload, FrameError> {
303        if self.frame_type != FrameType::Open {
304            return Err(FrameError::Io("not an OPEN frame".into()));
305        }
306        serde_json::from_slice(&self.payload).map_err(|e| FrameError::Io(e.to_string()))
307    }
308
309    /// Journal request: `seq` = after, payload = u32 limit (big-endian).
310    pub fn journal_req(channel_id: u32, after: u64, limit: u32) -> Self {
311        Self {
312            version: MULTIPLEX_VERSION,
313            frame_type: FrameType::JournalReq,
314            channel_id,
315            seq: after,
316            payload: limit.to_be_bytes().to_vec(),
317        }
318    }
319
320    pub fn journal_req_limit(&self) -> Option<u32> {
321        if self.frame_type != FrameType::JournalReq || self.payload.len() != 4 {
322            return None;
323        }
324        Some(u32::from_be_bytes([
325            self.payload[0],
326            self.payload[1],
327            self.payload[2],
328            self.payload[3],
329        ]))
330    }
331
332    pub fn journal_reply(channel_id: u32, body: &JournalReplyBody) -> Result<Self, FrameError> {
333        let payload = body.encode()?;
334        if payload.len() > MAX_DATA_PAYLOAD {
335            return Err(FrameError::PayloadTooLarge(payload.len()));
336        }
337        Ok(Self {
338            version: MULTIPLEX_VERSION,
339            frame_type: FrameType::JournalReply,
340            channel_id,
341            seq: body.after,
342            payload,
343        })
344    }
345}
346
347/// One journal entry on the wire (opaque payload).
348#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
349pub struct JournalWireEntry {
350    pub seq: u64,
351    pub payload: Vec<u8>,
352}
353
354/// Body of JournalReply: entries with seq > after for one channel.
355#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
356pub struct JournalReplyBody {
357    pub after: u64,
358    pub entries: Vec<JournalWireEntry>,
359}
360
361impl JournalReplyBody {
362    pub fn encode(&self) -> Result<Vec<u8>, FrameError> {
363        serde_json::to_vec(self).map_err(|e| FrameError::Io(e.to_string()))
364    }
365
366    pub fn decode(bytes: &[u8]) -> Result<Self, FrameError> {
367        serde_json::from_slice(bytes).map_err(|e| FrameError::Io(e.to_string()))
368    }
369}
370
371/// Drain-stream envelope: C2 prefixes each uplink DATA with an opaque node id for core.
372#[derive(Clone, Debug, PartialEq, Eq)]
373pub struct DrainEnvelope {
374    /// Opaque destination / source node id. Not an address.
375    pub node_id: String,
376    pub frame: Frame,
377}
378
379impl DrainEnvelope {
380    pub fn encode(&self) -> Result<Vec<u8>, FrameError> {
381        let id = self.node_id.as_bytes();
382        if id.len() > 255 {
383            return Err(FrameError::Io("node_id too long".into()));
384        }
385        let frame_bytes = self.frame.encode()?;
386        let mut out = Vec::with_capacity(1 + id.len() + frame_bytes.len());
387        out.push(id.len() as u8);
388        out.extend_from_slice(id);
389        out.extend_from_slice(&frame_bytes);
390        Ok(out)
391    }
392
393    pub fn decode(buf: &[u8]) -> Result<(Self, usize), FrameError> {
394        if buf.is_empty() {
395            return Err(FrameError::Truncated);
396        }
397        let id_len = buf[0] as usize;
398        if buf.len() < 1 + id_len {
399            return Err(FrameError::Truncated);
400        }
401        let node_id = String::from_utf8(buf[1..1 + id_len].to_vec())
402            .map_err(|_| FrameError::Io("node_id is not utf-8".into()))?;
403        let (frame, n) = Frame::decode(&buf[1 + id_len..])?;
404        Ok((Self { node_id, frame }, 1 + id_len + n))
405    }
406}
407
408
409/// Magic for node→C2 UDP data-plane datagrams (PTY bulk).
410pub const UDP_DATA_MAGIC: &[u8; 4] = b"H4UD";
411
412/// Encode one UDP datagram: magic + node_id + one DATA frame.
413/// Used for PTY/stream bulk so control stays on TCP.
414pub fn encode_udp_data(node_id: &str, frame: &Frame) -> Result<Vec<u8>, FrameError> {
415    if !matches!(frame.frame_type, FrameType::Data) {
416        return Err(FrameError::Io("udp data plane accepts DATA frames only".into()));
417    }
418    let nid = node_id.as_bytes();
419    if nid.is_empty() || nid.len() > 255 {
420        return Err(FrameError::Io("udp node_id length".into()));
421    }
422    let body = frame.encode()?;
423    let mut out = Vec::with_capacity(4 + 1 + nid.len() + body.len());
424    out.extend_from_slice(UDP_DATA_MAGIC);
425    out.push(nid.len() as u8);
426    out.extend_from_slice(nid);
427    out.extend_from_slice(&body);
428    Ok(out)
429}
430
431/// Decode a UDP datagram into `(node_id, frame)`.
432pub fn decode_udp_data(buf: &[u8]) -> Result<(String, Frame), FrameError> {
433    if buf.len() < 5 + FRAME_HEADER_LEN {
434        return Err(FrameError::Truncated);
435    }
436    if &buf[..4] != UDP_DATA_MAGIC {
437        return Err(FrameError::Io("udp magic".into()));
438    }
439    let nlen = buf[4] as usize;
440    if nlen == 0 || buf.len() < 5 + nlen + FRAME_HEADER_LEN {
441        return Err(FrameError::Truncated);
442    }
443    let node_id = std::str::from_utf8(&buf[5..5 + nlen])
444        .map_err(|_| FrameError::Io("udp node_id utf8".into()))?
445        .to_owned();
446    let (frame, consumed) = Frame::decode(&buf[5 + nlen..])?;
447    if 5 + nlen + consumed != buf.len() {
448        return Err(FrameError::Io("udp trailing bytes".into()));
449    }
450    if !matches!(frame.frame_type, FrameType::Data) {
451        return Err(FrameError::Io("udp data plane accepts DATA frames only".into()));
452    }
453    Ok((node_id, frame))
454}
455
456#[cfg(test)]
457mod tests {
458    use super::*;
459
460    #[test]
461    fn roundtrip_data_ack_resume() {
462        let data = Frame::data(7, 42, b"opaque-ciphertext".to_vec()).unwrap();
463        let bytes = data.encode().unwrap();
464        let (decoded, n) = Frame::decode(&bytes).unwrap();
465        assert_eq!(n, bytes.len());
466        assert_eq!(decoded, data);
467
468        let ack = Frame::ack(42);
469        let (d, _) = Frame::decode(&ack.encode().unwrap()).unwrap();
470        assert_eq!(d.frame_type, FrameType::Ack);
471        assert_eq!(d.seq, 42);
472
473        let resume = Frame::resume(41);
474        let (d, _) = Frame::decode(&resume.encode().unwrap()).unwrap();
475        assert_eq!(d.frame_type, FrameType::Resume);
476        assert_eq!(d.seq, 41);
477    }
478
479    #[test]
480    fn open_payload_roundtrip() {
481        let open = Frame::open(
482            1,
483            &OpenPayload {
484                channel_kind: ChannelKind::M4a,
485                node_id: "node-a".into(),
486            },
487        )
488        .unwrap();
489        let (decoded, _) = Frame::decode(&open.encode().unwrap()).unwrap();
490        let payload = decoded.parse_open().unwrap();
491        assert_eq!(payload.channel_kind, ChannelKind::M4a);
492        assert_eq!(payload.node_id, "node-a");
493    }
494
495
496    #[test]
497    fn grow_read_buf_accepts_frames_between_64kib_and_1mib() {
498        // Simulate the C2 uplink read loop: start at 64 KiB, receive a 100 KiB DATA.
499        let payload = vec![0xABu8; 100 * 1024];
500        let frame = Frame::data(1, 7, payload.clone()).unwrap();
501        let wire = frame.encode().unwrap();
502        assert!(wire.len() > 64 * 1024);
503
504        let mut buf = vec![0u8; 64 * 1024];
505        let mut filled = 0usize;
506        let mut offset = 0usize;
507        loop {
508            if filled == buf.len() {
509                grow_read_buf(&mut buf, filled).expect("must grow for 100 KiB frame");
510            }
511            let space = buf.len() - filled;
512            let take = space.min(wire.len() - offset);
513            buf[filled..filled + take].copy_from_slice(&wire[offset..offset + take]);
514            filled += take;
515            offset += take;
516            match Frame::decode(&buf[..filled]) {
517                Ok((decoded, consumed)) => {
518                    assert_eq!(decoded.payload, payload);
519                    assert_eq!(consumed, wire.len());
520                    break;
521                }
522                Err(FrameError::Truncated) => {
523                    assert!(offset < wire.len(), "still truncated after full wire");
524                    continue;
525                }
526                Err(e) => panic!("unexpected {e}"),
527            }
528        }
529    }
530
531    #[test]
532    fn rejects_oversized_payload() {
533        let big = vec![0u8; MAX_DATA_PAYLOAD + 1];
534        assert!(matches!(
535            Frame::data(1, 1, big),
536            Err(FrameError::PayloadTooLarge(_))
537        ));
538    }
539
540    #[test]
541    fn drain_envelope_prefixes_opaque_node_id() {
542        let env = DrainEnvelope {
543            node_id: "node-b".into(),
544            frame: Frame::data(3, 9, b"inner".to_vec()).unwrap(),
545        };
546        let bytes = env.encode().unwrap();
547        let (decoded, n) = DrainEnvelope::decode(&bytes).unwrap();
548        assert_eq!(n, bytes.len());
549        assert_eq!(decoded.node_id, "node-b");
550        assert_eq!(decoded.frame.seq, 9);
551        assert_eq!(decoded.frame.payload, b"inner");
552    }
553
554    #[test]
555    fn journal_req_roundtrip() {
556        let f = Frame::journal_req(3, 10, 64);
557        let (d, _) = Frame::decode(&f.encode().unwrap()).unwrap();
558        assert_eq!(d.frame_type, FrameType::JournalReq);
559        assert_eq!(d.seq, 10);
560        assert_eq!(d.journal_req_limit(), Some(64));
561        let body = JournalReplyBody {
562            after: 10,
563            entries: vec![JournalWireEntry {
564                seq: 11,
565                payload: b"x".to_vec(),
566            }],
567        };
568        let reply = Frame::journal_reply(3, &body).unwrap();
569        let (d, _) = Frame::decode(&reply.encode().unwrap()).unwrap();
570        let decoded = JournalReplyBody::decode(&d.payload).unwrap();
571        assert_eq!(decoded.entries[0].seq, 11);
572    }
573
574    #[test]
575    fn truncated_frame_errors() {
576        assert!(matches!(Frame::decode(&[1, 2, 3]), Err(FrameError::Truncated)));
577        let mut partial = Frame::data(1, 1, b"hello".to_vec()).unwrap().encode().unwrap();
578        partial.pop();
579        assert!(matches!(Frame::decode(&partial), Err(FrameError::Truncated)));
580    }
581
582    #[test]
583    fn udp_data_roundtrip() {
584        let frame = Frame::data(3, 0, vec![1, 2, 3, 4]).unwrap();
585        let dg = encode_udp_data("hatchery-box", &frame).unwrap();
586        let (nid, f) = decode_udp_data(&dg).unwrap();
587        assert_eq!(nid, "hatchery-box");
588        assert_eq!(f.channel_id, 3);
589        assert_eq!(f.payload, vec![1, 2, 3, 4]);
590    }
591}
hatchery_core/multiplex.rs

hatchery_core/
multiplex.rs
