dialcache/protocol/
envelope.rs

1//! Payload envelopes: escaping and zstd compression.
2//!
3//! Byte 0 of a binary payload written by this version or later selects the
4//! envelope: `0x00` escapes raw serializer output whose own first byte would
5//! collide, `0x01` marks a zstd frame of UTF-8 text and `0x02` a zstd frame of
6//! binary bytes. Raw output is escaped on every write, and every read
7//! interprets the envelope, so disabling compression never strands entries.
8//! See "Envelope selection and codec environment" in `formal/guides/PROTOCOL.md`.
9
10use crate::codec::Payload;
11use crate::limits::{DEFAULT_COMPRESSION_THRESHOLD_BYTES, DEFAULT_ZSTD_LEVEL, MAX_SAFE_INTEGER};
12use crate::observe::CompressionOutcome;
13
14use super::frame::ProtocolError;
15use super::text::replacement_utf8;
16
17/// Envelope byte `0x00`: the rest is raw serializer output whose own first
18/// byte would have collided with a marker.
19pub const MARKER_ESCAPED_RAW: u8 = 0x00;
20/// Envelope byte `0x01`: the rest is one zstd frame of UTF-8 text.
21pub const MARKER_ZSTD_UTF8: u8 = 0x01;
22/// Envelope byte `0x02`: the rest is one zstd frame of binary bytes.
23pub const MARKER_ZSTD_BINARY: u8 = 0x02;
24
25const MIN_ZSTD_LEVEL: i32 = 1;
26const MAX_ZSTD_LEVEL: i32 = 22;
27
28/// Write-side compression policy.
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub struct CompressionConfig {
31    /// Payloads of at least this many serialized bytes are compressed.
32    pub threshold_bytes: usize,
33    /// zstd level, `1..=22`.
34    pub level: i32,
35}
36
37impl Default for CompressionConfig {
38    fn default() -> Self {
39        CompressionConfig {
40            threshold_bytes: DEFAULT_COMPRESSION_THRESHOLD_BYTES,
41            level: DEFAULT_ZSTD_LEVEL,
42        }
43    }
44}
45
46impl CompressionConfig {
47    /// The threshold must be a positive safe integer and the level `1..=22`.
48    pub fn validate(&self) -> Result<(), ProtocolError> {
49        if self.threshold_bytes < 1 || self.threshold_bytes as u64 > MAX_SAFE_INTEGER {
50            return Err(ProtocolError::Compression(
51                "threshold_bytes must be a positive safe integer".to_string(),
52            ));
53        }
54        if !(MIN_ZSTD_LEVEL..=MAX_ZSTD_LEVEL).contains(&self.level) {
55            return Err(ProtocolError::Compression(
56                "level must be an integer between 1 and 22".to_string(),
57            ));
58        }
59        Ok(())
60    }
61}
62
63/// What [`compress_payload`] chose to store, and why.
64#[derive(Debug, Clone, PartialEq, Eq)]
65pub struct CompressionWriteResult {
66    /// The bytes to store: marked zstd output, or the escaped raw form.
67    pub payload: Payload,
68    /// One of the write outcomes: compressed, below threshold, not smaller
69    /// or over limit.
70    pub outcome: CompressionOutcome,
71    /// Length of the serializer output after text replacement, before escaping.
72    pub original_bytes: usize,
73    /// Length of `payload`.
74    pub stored_bytes: usize,
75}
76
77/// What [`decompress_payload`] produced from a stored payload.
78#[derive(Debug, Clone, PartialEq, Eq)]
79pub struct CompressionReadResult {
80    /// The decoded payload; the input with its `0x00` escape prefix removed
81    /// when it was escaped raw output; or the input unchanged when it passed
82    /// through untouched or zstd rejected it.
83    pub payload: Payload,
84    /// `None` when no zstd envelope was involved: the payload passed through
85    /// untouched or only had its escape prefix removed.
86    pub outcome: Option<CompressionOutcome>,
87}
88
89fn needs_raw_escape(payload: &Payload) -> bool {
90    payload.binary
91        && payload
92            .bytes
93            .first()
94            .is_some_and(|first| *first <= MARKER_ZSTD_BINARY)
95}
96
97/// Prefix raw binary output whose first byte would collide with the envelope.
98pub fn escape_raw_payload(payload: Payload) -> Payload {
99    if needs_raw_escape(&payload) {
100        let mut escaped = Vec::with_capacity(payload.bytes.len() + 1);
101        escaped.push(MARKER_ESCAPED_RAW);
102        escaped.extend_from_slice(&payload.bytes);
103        return Payload::binary(escaped);
104    }
105    payload
106}
107
108/// Compress when the payload meets the threshold and the marked result is
109/// strictly smaller than the escaped raw form.
110///
111/// Text payloads receive replacement UTF-8 conversion first, so the measured
112/// size is the stored text's. The threshold comparison precedes the output
113/// cap: a value the read side would refuse is stored raw (`write_over_limit`)
114/// and subject to Redis's own value limit instead. `original_bytes` is the
115/// payload's own length and `stored_bytes` the length actually written.
116pub fn compress_payload(
117    payload: Payload,
118    config: &CompressionConfig,
119    max_decompressed_bytes: usize,
120) -> Result<CompressionWriteResult, ProtocolError> {
121    config.validate()?;
122    let payload = if payload.binary {
123        payload
124    } else {
125        Payload {
126            bytes: replacement_utf8(&payload.bytes),
127            binary: false,
128        }
129    };
130    let original_bytes = payload.len();
131    let outcome = if original_bytes < config.threshold_bytes {
132        CompressionOutcome::BelowThreshold
133    } else if original_bytes > max_decompressed_bytes {
134        CompressionOutcome::WriteOverLimit
135    } else {
136        let compressed = zstd_compress(&payload.bytes, config.level)?;
137        let raw_bytes = original_bytes + usize::from(needs_raw_escape(&payload));
138        if compressed.len() + 1 < raw_bytes {
139            let mut marked = Vec::with_capacity(compressed.len() + 1);
140            marked.push(if payload.binary {
141                MARKER_ZSTD_BINARY
142            } else {
143                MARKER_ZSTD_UTF8
144            });
145            marked.extend_from_slice(&compressed);
146            return Ok(CompressionWriteResult {
147                stored_bytes: marked.len(),
148                payload: Payload::binary(marked),
149                outcome: CompressionOutcome::Compressed,
150                original_bytes,
151            });
152        }
153        CompressionOutcome::NotSmaller
154    };
155    let raw = escape_raw_payload(payload);
156    Ok(CompressionWriteResult {
157        stored_bytes: raw.len(),
158        payload: raw,
159        outcome,
160        original_bytes,
161    })
162}
163
164/// Reverse of [`compress_payload`], applied to every read.
165///
166/// Text payloads and empty binary payloads pass through. An escape prefix is
167/// removed only when the next byte is one the writer escapes; other `0x00`
168/// leads and unknown markers pass through unchanged. A marked payload decodes
169/// only its first zstd frame and ignores trailing bytes, like Node's
170/// synchronous decoder. Malformed or truncated zstd data keeps the original
171/// marked bytes with `fallback_raw`; output beyond `max_decompressed_bytes`
172/// (or a frame demanding more decoder memory than allowed) keeps them with
173/// `read_over_limit`. Successful `0x01` output receives replacement UTF-8
174/// conversion; `0x02` output is exact bytes. The input is never mutated, so
175/// repeated loads of a retained payload are independent.
176pub fn decompress_payload(
177    payload: Payload,
178    max_decompressed_bytes: usize,
179) -> CompressionReadResult {
180    let passthrough = |payload| CompressionReadResult {
181        payload,
182        outcome: None,
183    };
184    if !payload.binary {
185        return passthrough(payload);
186    }
187    let Some(&marker) = payload.bytes.first() else {
188        return passthrough(payload);
189    };
190    match marker {
191        MARKER_ESCAPED_RAW => {
192            if payload
193                .bytes
194                .get(1)
195                .is_some_and(|next| *next <= MARKER_ZSTD_BINARY)
196            {
197                return passthrough(Payload::binary(&payload.bytes[1..]));
198            }
199            passthrough(payload)
200        }
201        MARKER_ZSTD_UTF8 | MARKER_ZSTD_BINARY => {
202            match decode_first_zstd_frame(&payload.bytes[1..], max_decompressed_bytes) {
203                Ok(decoded) => {
204                    let decoded = if marker == MARKER_ZSTD_UTF8 {
205                        Payload {
206                            bytes: replacement_utf8(&decoded),
207                            binary: false,
208                        }
209                    } else {
210                        Payload::binary(decoded)
211                    };
212                    CompressionReadResult {
213                        payload: decoded,
214                        outcome: Some(CompressionOutcome::Decompressed),
215                    }
216                }
217                Err(outcome) => CompressionReadResult {
218                    payload,
219                    outcome: Some(outcome),
220                },
221            }
222        }
223        _ => passthrough(payload),
224    }
225}
226
227/// One zstd frame of `bytes` at `level`, without a checksum, as the writer
228/// stores it (single-shot, so the frame header carries the content size).
229fn zstd_compress(bytes: &[u8], level: i32) -> Result<Vec<u8>, ProtocolError> {
230    let codec =
231        |code: usize| ProtocolError::Compression(zstd_safe::get_error_name(code).to_string());
232    let mut context = zstd_safe::CCtx::create();
233    context
234        .set_parameter(zstd_safe::CParameter::CompressionLevel(level))
235        .map_err(codec)?;
236    context
237        .set_parameter(zstd_safe::CParameter::ChecksumFlag(false))
238        .map_err(codec)?;
239    let mut out = vec![0u8; zstd_safe::compress_bound(bytes.len())];
240    let written = context.compress2(&mut out[..], bytes).map_err(codec)?;
241    out.truncate(written);
242    Ok(out)
243}
244
245/// Decode only the first zstd frame of `body`, ignoring any trailer, with the
246/// output cap enforced while streaming so a small stored frame can never force
247/// a large allocation. Skippable frames decode to empty output.
248fn decode_first_zstd_frame(
249    body: &[u8],
250    max_decompressed_bytes: usize,
251) -> Result<Vec<u8>, CompressionOutcome> {
252    // Rejects truncated and malformed headers or block sequences up front,
253    // like Go's firstZstdFrame; only a complete first frame is decoded.
254    let frame_len =
255        zstd_safe::find_frame_compressed_size(body).map_err(|_| CompressionOutcome::FallbackRaw)?;
256    let frame = body
257        .get(..frame_len)
258        .ok_or(CompressionOutcome::FallbackRaw)?;
259
260    let mut context = zstd_safe::DCtx::create();
261    let mut out = Vec::new();
262    let mut input = zstd_safe::InBuffer::around(frame);
263    // Each step writes at most one chunk; a cap smaller than a chunk still
264    // only needs one chunk to observe the overflow.
265    let chunk_len = zstd_safe::DCtx::out_size().max(1);
266    let mut chunk = vec![0u8; chunk_len];
267    loop {
268        let mut output = zstd_safe::OutBuffer::around(&mut chunk[..]);
269        let consumed_before = input.pos();
270        let hint = context
271            .decompress_stream(&mut output, &mut input)
272            .map_err(classify_decoder_error)?;
273        let produced = output.as_slice();
274        if out.len() + produced.len() > max_decompressed_bytes {
275            return Err(CompressionOutcome::ReadOverLimit);
276        }
277        out.extend_from_slice(produced);
278        if hint == 0 {
279            return Ok(out);
280        }
281        if input.pos() == frame.len() && produced.is_empty() {
282            // The verified frame ended without the decoder finishing it.
283            return Err(CompressionOutcome::FallbackRaw);
284        }
285        if input.pos() == consumed_before && produced.is_empty() {
286            return Err(CompressionOutcome::FallbackRaw);
287        }
288    }
289}
290
291/// zstd errors that mean the frame exceeds the decoder's resource limits are
292/// read-limit outcomes; anything else falls back to the raw bytes.
293fn classify_decoder_error(code: usize) -> CompressionOutcome {
294    let name = zstd_safe::get_error_name(code);
295    if name.contains("too much memory") || name.contains("Destination buffer is too small") {
296        CompressionOutcome::ReadOverLimit
297    } else {
298        CompressionOutcome::FallbackRaw
299    }
300}
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305    use crate::limits::MAX_DECOMPRESSED_BYTES;
306
307    fn text(len: usize) -> Payload {
308        Payload::text("x".repeat(len))
309    }
310
311    fn compressed(payload: Payload) -> CompressionWriteResult {
312        let result = compress_payload(
313            payload,
314            &CompressionConfig {
315                threshold_bytes: 1,
316                level: 3,
317            },
318            MAX_DECOMPRESSED_BYTES,
319        )
320        .unwrap();
321        assert_eq!(result.outcome, CompressionOutcome::Compressed);
322        result
323    }
324
325    #[test]
326    fn config_validation() {
327        assert_eq!(CompressionConfig::default().validate(), Ok(()));
328        assert!(CompressionConfig {
329            threshold_bytes: 0,
330            level: 3
331        }
332        .validate()
333        .is_err());
334        assert!(CompressionConfig {
335            threshold_bytes: usize::MAX,
336            level: 3
337        }
338        .validate()
339        .is_err());
340        assert!(CompressionConfig {
341            threshold_bytes: 1,
342            level: 0
343        }
344        .validate()
345        .is_err());
346        assert!(CompressionConfig {
347            threshold_bytes: 1,
348            level: 23
349        }
350        .validate()
351        .is_err());
352        assert_eq!(
353            CompressionConfig {
354                threshold_bytes: 1,
355                level: 22
356            }
357            .validate(),
358            Ok(())
359        );
360        let invalid = compress_payload(
361            text(1),
362            &CompressionConfig {
363                threshold_bytes: 1,
364                level: 23,
365            },
366            MAX_DECOMPRESSED_BYTES,
367        );
368        assert!(matches!(invalid, Err(ProtocolError::Compression(_))));
369    }
370
371    #[test]
372    fn escaping_and_passthrough() {
373        assert_eq!(
374            escape_raw_payload(Payload::binary(vec![0, 9])).bytes,
375            [0, 0, 9]
376        );
377        assert_eq!(escape_raw_payload(Payload::binary(vec![2])).bytes, [0, 2]);
378        assert_eq!(escape_raw_payload(Payload::binary(vec![3])).bytes, [3]);
379        assert_eq!(
380            escape_raw_payload(Payload::binary(Vec::new())).bytes,
381            Vec::<u8>::new()
382        );
383        assert_eq!(escape_raw_payload(Payload::text("\u{0}")).bytes, [0]);
384        // Only a prefix the writer produces is stripped; legacy 0x00 leads stay.
385        assert_eq!(
386            decompress_payload(Payload::binary(vec![0, 1, 7]), 10),
387            CompressionReadResult {
388                payload: Payload::binary(vec![1, 7]),
389                outcome: None
390            }
391        );
392        for untouched in [
393            Payload::binary(vec![0, 3]),
394            Payload::binary(vec![0]),
395            Payload::binary(vec![3, 1]),
396            Payload::binary(Vec::new()),
397            Payload::text("\u{1}x"),
398        ] {
399            assert_eq!(
400                decompress_payload(untouched.clone(), 10),
401                CompressionReadResult {
402                    payload: untouched,
403                    outcome: None
404                }
405            );
406        }
407    }
408
409    /// Port of go/internal/testsuite/codec_test.go TestCompressionLimitsAndIndependentLoads.
410    #[test]
411    fn compression_limits_and_independent_loads() {
412        let raw = text(4096);
413        let written = compressed(raw.clone());
414        let too_large = decompress_payload(written.payload.clone(), 4095);
415        assert_eq!(too_large.outcome, Some(CompressionOutcome::ReadOverLimit));
416        assert_eq!(too_large.payload, written.payload);
417        let exact = decompress_payload(written.payload.clone(), 4096);
418        assert_eq!(
419            exact,
420            CompressionReadResult {
421                payload: raw.clone(),
422                outcome: Some(CompressionOutcome::Decompressed)
423            }
424        );
425        let refused = compress_payload(
426            raw.clone(),
427            &CompressionConfig {
428                threshold_bytes: 1,
429                level: 3,
430            },
431            4095,
432        )
433        .unwrap();
434        assert_eq!(refused.outcome, CompressionOutcome::WriteOverLimit);
435        assert_eq!(refused.payload, raw);
436        assert_eq!((refused.original_bytes, refused.stored_bytes), (4096, 4096));
437        // Repeated loads of a retained input are independent.
438        let mut first = decompress_payload(written.payload.clone(), MAX_DECOMPRESSED_BYTES);
439        first.payload.bytes[0] = b'y';
440        assert_eq!(
441            decompress_payload(written.payload.clone(), MAX_DECOMPRESSED_BYTES).payload,
442            raw
443        );
444        let escaped = escape_raw_payload(Payload::binary(vec![1, 2, 3]));
445        let mut first = decompress_payload(escaped.clone(), MAX_DECOMPRESSED_BYTES);
446        first.payload.bytes[0] = 9;
447        assert_eq!(
448            decompress_payload(escaped, MAX_DECOMPRESSED_BYTES)
449                .payload
450                .bytes,
451            [1, 2, 3]
452        );
453        // A zero cap admits only empty output.
454        assert_eq!(
455            decompress_payload(written.payload.clone(), 0).outcome,
456            Some(CompressionOutcome::ReadOverLimit)
457        );
458    }
459
460    /// Port of go/internal/testsuite/codec_test.go TestZstdFirstFrameAndMalformedBodies.
461    #[test]
462    fn first_frame_and_malformed_bodies() {
463        let source = Payload::text("first".repeat(100));
464        let written = compressed(source.clone());
465        for trailer in [&b"CRC!"[..], &written.payload.bytes[1..]] {
466            let mut marked = written.payload.bytes.clone();
467            marked.extend_from_slice(trailer);
468            let result = decompress_payload(Payload::binary(marked), MAX_DECOMPRESSED_BYTES);
469            assert_eq!(
470                result,
471                CompressionReadResult {
472                    payload: source.clone(),
473                    outcome: Some(CompressionOutcome::Decompressed)
474                }
475            );
476        }
477        let half = written.payload.bytes[..written.payload.bytes.len() / 2].to_vec();
478        for marked in [
479            vec![1],
480            vec![2],
481            half,
482            vec![1, 0x28, 0xB5, 0x2F, 0xFD],
483            vec![2, 0xFF, 0xFF],
484        ] {
485            let got = decompress_payload(Payload::binary(marked.clone()), MAX_DECOMPRESSED_BYTES);
486            assert_eq!(
487                got,
488                CompressionReadResult {
489                    payload: Payload::binary(marked),
490                    outcome: Some(CompressionOutcome::FallbackRaw)
491                }
492            );
493        }
494    }
495
496    #[test]
497    fn skippable_frames_decode_to_empty_output() {
498        // Magic 0x184D2A50, 4-byte little-endian size, then that many bytes.
499        let skippable = [
500            0x01, 0x50, 0x2A, 0x4D, 0x18, 0x02, 0x00, 0x00, 0x00, 0xAA, 0xBB, 0xCC,
501        ];
502        let got = decompress_payload(Payload::binary(skippable.to_vec()), MAX_DECOMPRESSED_BYTES);
503        assert_eq!(
504            got,
505            CompressionReadResult {
506                payload: Payload::text(""),
507                outcome: Some(CompressionOutcome::Decompressed)
508            }
509        );
510        // Zero cap with empty output is still a successful decode.
511        assert_eq!(
512            decompress_payload(Payload::binary(skippable.to_vec()), 0).outcome,
513            Some(CompressionOutcome::Decompressed)
514        );
515        let truncated = skippable[..8].to_vec();
516        assert_eq!(
517            decompress_payload(Payload::binary(truncated), MAX_DECOMPRESSED_BYTES).outcome,
518            Some(CompressionOutcome::FallbackRaw)
519        );
520    }
521
522    #[test]
523    fn selection_rule_and_types() {
524        let config = CompressionConfig {
525            threshold_bytes: 4,
526            level: 3,
527        };
528        let below = compress_payload(text(3), &config, MAX_DECOMPRESSED_BYTES).unwrap();
529        assert_eq!(
530            (below.outcome, below.original_bytes, below.stored_bytes),
531            (CompressionOutcome::BelowThreshold, 3, 3)
532        );
533        let incompressible = compress_payload(
534            Payload::binary(vec![1, 2, 3, 4]),
535            &config,
536            MAX_DECOMPRESSED_BYTES,
537        )
538        .unwrap();
539        assert_eq!(incompressible.outcome, CompressionOutcome::NotSmaller);
540        assert_eq!(incompressible.payload.bytes, [0, 1, 2, 3, 4]);
541        assert_eq!(
542            (incompressible.original_bytes, incompressible.stored_bytes),
543            (4, 5)
544        );
545        let binary = compressed(Payload::binary(vec![7; 500]));
546        assert_eq!(binary.payload.bytes[0], MARKER_ZSTD_BINARY);
547        assert!(binary.payload.binary && binary.stored_bytes < 500);
548        assert_eq!(
549            decompress_payload(binary.payload, MAX_DECOMPRESSED_BYTES).payload,
550            Payload::binary(vec![7; 500])
551        );
552        let text_result = compressed(text(500));
553        assert_eq!(text_result.payload.bytes[0], MARKER_ZSTD_UTF8);
554        assert_eq!(
555            decompress_payload(text_result.payload, MAX_DECOMPRESSED_BYTES).payload,
556            text(500)
557        );
558        // Ill-formed text is replaced before measuring and compressing.
559        let ill_formed = Payload {
560            bytes: vec![0xE2, 0x82, b'A'],
561            binary: false,
562        };
563        let stored = compress_payload(
564            ill_formed,
565            &CompressionConfig {
566                threshold_bytes: 100,
567                level: 3,
568            },
569            MAX_DECOMPRESSED_BYTES,
570        )
571        .unwrap();
572        assert_eq!(stored.payload, Payload::text("\u{FFFD}A"));
573        assert_eq!((stored.original_bytes, stored.stored_bytes), (4, 4));
574        // Decoded 0x01 text is replaced; 0x02 bytes are exact.
575        let mut marked_text =
576            compressed(Payload::binary([0xE2, 0x82, b'A', 0xFF].repeat(50))).payload;
577        marked_text.bytes[0] = MARKER_ZSTD_UTF8;
578        let decoded = decompress_payload(marked_text, MAX_DECOMPRESSED_BYTES);
579        assert_eq!(
580            decoded.payload,
581            Payload::text("\u{FFFD}A\u{FFFD}".repeat(50))
582        );
583    }
584}
dialcache/protocol/envelope.rs

dialcache/protocol/
envelope.rs
