reddb_wire/redwire/
codec.rs

1//! Hand-rolled binary codec for v2 frames. No serde — the on-wire
2//! shape is fixed by ADR 0001, kept simple so a hex-dump is
3//! readable.
4
5use super::frame::{Flags, Frame, MessageKind, FRAME_HEADER_SIZE, MAX_FRAME_SIZE};
6
7/// Upper bound on the inflated size of a COMPRESSED payload. Four times
8/// the on-wire frame cap: enough headroom for compression to be useful,
9/// small enough that a hostile peer cannot exhaust memory with one frame.
10pub const MAX_DECOMPRESSED_PAYLOAD: usize = 4 * (MAX_FRAME_SIZE as usize);
11
12fn decompress_bounded(on_wire: &[u8], payload_len: usize) -> Result<Vec<u8>, FrameError> {
13    use std::io::Read;
14    let limit = MAX_DECOMPRESSED_PAYLOAD as u64;
15    let unparseable = |e: std::io::Error| FrameError::PayloadTruncated {
16        // Reuse PayloadTruncated for "decompression failed" rather than
17        // introduce a new variant — the wire-layer outcome is the same:
18        // the body is unparseable, drop the connection.
19        expected: payload_len as u32,
20        available: e.to_string().len() as u32,
21    };
22    let decoder = zstd::stream::Decoder::new(on_wire).map_err(unparseable)?;
23    let mut plain = Vec::new();
24    // Read one byte past the limit so an exactly-at-limit payload is
25    // accepted while anything larger is detected without buffering it.
26    decoder
27        .take(limit + 1)
28        .read_to_end(&mut plain)
29        .map_err(unparseable)?;
30    if plain.len() as u64 > limit {
31        return Err(FrameError::DecompressedTooLarge {
32            limit: MAX_DECOMPRESSED_PAYLOAD as u32,
33        });
34    }
35    Ok(plain)
36}
37
38#[derive(Debug, Clone, PartialEq, Eq)]
39pub enum FrameError {
40    Truncated,
41    InvalidLength(u32),
42    PayloadTruncated {
43        expected: u32,
44        available: u32,
45    },
46    UnknownKind(u8),
47    UnknownFlags(u8),
48    /// Catalog cross-check failed: the flag bits set on the frame are
49    /// not in `MessageKind::allowed_flags()` for this kind. The wire
50    /// catalog is the single source of truth for which bits a kind
51    /// may carry — see `frame.rs::MessageKind::allowed_flags`.
52    /// A COMPRESSED payload inflated past [`MAX_DECOMPRESSED_PAYLOAD`].
53    DecompressedTooLarge {
54        limit: u32,
55    },
56    FlagsNotAllowedForKind {
57        kind: u8,
58        flags: u8,
59    },
60}
61
62impl std::fmt::Display for FrameError {
63    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64        match self {
65            Self::Truncated => write!(f, "frame header truncated (< 16 bytes)"),
66            Self::InvalidLength(n) => write!(f, "frame length field invalid: {n}"),
67            Self::PayloadTruncated {
68                expected,
69                available,
70            } => write!(
71                f,
72                "frame payload truncated: expected {expected} bytes, got {available}"
73            ),
74            Self::UnknownKind(byte) => write!(f, "unknown message kind 0x{byte:02x}"),
75            Self::UnknownFlags(byte) => write!(f, "unknown flag bits 0x{byte:02x}"),
76            Self::DecompressedTooLarge { limit } => {
77                write!(f, "compressed payload inflates past the {limit}-byte limit")
78            }
79            Self::FlagsNotAllowedForKind { kind, flags } => write!(
80                f,
81                "flag bits 0x{flags:02x} not allowed on kind 0x{kind:02x}"
82            ),
83        }
84    }
85}
86
87impl std::error::Error for FrameError {}
88
89pub fn frame_len_from_header(header: &[u8; FRAME_HEADER_SIZE]) -> Result<usize, FrameError> {
90    let length = u32::from_le_bytes([header[0], header[1], header[2], header[3]]);
91    if length < FRAME_HEADER_SIZE as u32 || length > MAX_FRAME_SIZE {
92        return Err(FrameError::InvalidLength(length));
93    }
94    Ok(length as usize)
95}
96
97pub fn decode_frame_parts(
98    header: &[u8; FRAME_HEADER_SIZE],
99    payload: &[u8],
100) -> Result<Frame, FrameError> {
101    let length = frame_len_from_header(header)?;
102    let expected_payload_len = length - FRAME_HEADER_SIZE;
103    if payload.len() < expected_payload_len {
104        return Err(FrameError::PayloadTruncated {
105            expected: expected_payload_len as u32,
106            available: payload.len() as u32,
107        });
108    }
109
110    let mut bytes = Vec::with_capacity(length);
111    bytes.extend_from_slice(header);
112    bytes.extend_from_slice(&payload[..expected_payload_len]);
113    decode_frame(&bytes).map(|(frame, _)| frame)
114}
115
116pub fn encode_frame(frame: &Frame) -> Vec<u8> {
117    // The frame's `payload` is always the plaintext form. If the
118    // COMPRESSED flag is set we compress on the wire and rewrite
119    // the length header to match the compressed size — the
120    // receiver inflates before delivering to the dispatch loop.
121    if frame.flags.contains(Flags::COMPRESSED) {
122        return encode_compressed(frame);
123    }
124    let total = frame.encoded_len() as usize;
125    let mut buf = Vec::with_capacity(total);
126    buf.extend_from_slice(&frame.encoded_len().to_le_bytes());
127    buf.push(frame.kind as u8);
128    buf.push(frame.flags.bits());
129    buf.extend_from_slice(&frame.stream_id.to_le_bytes());
130    buf.extend_from_slice(&frame.correlation_id.to_le_bytes());
131    buf.extend_from_slice(&frame.payload);
132    buf
133}
134
135fn encode_compressed(frame: &Frame) -> Vec<u8> {
136    // zstd level 1 — keeps CPU low while still cutting JSON +
137    // BulkInsertBinary by 60-80%. Operators that want max ratio
138    // can flip to level 3+ via `RED_REDWIRE_ZSTD_LEVEL` env.
139    let level = std::env::var("RED_REDWIRE_ZSTD_LEVEL")
140        .ok()
141        .and_then(|s| s.parse::<i32>().ok())
142        .unwrap_or(1);
143    let compressed = match zstd::stream::encode_all(frame.payload.as_slice(), level) {
144        Ok(buf) => buf,
145        Err(_) => {
146            // Fallback: drop the COMPRESSED flag and ship plaintext.
147            // Compression failures are rare (level 1 effectively
148            // never fails on bytes), but the fallback is safer
149            // than panicking inside the framing layer.
150            let mut clone = frame.clone();
151            clone.flags = Flags::from_bits(clone.flags.bits() & !Flags::COMPRESSED.bits());
152            return encode_frame(&clone);
153        }
154    };
155    let total = (FRAME_HEADER_SIZE + compressed.len()) as u32;
156    let mut buf = Vec::with_capacity(total as usize);
157    buf.extend_from_slice(&total.to_le_bytes());
158    buf.push(frame.kind as u8);
159    buf.push(frame.flags.bits());
160    buf.extend_from_slice(&frame.stream_id.to_le_bytes());
161    buf.extend_from_slice(&frame.correlation_id.to_le_bytes());
162    buf.extend_from_slice(&compressed);
163    buf
164}
165
166pub fn decode_frame(bytes: &[u8]) -> Result<(Frame, usize), FrameError> {
167    if bytes.len() < FRAME_HEADER_SIZE {
168        return Err(FrameError::Truncated);
169    }
170    let mut header = [0u8; FRAME_HEADER_SIZE];
171    header.copy_from_slice(&bytes[..FRAME_HEADER_SIZE]);
172    let length = frame_len_from_header(&header)? as u32;
173    if (bytes.len() as u32) < length {
174        return Err(FrameError::PayloadTruncated {
175            expected: length,
176            available: bytes.len() as u32,
177        });
178    }
179    let kind = MessageKind::from_u8(bytes[4]).ok_or(FrameError::UnknownKind(bytes[4]))?;
180    let flag_bits = bytes[5];
181    const KNOWN_FLAGS: u8 = 0b0000_0011;
182    if flag_bits & !KNOWN_FLAGS != 0 {
183        return Err(FrameError::UnknownFlags(flag_bits));
184    }
185    let flags = Flags::from_bits(flag_bits);
186    // Catalog cross-check: the kind's `allowed_flags()` is the single
187    // source of truth. Reject combinations the catalog forbids
188    // (e.g. COMPRESSED on tiny handshake payloads) so misframed
189    // frames fail at the boundary instead of reaching dispatch.
190    if !kind.permits_flags(flags) {
191        return Err(FrameError::FlagsNotAllowedForKind {
192            kind: bytes[4],
193            flags: flag_bits,
194        });
195    }
196    let stream_id = u16::from_le_bytes([bytes[6], bytes[7]]);
197    let correlation_id = u64::from_le_bytes([
198        bytes[8], bytes[9], bytes[10], bytes[11], bytes[12], bytes[13], bytes[14], bytes[15],
199    ]);
200    let payload_len = (length as usize) - FRAME_HEADER_SIZE;
201    let on_wire = &bytes[FRAME_HEADER_SIZE..FRAME_HEADER_SIZE + payload_len];
202    let payload = if flags.contains(Flags::COMPRESSED) {
203        // Decompress on read so the rest of the dispatch loop
204        // sees plaintext bytes regardless of how they arrived.
205        // `MAX_FRAME_SIZE` only bounds the *compressed* bytes; a
206        // streaming decoder capped at `MAX_DECOMPRESSED_PAYLOAD` keeps a
207        // 16 MiB zstd bomb from inflating into gigabytes of heap.
208        decompress_bounded(on_wire, payload_len)?
209    } else {
210        on_wire.to_vec()
211    };
212    Ok((
213        Frame {
214            kind,
215            // The flag stays on the decoded frame so dispatch can
216            // see it was compressed if it cares (audit, metrics).
217            flags,
218            stream_id,
219            correlation_id,
220            payload,
221        },
222        length as usize,
223    ))
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229
230    fn round_trip(frame: Frame) {
231        let bytes = encode_frame(&frame);
232        let (decoded, consumed) = decode_frame(&bytes).expect("decode");
233        assert_eq!(consumed, bytes.len());
234        assert_eq!(decoded, frame);
235    }
236
237    #[test]
238    fn round_trip_empty_payload() {
239        round_trip(Frame::new(MessageKind::Ping, 1, vec![]));
240    }
241
242    #[test]
243    fn frame_len_from_header_validates_bounds() {
244        let mut header = [0u8; FRAME_HEADER_SIZE];
245        header[..4].copy_from_slice(&(FRAME_HEADER_SIZE as u32).to_le_bytes());
246        assert_eq!(frame_len_from_header(&header).unwrap(), FRAME_HEADER_SIZE);
247
248        header[..4].copy_from_slice(&15u32.to_le_bytes());
249        assert_eq!(
250            frame_len_from_header(&header),
251            Err(FrameError::InvalidLength(15))
252        );
253
254        header[..4].copy_from_slice(&(MAX_FRAME_SIZE + 1).to_le_bytes());
255        assert_eq!(
256            frame_len_from_header(&header),
257            Err(FrameError::InvalidLength(MAX_FRAME_SIZE + 1))
258        );
259    }
260
261    #[test]
262    fn round_trip_with_payload() {
263        round_trip(Frame::new(MessageKind::Query, 42, b"SELECT 1".to_vec()));
264    }
265
266    #[test]
267    fn decode_frame_parts_matches_full_buffer_decode() {
268        let frame = Frame::new(MessageKind::Result, 42, br#"{"ok":true}"#.to_vec());
269        let bytes = encode_frame(&frame);
270        let mut header = [0u8; FRAME_HEADER_SIZE];
271        header.copy_from_slice(&bytes[..FRAME_HEADER_SIZE]);
272        let payload = &bytes[FRAME_HEADER_SIZE..];
273
274        let decoded = decode_frame_parts(&header, payload).expect("decode parts");
275        assert_eq!(decoded, frame);
276    }
277
278    #[test]
279    fn round_trip_with_stream_and_flags() {
280        let frame = Frame::new(MessageKind::Result, 999, vec![0xab; 256])
281            .with_stream(7)
282            .with_flags(Flags::COMPRESSED | Flags::MORE_FRAMES);
283        round_trip(frame);
284    }
285
286    #[test]
287    fn truncated_header_rejected() {
288        assert_eq!(decode_frame(&[]), Err(FrameError::Truncated));
289        assert_eq!(decode_frame(&[0; 15]), Err(FrameError::Truncated));
290    }
291
292    #[test]
293    fn length_below_header_rejected() {
294        let mut bytes = vec![0u8; 16];
295        bytes[..4].copy_from_slice(&15u32.to_le_bytes());
296        assert!(matches!(
297            decode_frame(&bytes),
298            Err(FrameError::InvalidLength(15))
299        ));
300    }
301
302    #[test]
303    fn unknown_kind_rejected() {
304        let mut bytes = vec![0u8; 16];
305        bytes[..4].copy_from_slice(&16u32.to_le_bytes());
306        bytes[4] = 0xff;
307        assert_eq!(decode_frame(&bytes), Err(FrameError::UnknownKind(0xff)));
308    }
309
310    #[test]
311    fn unknown_flag_bits_rejected() {
312        let mut bytes = vec![0u8; 16];
313        bytes[..4].copy_from_slice(&16u32.to_le_bytes());
314        bytes[4] = MessageKind::Ping as u8;
315        bytes[5] = 0b1000_0000;
316        assert!(matches!(
317            decode_frame(&bytes),
318            Err(FrameError::UnknownFlags(_))
319        ));
320    }
321
322    #[test]
323    fn flags_not_allowed_for_kind_rejected() {
324        // Ping is a handshake kind — the catalog forbids COMPRESSED
325        // on tiny handshake payloads. A frame with kind=Ping and the
326        // COMPRESSED bit set must be rejected at the boundary.
327        let mut bytes = vec![0u8; 16];
328        bytes[..4].copy_from_slice(&16u32.to_le_bytes());
329        bytes[4] = MessageKind::Ping as u8;
330        bytes[5] = Flags::COMPRESSED.bits();
331        match decode_frame(&bytes) {
332            Err(FrameError::FlagsNotAllowedForKind { kind, flags }) => {
333                assert_eq!(kind, MessageKind::Ping as u8);
334                assert_eq!(flags, Flags::COMPRESSED.bits());
335            }
336            other => panic!("expected FlagsNotAllowedForKind, got {other:?}"),
337        }
338    }
339
340    #[test]
341    fn streaming_decode_two_frames_back_to_back() {
342        let f1 = Frame::new(MessageKind::Query, 1, b"a".to_vec());
343        let f2 = Frame::new(MessageKind::Query, 2, b"b".to_vec());
344        let mut buf = encode_frame(&f1);
345        buf.extend(encode_frame(&f2));
346        let (got1, n1) = decode_frame(&buf).unwrap();
347        let (got2, _n2) = decode_frame(&buf[n1..]).unwrap();
348        assert_eq!(got1, f1);
349        assert_eq!(got2, f2);
350    }
351
352    #[test]
353    fn compressed_round_trip_recovers_plaintext() {
354        // A compressible payload — a kilobyte of repeating text.
355        let payload = b"abcabcabcabc".repeat(100);
356        let frame =
357            Frame::new(MessageKind::Result, 7, payload.clone()).with_flags(Flags::COMPRESSED);
358        let bytes = encode_frame(&frame);
359        // Wire form should be smaller than the plaintext frame.
360        assert!(
361            bytes.len() < FRAME_HEADER_SIZE + payload.len(),
362            "compressed frame ({}) must be smaller than plaintext payload ({})",
363            bytes.len(),
364            payload.len(),
365        );
366        let (decoded, _) = decode_frame(&bytes).expect("decode compressed");
367        assert_eq!(decoded.payload, payload);
368        assert!(decoded.flags.contains(Flags::COMPRESSED));
369    }
370
371    #[test]
372    fn output_stream_lifecycle_envelopes_round_trip() {
373        // Golden encode/decode for every new variant added in
374        // issue #762 / PRD #759 S3. Pins the exact byte values
375        // and confirms `stream_id` multiplex routing survives the
376        // round trip (it is what `StreamCancel`'s per-stream
377        // targeting relies on).
378        let open = Frame::new(
379            MessageKind::OpenStream,
380            10,
381            br#"{"sql":"SELECT 1","opts":{}}"#.to_vec(),
382        )
383        .with_stream(7);
384        round_trip(open.clone());
385        assert_eq!(encode_frame(&open)[4], 0x29);
386
387        let ack = Frame::new(
388            MessageKind::OpenAck,
389            10,
390            br#"{"lease_handle":"42","resumable":false,"snapshot_lsn":1234}"#.to_vec(),
391        )
392        .with_stream(7);
393        round_trip(ack.clone());
394        assert_eq!(encode_frame(&ack)[4], 0x2A);
395
396        let chunk = Frame::new(
397            MessageKind::StreamChunk,
398            10,
399            br#"{"seq":0,"rows":[{"a":1}],"terminal":false}"#.to_vec(),
400        )
401        .with_stream(7);
402        round_trip(chunk.clone());
403        assert_eq!(encode_frame(&chunk)[4], 0x2B);
404
405        let serr = Frame::new(
406            MessageKind::StreamError,
407            10,
408            br#"{"code":"unknown_stream","message":"x"}"#.to_vec(),
409        )
410        .with_stream(7);
411        round_trip(serr.clone());
412        assert_eq!(encode_frame(&serr)[4], 0x2C);
413
414        let end = Frame::new(
415            MessageKind::StreamEnd,
416            10,
417            br#"{"stats":{"row_count":1}}"#.to_vec(),
418        )
419        .with_stream(7);
420        round_trip(end.clone());
421        assert_eq!(encode_frame(&end)[4], 0x25);
422
423        let cancel = Frame::new(
424            MessageKind::StreamCancel,
425            10,
426            br#"{"reason":"client-abort"}"#.to_vec(),
427        )
428        .with_stream(7);
429        round_trip(cancel.clone());
430        assert_eq!(encode_frame(&cancel)[4], 0x2D);
431    }
432
433    #[test]
434    fn input_stream_envelopes_round_trip() {
435        // Golden encode/decode for the input-direction envelopes
436        // added in issue #764 / PRD #759 S5. The envelope *vocabulary*
437        // is reused from S3 — only the payload shapes and the
438        // direction of `StreamChunk` differ — so the byte values are
439        // pinned to the same kinds (no new MessageKind bytes).
440
441        // OpenStream with `direction:"in"` + target/columns instead of
442        // a `sql` field. Still kind 0x29, still multiplex via stream_id.
443        let open_in = Frame::new(
444            MessageKind::OpenStream,
445            20,
446            br#"{"direction":"in","target":"t","columns":["id","name"]}"#.to_vec(),
447        )
448        .with_stream(5);
449        round_trip(open_in.clone());
450        assert_eq!(encode_frame(&open_in)[4], 0x29);
451
452        // Client-originated chunk of rows on the input stream. Same
453        // 0x2B kind the server uses on output streams; the rows are
454        // JSON objects keyed by column.
455        let chunk_in = Frame::new(
456            MessageKind::StreamChunk,
457            20,
458            br#"{"seq":0,"rows":[{"id":1,"name":"a"}],"terminal":false}"#.to_vec(),
459        )
460        .with_stream(5);
461        round_trip(chunk_in.clone());
462        assert_eq!(encode_frame(&chunk_in)[4], 0x2B);
463
464        // Terminal chunk closes the input stream.
465        let chunk_terminal = Frame::new(
466            MessageKind::StreamChunk,
467            20,
468            br#"{"seq":2,"rows":[],"terminal":true}"#.to_vec(),
469        )
470        .with_stream(5);
471        round_trip(chunk_terminal.clone());
472        assert_eq!(encode_frame(&chunk_terminal)[4], 0x2B);
473
474        // Server StreamEnd carries the committed RID range + stats.
475        let end = Frame::new(
476            MessageKind::StreamEnd,
477            20,
478            br#"{"stats":{"row_count":3,"chunk_count":2,"committed_rid":42,"snapshot_lsn":40,"cancelled":false}}"#.to_vec(),
479        )
480        .with_stream(5);
481        round_trip(end.clone());
482        assert_eq!(encode_frame(&end)[4], 0x25);
483
484        // Server StreamError carries the recoverable_rid prefix.
485        let serr = Frame::new(
486            MessageKind::StreamError,
487            20,
488            br#"{"code":"invalid_row","message":"x","chunk_seq":1,"recoverable_rid":41}"#.to_vec(),
489        )
490        .with_stream(5);
491        round_trip(serr.clone());
492        assert_eq!(encode_frame(&serr)[4], 0x2C);
493    }
494
495    #[test]
496    fn queue_wait_envelopes_round_trip() {
497        // Golden encode/decode for the live queue-wait envelopes added
498        // in issue #917 / PRD #915. Pins the new byte values and
499        // confirms the request/push pair round-trips equal through the
500        // codec, multiplexed over `stream_id` like the other streamed
501        // envelopes.
502        let open = Frame::new(
503            MessageKind::QueueWaitOpen,
504            10,
505            br#"{"queue":"jobs","consumer":"w1","count":1,"wait_ms":5000}"#.to_vec(),
506        )
507        .with_stream(3);
508        round_trip(open.clone());
509        assert_eq!(encode_frame(&open)[4], 0x2E);
510
511        let push = Frame::new(
512            MessageKind::QueueEventPush,
513            10,
514            br#"{"message_id":"42","payload":{"hello":"world"},"consumer":"w1","delivery_count":1}"#
515                .to_vec(),
516        )
517        .with_stream(3);
518        round_trip(push.clone());
519        assert_eq!(encode_frame(&push)[4], 0x2F);
520    }
521
522    #[test]
523    fn uncompressed_frame_decodes_unchanged_when_flag_unset() {
524        let payload = b"hello world".to_vec();
525        let frame = Frame::new(MessageKind::Result, 1, payload.clone());
526        let bytes = encode_frame(&frame);
527        let (decoded, _) = decode_frame(&bytes).unwrap();
528        assert_eq!(decoded.payload, payload);
529        assert!(!decoded.flags.contains(Flags::COMPRESSED));
530    }
531
532    #[test]
533    fn compressed_payload_that_inflates_past_the_cap_is_refused() {
534        // MAX_FRAME_SIZE bounds only the *compressed* bytes. A highly
535        // compressible payload therefore lets a peer turn a small frame into
536        // a multi-gigabyte allocation unless the decoder is bounded.
537        let plain = vec![0u8; MAX_DECOMPRESSED_PAYLOAD + 1024];
538        let compressed = zstd::stream::encode_all(plain.as_slice(), 1).expect("compress");
539        assert!(
540            compressed.len() < MAX_FRAME_SIZE as usize,
541            "the bomb must fit in one legal frame to be a realistic attack"
542        );
543
544        let total = (FRAME_HEADER_SIZE + compressed.len()) as u32;
545        let mut bytes = Vec::with_capacity(total as usize);
546        bytes.extend_from_slice(&total.to_le_bytes());
547        bytes.push(MessageKind::Query as u8);
548        bytes.push(Flags::COMPRESSED.bits());
549        bytes.extend_from_slice(&0u16.to_le_bytes());
550        bytes.extend_from_slice(&0u64.to_le_bytes());
551        bytes.extend_from_slice(&compressed);
552
553        match decode_frame(&bytes) {
554            Err(FrameError::DecompressedTooLarge { limit }) => {
555                assert_eq!(limit as usize, MAX_DECOMPRESSED_PAYLOAD);
556            }
557            other => panic!("expected DecompressedTooLarge, got {other:?}"),
558        }
559    }
560}
reddb_wire/redwire/codec.rs

reddb_wire/redwire/
codec.rs
