rws_http/compression/
dictionary.rs1use crate::header::Header;
30use crate::request::Request;
31use crate::response::Response;
32
33use super::Level;
34
35pub const DCZ_MAGIC: [u8; 8] = [0x5e, 0x2a, 0x4d, 0x18, 0x20, 0x00, 0x00, 0x00];
37pub const DCB_MAGIC: [u8; 4] = [0xff, 0x44, 0x43, 0x42];
39pub const MAX_SIZE: usize = 1024 * 1024;
42
43pub type DictionaryHash = [u8; 32];
45
46#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
48pub enum DictionaryEncoding {
49 Dcz,
51 Dcb,
53}
54
55impl DictionaryEncoding {
56 pub fn token(self) -> &'static str {
58 match self {
59 DictionaryEncoding::Dcz => "dcz",
60 DictionaryEncoding::Dcb => "dcb",
61 }
62 }
63
64 pub fn available(self) -> bool {
66 match self {
67 DictionaryEncoding::Dcz => cfg!(feature = "zstd"),
68 DictionaryEncoding::Dcb => cfg!(feature = "brotli"),
69 }
70 }
71
72 fn magic(self) -> &'static [u8] {
73 match self {
74 DictionaryEncoding::Dcz => &DCZ_MAGIC,
75 DictionaryEncoding::Dcb => &DCB_MAGIC,
76 }
77 }
78
79 pub fn header_len(self) -> usize {
81 self.magic().len() + 32
82 }
83}
84
85pub const PREFERENCE: &[DictionaryEncoding] = &[DictionaryEncoding::Dcz, DictionaryEncoding::Dcb];
88
89pub fn hash(dictionary: &[u8]) -> DictionaryHash {
91 crate::core::sha256::sha256(dictionary)
92}
93
94pub fn parse_available_dictionary(value: &str) -> Option<DictionaryHash> {
97 let rest = value.trim().strip_prefix(':')?;
98 let end = rest.find(':')?;
99 let (encoded, after) = (&rest[..end], &rest[end + 1..]);
100 if !(after.is_empty() || after.starts_with(';')) || encoded.len() != 44 {
101 return None;
102 }
103 if !encoded.bytes().all(|b| b.is_ascii_alphanumeric() || matches!(b, b'+' | b'/' | b'=')) {
104 return None;
105 }
106 let bytes = crate::core::base64::Base64::decode(encoded.to_string()).ok()?;
107 bytes.try_into().ok()
108}
109
110pub fn available_dictionary(request: &Request) -> Option<DictionaryHash> {
112 parse_available_dictionary(&request.header("Available-Dictionary")?.value)
113}
114
115pub fn negotiate(request: &Request) -> Option<DictionaryEncoding> {
119 let header = request.header(Header::_ACCEPT_ENCODING)?;
120 let mut best: Option<(DictionaryEncoding, f32)> = None;
121 for &encoding in PREFERENCE.iter().filter(|e| e.available()) {
122 let listed = header.value.split(',').any(|c| c.split(';').next().unwrap_or("").trim().eq_ignore_ascii_case(encoding.token()));
123 let q = if listed { super::accepted_weight(&header.value, encoding.token()) } else { 0.0 };
124 if q > 0.0 && best.is_none_or(|(_, best_q)| q > best_q) {
125 best = Some((encoding, q));
126 }
127 }
128 best.map(|(encoding, _)| encoding)
129}
130
131fn window_log(bytes: usize) -> u32 {
133 (usize::BITS - bytes.max(1).saturating_sub(1).leading_zeros()).max(10)
134}
135
136pub fn encode(encoding: DictionaryEncoding, body: &[u8], dictionary: &[u8], level: Level) -> Option<Vec<u8>> {
140 if body.len() > MAX_SIZE || dictionary.len() > MAX_SIZE || !encoding.available() {
141 return None;
142 }
143 let mut out = Vec::with_capacity(encoding.header_len() + body.len() / 4 + 64);
144 out.extend_from_slice(encoding.magic());
145 out.extend_from_slice(&hash(dictionary));
146 let window = window_log(dictionary.len() + body.len() + 16);
149 match encoding {
150 #[cfg(feature = "zstd")]
151 DictionaryEncoding::Dcz => {
152 let mut compressor = zstd::bulk::Compressor::new(super::ZSTD_LEVEL[level.index()]).ok()?;
153 compressor.set_parameter(zstd::zstd_safe::CParameter::WindowLog(window)).ok()?;
154 compressor.context_mut().ref_prefix(dictionary).ok()?;
155 out.extend_from_slice(&compressor.compress(body).ok()?);
156 }
157 #[cfg(feature = "brotli")]
158 DictionaryEncoding::Dcb => {
159 let params = brotli::enc::BrotliEncoderParams {
160 quality: super::BROTLI_QUALITY[level.index()] as i32,
161 lgwin: window.min(24) as i32,
162 ..Default::default()
163 };
164 let mut nop = |_: &mut brotli::interface::PredictionModeContextMap<brotli::InputReferenceMut>,
165 _: &mut [brotli::interface::StaticCommand],
166 _: brotli::InputPair,
167 _: &mut brotli::enc::StandardAlloc| ();
168 brotli::BrotliCompressCustomIoCustomDict(
169 &mut brotli::enc::IoReaderWrapper(&mut &body[..]),
170 &mut brotli::enc::IoWriterWrapper(&mut out),
171 &mut [0u8; 4096],
172 &mut [0u8; 4096],
173 ¶ms,
174 brotli::enc::StandardAlloc::default(),
175 &mut nop,
176 dictionary,
177 std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "unexpected end of input"),
178 )
179 .ok()?;
180 }
181 #[allow(unreachable_patterns)]
183 _ => {
184 let _ = (level, window);
185 }
186 }
187 Some(out)
188}
189
190#[cfg_attr(not(any(feature = "zstd", feature = "brotli")), allow(unused_variables, unused_mut, unreachable_code))]
194pub fn decode(encoding: DictionaryEncoding, data: &[u8], dictionary: &[u8], max_len: usize) -> Option<Vec<u8>> {
195 let header = encoding.header_len();
196 if data.len() < header || !data.starts_with(encoding.magic()) || data[encoding.magic().len()..header] != hash(dictionary) {
197 return None;
198 }
199 let stream = &data[header..];
200 let mut out = Vec::new();
201 match encoding {
202 #[cfg(feature = "zstd")]
203 DictionaryEncoding::Dcz => {
204 use std::io::Read;
205 let mut decoder = zstd::stream::read::Decoder::with_ref_prefix(stream, dictionary).ok()?;
206 decoder.window_log_max(23).ok()?;
207 decoder.take(max_len as u64 + 1).read_to_end(&mut out).ok()?;
208 }
209 #[cfg(feature = "brotli")]
210 DictionaryEncoding::Dcb => {
211 use brotli::{Allocator, SliceWrapperMut};
212 let mut alloc = brotli::enc::StandardAlloc::default();
213 let mut dict = <brotli::enc::StandardAlloc as Allocator<u8>>::alloc_cell(&mut alloc, dictionary.len());
214 dict.slice_mut().copy_from_slice(dictionary);
215 let mut writer = LimitedWriter { out: &mut out, max: max_len + 1 };
216 brotli::BrotliDecompressCustomIoCustomDict(
217 &mut brotli::enc::IoReaderWrapper(&mut &stream[..]),
218 &mut brotli::enc::IoWriterWrapper(&mut writer),
219 &mut [0u8; 4096],
220 &mut [0u8; 4096],
221 brotli::enc::StandardAlloc::default(),
222 brotli::enc::StandardAlloc::default(),
223 brotli::enc::StandardAlloc::default(),
224 dict,
225 std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "unexpected end of input"),
226 )
227 .ok()?;
228 }
229 #[allow(unreachable_patterns)]
230 _ => {
231 let _ = stream;
232 return None;
233 }
234 }
235 (out.len() <= max_len).then_some(out)
236}
237
238#[cfg(feature = "brotli")]
240struct LimitedWriter<'a> {
241 out: &'a mut Vec<u8>,
242 max: usize,
243}
244
245#[cfg(feature = "brotli")]
246impl std::io::Write for LimitedWriter<'_> {
247 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
248 if self.out.len() + buf.len() > self.max {
249 return Err(std::io::Error::other("decoded body too large"));
250 }
251 self.out.extend_from_slice(buf);
252 Ok(buf.len())
253 }
254
255 fn flush(&mut self) -> std::io::Result<()> {
256 Ok(())
257 }
258}
259
260pub fn mark_encoded(response: &mut Response, encoding: DictionaryEncoding) {
263 response.headers.push(Header { name: Header::_CONTENT_ENCODING.into(), value: encoding.token().into() });
264 add_vary(response, "Accept-Encoding");
265 add_vary(response, "Available-Dictionary");
266}
267
268#[doc(hidden)]
270pub fn add_vary(response: &mut Response, name: &'static str) {
271 if let Some(vary) = response.headers.iter_mut().find(|h| h.name.eq_ignore_ascii_case("vary")) {
272 if !vary.value.split(',').any(|v| v.trim().eq_ignore_ascii_case(name) || v.trim() == "*") {
273 vary.value = format!("{}, {name}", vary.value).into();
274 }
275 } else {
276 response.headers.push(Header { name: "Vary".into(), value: name.into() });
277 }
278}
279
280#[cfg(test)]
281mod tests;