1use std::sync::Mutex;
20
21use flate2::{Compress, Compression, Decompress, FlushCompress, FlushDecompress, Status};
22
23const TRAILER: [u8; 4] = [0x00, 0x00, 0xFF, 0xFF];
25
26#[derive(Debug, Clone, Copy)]
27pub struct DeflateConfig {
28 pub enabled: bool,
29 pub force_client_no_context_takeover: bool,
32 pub min_compress_size: usize,
34 pub level: u32,
36}
37
38impl Default for DeflateConfig {
39 fn default() -> Self {
40 DeflateConfig {
41 enabled: true,
42 force_client_no_context_takeover: false,
43 min_compress_size: 128,
44 level: 6,
45 }
46 }
47}
48
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
51pub struct DeflateParams {
52 pub client_no_context_takeover: bool,
54 pub server_no_context_takeover: bool,
56}
57
58pub fn negotiate_server<'a, I>(offers: I, cfg: &DeflateConfig) -> Option<(DeflateParams, String)>
61where
62 I: IntoIterator<Item = &'a str>,
63{
64 if !cfg.enabled {
65 return None;
66 }
67 for header in offers {
68 for offer in header.split(',') {
69 let mut parts = offer.split(';').map(str::trim);
70 if !parts
71 .next()
72 .is_some_and(|n| n.eq_ignore_ascii_case("permessage-deflate"))
73 {
74 continue;
75 }
76 let mut seen = Vec::new();
77 let mut client_nct = false;
78 let mut acceptable = true;
79 for p in parts.filter(|p| !p.is_empty()) {
80 let (name, value) = match p.split_once('=') {
81 Some((n, v)) => (
82 n.trim().to_ascii_lowercase(),
83 Some(v.trim().trim_matches('"')),
84 ),
85 None => (p.to_ascii_lowercase(), None),
86 };
87 if seen.contains(&name) {
88 acceptable = false;
89 break;
90 }
91 match (name.as_str(), value) {
92 ("server_no_context_takeover", None) => {}
93 ("client_no_context_takeover", None) => client_nct = true,
94 ("server_max_window_bits", Some(v)) => acceptable &= v == "15",
95 ("client_max_window_bits", None) => {}
96 ("client_max_window_bits", Some(v)) => {
97 acceptable &= v.parse::<u8>().is_ok_and(|b| (8..=15).contains(&b))
98 }
99 _ => acceptable = false,
100 }
101 seen.push(name);
102 }
103 if !acceptable {
104 continue;
105 }
106 let client_nct = client_nct || cfg.force_client_no_context_takeover;
107 let mut response = String::from("permessage-deflate; server_no_context_takeover");
108 if client_nct {
109 response.push_str("; client_no_context_takeover");
110 }
111 return Some((
112 DeflateParams {
113 client_no_context_takeover: client_nct,
114 server_no_context_takeover: true,
115 },
116 response,
117 ));
118 }
119 }
120 None
121}
122
123pub fn accept_client_response(header: &str) -> Option<DeflateParams> {
125 let mut parts = header.split(';').map(str::trim);
126 if !parts.next()?.eq_ignore_ascii_case("permessage-deflate") {
127 return None;
128 }
129 let mut params = DeflateParams {
130 client_no_context_takeover: false,
131 server_no_context_takeover: false,
132 };
133 for p in parts {
134 match p
135 .split('=')
136 .next()
137 .unwrap_or("")
138 .trim()
139 .to_ascii_lowercase()
140 .as_str()
141 {
142 "client_no_context_takeover" => params.client_no_context_takeover = true,
143 "server_no_context_takeover" => params.server_no_context_takeover = true,
144 _ => {}
145 }
146 }
147 Some(params)
148}
149
150static COMPRESSOR_POOL: Mutex<Vec<(u32, Compress)>> = Mutex::new(Vec::new());
151const POOL_CAP: usize = 64;
152
153pub fn compress_message(data: &[u8], level: u32) -> Vec<u8> {
155 let mut c = {
156 let mut pool = COMPRESSOR_POOL.lock().unwrap_or_else(|e| e.into_inner());
157 match pool.iter().position(|(l, _)| *l == level) {
158 Some(i) => pool.swap_remove(i).1,
159 None => Compress::new(Compression::new(level), false),
160 }
161 };
162 let start_in = c.total_in();
163 let mut out = Vec::with_capacity(data.len() / 2 + 64);
164 loop {
165 if out.capacity() - out.len() < 64 {
166 out.reserve(out.capacity().max(256));
167 }
168 let consumed = (c.total_in() - start_in) as usize;
169 c.compress_vec(&data[consumed..], &mut out, FlushCompress::Sync)
170 .expect("deflate never fails on valid state");
171 let all_in = (c.total_in() - start_in) as usize == data.len();
172 if all_in && out.len() < out.capacity() {
174 break;
175 }
176 }
177 if out.ends_with(&TRAILER) {
178 out.truncate(out.len() - TRAILER.len());
179 }
180 c.reset();
181 let mut pool = COMPRESSOR_POOL.lock().unwrap_or_else(|e| e.into_inner());
182 if pool.len() < POOL_CAP {
183 pool.push((level, c));
184 }
185 out
186}
187
188#[derive(Debug, Clone, PartialEq, Eq)]
189pub enum InflateError {
190 TooBig,
191 Corrupt,
192}
193
194pub struct Inflater {
196 state: Option<Box<Decompress>>,
197 no_context_takeover: bool,
198}
199
200impl Inflater {
201 pub fn new(no_context_takeover: bool) -> Self {
202 Inflater {
203 state: None,
204 no_context_takeover,
205 }
206 }
207
208 pub fn decompress(&mut self, data: &[u8], max_size: usize) -> Result<Vec<u8>, InflateError> {
209 let d = self
210 .state
211 .get_or_insert_with(|| Box::new(Decompress::new(false)));
212 let result = inflate(d, data, max_size);
213 match &result {
214 Err(_) => self.state = None,
216 Ok(_) if self.no_context_takeover => d.reset(false),
217 Ok(_) => {}
218 }
219 result
220 }
221}
222
223fn inflate(d: &mut Decompress, data: &[u8], max_size: usize) -> Result<Vec<u8>, InflateError> {
224 let mut input = Vec::with_capacity(data.len() + TRAILER.len());
225 input.extend_from_slice(data);
226 input.extend_from_slice(&TRAILER);
227 let start_in = d.total_in();
228 let consumed = |d: &Decompress| (d.total_in() - start_in) as usize;
229 let mut out = Vec::with_capacity((data.len() * 3).clamp(256, max_size.max(256)));
230 loop {
231 if out.capacity() - out.len() < 256 {
232 if out.len() >= max_size {
233 return Err(InflateError::TooBig);
234 }
235 out.reserve(out.capacity().min(max_size));
236 }
237 let (in_before, out_before) = (consumed(d), out.len());
238 let status = d
239 .decompress_vec(&input[in_before..], &mut out, FlushDecompress::Sync)
240 .map_err(|_| InflateError::Corrupt)?;
241 if out.len() > max_size {
242 return Err(InflateError::TooBig);
243 }
244 let all_in = consumed(d) == input.len();
245 let output_has_room = out.len() < out.capacity();
246 if status == Status::StreamEnd || (all_in && output_has_room) {
247 return Ok(out);
248 }
249 let progressed = consumed(d) != in_before || out.len() != out_before;
250 if !progressed && output_has_room {
251 return Err(InflateError::Corrupt);
253 }
254 }
255}
256
257#[cfg(test)]
258mod tests {
259 use super::*;
260
261 #[test]
262 fn negotiation() {
263 let cfg = DeflateConfig::default();
264 let (p, r) =
265 negotiate_server(["permessage-deflate; client_max_window_bits"], &cfg).unwrap();
266 assert_eq!(r, "permessage-deflate; server_no_context_takeover");
267 assert!(!p.client_no_context_takeover);
268 let (p, r) = negotiate_server(["permessage-deflate; server_max_window_bits=10, permessage-deflate; client_no_context_takeover"], &cfg).unwrap();
269 assert!(p.client_no_context_takeover);
270 assert_eq!(
271 r,
272 "permessage-deflate; server_no_context_takeover; client_no_context_takeover"
273 );
274 assert!(negotiate_server(["x-webkit-deflate-frame"], &cfg).is_none());
275 assert!(negotiate_server(["permessage-deflate; foo"], &cfg).is_none());
276 assert!(
277 negotiate_server(
278 ["permessage-deflate; server_no_context_takeover; server_no_context_takeover"],
279 &cfg
280 )
281 .is_none()
282 );
283 assert!(
284 negotiate_server(
285 ["permessage-deflate"],
286 &DeflateConfig {
287 enabled: false,
288 ..cfg
289 }
290 )
291 .is_none()
292 );
293 let forced = DeflateConfig {
294 force_client_no_context_takeover: true,
295 ..cfg
296 };
297 assert!(
298 negotiate_server(["permessage-deflate"], &forced)
299 .unwrap()
300 .0
301 .client_no_context_takeover
302 );
303 }
304
305 #[test]
306 fn roundtrip_with_and_without_context() {
307 for nct in [false, true] {
308 let mut inf = Inflater::new(nct);
309 for i in 0..20 {
310 let msg = format!(
311 r#"[2,"{i}","MeterValues",{{"connectorId":1,"meterValue":[{{"timestamp":"2024-01-01T00:00:00Z","sampledValue":[{{"value":"{}"}}]}}]}}]"#,
312 i * 1000
313 );
314 let c = compress_message(msg.as_bytes(), 6);
315 assert!(!c.ends_with(&TRAILER));
316 assert_eq!(inf.decompress(&c, 1 << 20).unwrap(), msg.as_bytes());
317 }
318 }
319 }
320
321 #[test]
322 fn large_and_bomb() {
323 let big = "x".repeat(3_000_000);
324 let c = compress_message(big.as_bytes(), 6);
325 assert!(c.len() < 20_000);
326 let mut inf = Inflater::new(true);
327 assert_eq!(inf.decompress(&c, 1 << 20), Err(InflateError::TooBig));
328 assert_eq!(inf.decompress(&c, 4 << 20).unwrap().len(), 3_000_000);
329 assert_eq!(
330 Inflater::new(true).decompress(&[0xFF, 0xFF, 0xFF], 1024),
331 Err(InflateError::Corrupt)
332 );
333 }
334}