hpke_http/
compression.rs

1//! Per-record zstd coding for clear payload bytes.
2
3use std::{borrow::Cow, io::Read};
4
5use ruzstd::{
6    decoding::StreamingDecoder,
7    encoding::{CompressionLevel, compress_to_vec},
8};
9
10use crate::Error;
11
12/// Raw DATA bytes.
13pub(crate) const RAW: u8 = 0;
14/// Zstd DATA bytes.
15pub(crate) const ZSTD: u8 = 1;
16
17/// Compress only when the coded record is shorter than its clear data.
18pub(crate) fn encode(body: &[u8]) -> (u8, Cow<'_, [u8]>) {
19    if body.len() < 64 {
20        return (RAW, Cow::Borrowed(body));
21    }
22    let coded = compress_to_vec(body, CompressionLevel::Fastest);
23    if coded.len() < body.len() {
24        (ZSTD, Cow::Owned(coded))
25    } else {
26        (RAW, Cow::Borrowed(body))
27    }
28}
29
30/// Decode one bounded record. The caller checks the total clear-body limit.
31pub(crate) fn decode(kind: u8, body: &[u8], max_clear: usize) -> Result<Vec<u8>, Error> {
32    match kind {
33        RAW => {
34            if body.is_empty() || body.len() > max_clear {
35                return Err(if body.is_empty() {
36                    Error::MalformedEnvelope
37                } else {
38                    Error::LimitExceeded
39                });
40            }
41            Ok(body.to_vec())
42        }
43        ZSTD => {
44            if body.is_empty() {
45                return Err(Error::MalformedEnvelope);
46            }
47            let mut decoder =
48                StreamingDecoder::new_with_max_window_size(body, max_clear.max(1024 * 1024) as u64)
49                    .map_err(|_| Error::MalformedEnvelope)?;
50            let declared_size = decoder.decoder.content_size();
51            if declared_size > max_clear as u64 {
52                return Err(Error::LimitExceeded);
53            }
54            let mut output = Vec::new();
55            let mut buffer = [0_u8; 8192];
56            loop {
57                let count = decoder
58                    .read(&mut buffer)
59                    .map_err(|_| Error::MalformedEnvelope)?;
60                if count == 0 {
61                    break;
62                }
63                if count > max_clear.saturating_sub(output.len()) {
64                    return Err(Error::LimitExceeded);
65                }
66                output.extend_from_slice(&buffer[..count]);
67            }
68            if declared_size != 0 && declared_size != output.len() as u64 {
69                return Err(Error::MalformedEnvelope);
70            }
71            if decoder.decoder.get_checksum_from_data().is_some()
72                && decoder.decoder.get_checksum_from_data()
73                    != decoder.decoder.get_calculated_checksum()
74            {
75                return Err(Error::MalformedEnvelope);
76            }
77            if output.is_empty() || !decoder.into_inner().is_empty() {
78                return Err(Error::MalformedEnvelope);
79            }
80            Ok(output)
81        }
82        _ => Err(Error::MalformedEnvelope),
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::{RAW, ZSTD, decode, encode};
89    use crate::Error;
90    use ruzstd::encoding::{CompressionLevel, compress_to_vec};
91
92    #[test]
93    fn chooses_raw_or_zstd_and_bounds_expansion() {
94        let small = b"abc";
95        assert_eq!(encode(small).0, RAW);
96        let large = vec![b'a'; 4096];
97        let (kind, coded) = encode(&large);
98        assert_eq!(kind, ZSTD);
99        assert_eq!(decode(kind, &coded, large.len()), Ok(large.clone()));
100        assert_eq!(decode(kind, &coded, 128), Err(Error::LimitExceeded));
101        assert_eq!(decode(9, &coded, 4096), Err(Error::MalformedEnvelope));
102        assert_eq!(
103            decode(ZSTD, &coded[..coded.len() - 1], 4096),
104            Err(Error::MalformedEnvelope)
105        );
106        let mut with_extra = coded.clone().into_owned();
107        with_extra.push(0);
108        assert_eq!(
109            decode(ZSTD, &with_extra, 4096),
110            Err(Error::MalformedEnvelope)
111        );
112        let mut corrupt = coded.into_owned();
113        corrupt[0] ^= 1;
114        assert_eq!(decode(ZSTD, &corrupt, 4096), Err(Error::MalformedEnvelope));
115    }
116
117    #[test]
118    fn accepts_valid_zstd_even_when_it_does_not_save_bytes() {
119        let clear = b"abc";
120        let coded = compress_to_vec(clear.as_slice(), CompressionLevel::Fastest);
121        assert!(coded.len() >= clear.len());
122        assert_eq!(decode(ZSTD, &coded, clear.len()), Ok(clear.to_vec()));
123    }
124
125    #[test]
126    fn rejects_bad_checksum_and_oversized_window() {
127        // Python zstandard made this single frame with a checksum and 4096 'a' bytes.
128        let mut checksum_frame = [
129            0x28, 0xb5, 0x2f, 0xfd, 0x64, 0x00, 0x0f, 0x4d, 0x00, 0x00, 0x10, 0x61, 0x61, 0x01,
130            0x00, 0xfb, 0xf7, 0x01, 0x16, 0x03, 0x16, 0x2c, 0x89,
131        ];
132        assert_eq!(decode(ZSTD, &checksum_frame, 4096), Ok(vec![b'a'; 4096]));
133        checksum_frame[22] ^= 1;
134        assert_eq!(
135            decode(ZSTD, &checksum_frame, 4096),
136            Err(Error::MalformedEnvelope)
137        );
138
139        // This valid frame declares a 2 MiB window but expands to only 4096 bytes.
140        let large_window_frame = [
141            0x28, 0xb5, 0x2f, 0xfd, 0x00, 0x58, 0x4d, 0x00, 0x00, 0x10, 0x61, 0x61, 0x01, 0x00,
142            0xfb, 0xf7, 0x01, 0x16,
143        ];
144        assert_eq!(
145            decode(ZSTD, &large_window_frame, 4096),
146            Err(Error::MalformedEnvelope)
147        );
148        assert_eq!(
149            decode(ZSTD, &large_window_frame, 2 * 1024 * 1024),
150            Ok(vec![b'a'; 4096])
151        );
152    }
153}
hpke_http/compression.rs

hpke_http/
compression.rs
