sark_core/http/
compress.rs

1//! ```no_run
2//! use sark_core::http::compress::Gzip;
3//! let body: &[u8] = b"hello";
4//! let zipped = Gzip::new().encode(body).unwrap();
5//! ```
6
7use libdeflater::{CompressionLvl, Compressor, DecompressionError, Decompressor};
8use o3::buffer::{
9    Bytes, Initialized, PoolLayoutError, Pooled, Retained, SharedPool, SharedPoolLayout,
10};
11use thiserror::Error;
12
13use crate::http::Body;
14
15const GZIP_SLOTS: usize = 32;
16const GZIP_CAPACITY: usize = 256 * 1024;
17
18pub struct Gzip {
19    inner: Option<GzipInner>,
20    layout: SharedPoolLayout,
21}
22
23struct GzipInner {
24    encoder: Compressor,
25    pool: SharedPool<Initialized>,
26}
27
28impl Gzip {
29    const LEVEL: i32 = 3;
30    const COMPRESSION_LEVEL: CompressionLvl = match CompressionLvl::new(Self::LEVEL) {
31        Ok(level) => level,
32        Err(_) => CompressionLvl::fastest(),
33    };
34
35    pub fn new() -> Self {
36        Self {
37            inner: None,
38            layout: SharedPoolLayout::fixed::<GZIP_SLOTS, GZIP_CAPACITY>(),
39        }
40    }
41
42    pub fn with_pool(slots: usize, capacity: usize) -> Result<Self, PoolLayoutError> {
43        Ok(Self {
44            inner: None,
45            layout: SharedPoolLayout::new(slots, capacity)?,
46        })
47    }
48
49    fn inner(&mut self) -> &mut GzipInner {
50        let layout = self.layout;
51        self.inner.get_or_insert_with(|| GzipInner {
52            encoder: Compressor::new(Self::COMPRESSION_LEVEL),
53            pool: SharedPool::<Initialized>::from_layout(layout),
54        })
55    }
56
57    pub fn encode(&mut self, src: &[u8]) -> Option<Pooled> {
58        let inner = self.inner();
59        let cap = inner.encoder.gzip_compress_bound(src.len());
60        if cap > inner.pool.capacity() {
61            return None;
62        }
63        let mut lease = inner.pool.try_acquire()?;
64        let n = inner.encoder.gzip_compress(src, lease.spare_mut()).ok()?;
65        lease.try_advance(n).ok()?;
66        Some(lease.freeze())
67    }
68}
69
70impl Default for Gzip {
71    fn default() -> Self {
72        Self::new()
73    }
74}
75
76pub enum GunzipOutput {
77    Pooled(Pooled),
78    Owned(Vec<u8>),
79}
80
81impl From<GunzipOutput> for Body<'static> {
82    fn from(output: GunzipOutput) -> Self {
83        match output {
84            GunzipOutput::Pooled(body) => Body::from(Bytes::<Retained>::from(body)),
85            GunzipOutput::Owned(body) => Body::from(body),
86        }
87    }
88}
89
90#[derive(Debug, Error)]
91pub enum GunzipError {
92    #[error("gzip decompression failed: {0}")]
93    Invalid(#[from] DecompressionError),
94    #[error("decompressed response body exceeds size limit")]
95    SizeLimit,
96}
97
98pub struct Gunzip {
99    decoder: Decompressor,
100    pool: SharedPool<Initialized>,
101}
102
103impl Gunzip {
104    pub fn new() -> Self {
105        Self {
106            decoder: Decompressor::new(),
107            pool: SharedPool::<Initialized>::from_layout(SharedPoolLayout::fixed::<
108                GZIP_SLOTS,
109                GZIP_CAPACITY,
110            >()),
111        }
112    }
113
114    pub fn with_pool(slots: usize, capacity: usize) -> Result<Self, PoolLayoutError> {
115        Ok(Self {
116            decoder: Decompressor::new(),
117            pool: SharedPool::<Initialized>::from_layout(SharedPoolLayout::new(slots, capacity)?),
118        })
119    }
120
121    pub fn decode(&mut self, src: &[u8], max_size: usize) -> Result<GunzipOutput, GunzipError> {
122        let expected = Self::decoded_size(src)?;
123        if expected > max_size {
124            return Err(GunzipError::SizeLimit);
125        }
126
127        if expected <= self.pool.capacity()
128            && let Some(mut lease) = self.pool.try_acquire()
129        {
130            let written = self
131                .decoder
132                .gzip_decompress(src, &mut lease.spare_mut()[..expected])?;
133            lease
134                .try_advance(written)
135                .map_err(|_| GunzipError::SizeLimit)?;
136            return Ok(GunzipOutput::Pooled(lease.freeze()));
137        }
138
139        let mut output = vec![0; expected];
140        let written = self.decoder.gzip_decompress(src, &mut output)?;
141        output.truncate(written);
142        Ok(GunzipOutput::Owned(output))
143    }
144
145    fn decoded_size(src: &[u8]) -> Result<usize, GunzipError> {
146        if src.len() < 18 || !src.starts_with(&[0x1f, 0x8b, 0x08]) {
147            return Err(DecompressionError::BadData.into());
148        }
149        let mut footer = [0; 4];
150        footer.copy_from_slice(&src[src.len() - 4..]);
151        Ok(u32::from_le_bytes(footer) as usize)
152    }
153}
154
155impl Default for Gunzip {
156    fn default() -> Self {
157        Self::new()
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use super::Gzip;
164
165    #[test]
166    fn gzip_allocates_its_pool_on_first_use() {
167        let mut gzip = Gzip::new();
168        assert!(gzip.inner.is_none());
169
170        assert!(gzip.encode(b"hello").is_some());
171        assert!(gzip.inner.is_some());
172    }
173}
sark_core/http/compress.rs

sark_core/http/
compress.rs
