1use std::io::{Read, Write};
8
9use flate2::{Compression, bufread::MultiGzDecoder, write::GzEncoder};
10use ruzstd::{
11 decoding::StreamingDecoder,
12 encoding::{CompressionLevel, compress_to_vec},
13};
14
15use crate::{
16 Error, Limits, Method,
17 message::{self, Request, Response},
18};
19
20pub(crate) const MARKER: u8 = 0x04;
21const MIN_BODY_LEN: usize = 64;
22const MIN_ZSTD_WINDOW: usize = 1024 * 1024;
26const MAX_ZSTD_WINDOW: usize = 8 * 1024 * 1024;
27const MAX_ZSTD_FRAMES: usize = 16;
28
29#[derive(Clone, Copy, Debug, Eq, PartialEq)]
31pub enum CompressionCoding {
32 Gzip,
34 Zstd,
36}
37
38impl CompressionCoding {
39 pub(crate) const fn wire_id(self) -> u8 {
40 match self {
41 Self::Gzip => 1,
42 Self::Zstd => 2,
43 }
44 }
45
46 pub(crate) const fn from_wire_id(id: u8) -> Result<Option<Self>, Error> {
47 match id {
48 0 => Ok(None),
49 1 => Ok(Some(Self::Gzip)),
50 2 => Ok(Some(Self::Zstd)),
51 _ => Err(Error::MalformedEnvelope),
52 }
53 }
54}
55
56pub(crate) fn encode_request(
57 request: &Request,
58 limits: Limits,
59 preferred: Option<CompressionCoding>,
60) -> Result<Vec<u8>, Error> {
61 let Some(preferred) = preferred else {
62 return message::encode_request(request, limits);
63 };
64 let coded = maybe_compress(&request.body, preferred)?;
65 let body_coding = coded.as_ref().map(|_| preferred);
66 let mut output = vec![
67 MARKER,
68 body_coding.map_or(0, CompressionCoding::wire_id),
69 preferred.wire_id(),
70 ];
71 let bhttp = message::encode_request_with_body(
72 request,
73 coded.as_deref().unwrap_or(&request.body),
74 limits,
75 )?;
76 output.extend_from_slice(&bhttp);
77 Ok(output)
78}
79
80pub(crate) fn decode_request(
81 input: &[u8],
82 limits: Limits,
83 enabled: bool,
84) -> Result<
85 (
86 Request,
87 Option<CompressionCoding>,
88 Option<CompressionCoding>,
89 ),
90 Error,
91> {
92 if !input.starts_with(&[MARKER]) {
93 return Ok((message::decode_request(input, limits)?, None, None));
94 }
95 if !enabled || input.len() < 4 {
96 return Err(Error::MalformedEnvelope);
97 }
98 let body_coding = CompressionCoding::from_wire_id(input[1])?;
99 let response_coding = CompressionCoding::from_wire_id(input[2])?;
100 if body_coding.is_none() && response_coding.is_none() {
101 return Err(Error::MalformedEnvelope);
102 }
103 let request = if body_coding.is_some() {
104 message::decode_coded_request(&input[3..], limits)?
105 } else {
106 message::decode_request(&input[3..], limits)?
107 };
108 Ok((request, body_coding, response_coding))
109}
110
111pub(crate) fn finish_request(
112 mut request: Request,
113 coding: Option<CompressionCoding>,
114 limits: Limits,
115) -> Result<Request, Error> {
116 if let Some(coding) = coding {
117 request.body = decompress(&request.body, coding, limits.max_body_len)?;
118 message::validate_decoded_request_body(&request)?;
119 }
120 Ok(request)
121}
122
123pub(crate) fn encode_response(
124 response: &Response,
125 method: Method,
126 limits: Limits,
127 negotiated: Option<CompressionCoding>,
128) -> Result<Vec<u8>, Error> {
129 if let Some(coding) = negotiated {
130 if let Some(coded) = maybe_compress(&response.body, coding)? {
131 let mut output = vec![MARKER, coding.wire_id()];
132 output.extend_from_slice(&message::encode_response_with_body(
133 response, &coded, method, limits,
134 )?);
135 return Ok(output);
136 }
137 }
138 message::encode_response(response, method, limits)
139}
140
141pub(crate) fn decode_response(
142 input: &[u8],
143 method: Method,
144 limits: Limits,
145 negotiated: Option<CompressionCoding>,
146) -> Result<Response, Error> {
147 if !input.starts_with(&[MARKER]) {
148 return message::decode_response(input, method, limits);
149 }
150 if input.len() < 3 {
151 return Err(Error::MalformedEnvelope);
152 }
153 let coding = CompressionCoding::from_wire_id(input[1])?.ok_or(Error::MalformedEnvelope)?;
154 if Some(coding) != negotiated {
155 return Err(Error::MalformedEnvelope);
156 }
157 let mut response = message::decode_coded_response(&input[2..], method, limits)?;
158 response.body = decompress(&response.body, coding, limits.max_body_len)?;
159 message::validate_decoded_response_body(&response, method)?;
160 Ok(response)
161}
162
163pub(crate) fn maybe_compress(
164 body: &[u8],
165 coding: CompressionCoding,
166) -> Result<Option<Vec<u8>>, Error> {
167 if body.len() < MIN_BODY_LEN {
168 return Ok(None);
169 }
170 let coded = match coding {
171 CompressionCoding::Gzip => {
172 let mut encoder = GzEncoder::new(Vec::new(), Compression::fast());
173 encoder
174 .write_all(body)
175 .map_err(|_| Error::CompressionFailure)?;
176 encoder.finish().map_err(|_| Error::CompressionFailure)?
177 }
178 CompressionCoding::Zstd => compress_to_vec(body, CompressionLevel::Fastest),
179 };
180 Ok((coded.len() < body.len()).then_some(coded))
181}
182
183pub(crate) fn decompress(
184 coded: &[u8],
185 coding: CompressionCoding,
186 maximum_body_len: usize,
187) -> Result<Vec<u8>, Error> {
188 if coded.is_empty() {
189 return Err(Error::MalformedEnvelope);
190 }
191 match coding {
192 CompressionCoding::Gzip => {
193 let mut decoder = MultiGzDecoder::new(coded);
194 let body = read_bounded(&mut decoder, maximum_body_len)?;
195 if !decoder.into_inner().is_empty() {
196 return Err(Error::MalformedEnvelope);
197 }
198 Ok(body)
199 }
200 CompressionCoding::Zstd => {
201 let mut remaining = coded;
202 let mut body = Vec::new();
203 let mut decoded_frames = 0_usize;
204 let mut seen_frames = 0_usize;
205 while !remaining.is_empty() {
206 seen_frames += 1;
207 if seen_frames > MAX_ZSTD_FRAMES {
208 return Err(Error::LimitExceeded);
209 }
210 if let Some(skip_len) = skippable_frame_len(remaining)? {
212 remaining = &remaining[skip_len..];
213 continue;
214 }
215 let mut decoder = StreamingDecoder::new_with_max_window_size(
216 remaining,
217 maximum_body_len.clamp(MIN_ZSTD_WINDOW, MAX_ZSTD_WINDOW) as u64,
218 )
219 .map_err(|_| Error::MalformedEnvelope)?;
220 let declared_size = decoder.decoder.content_size();
221 if declared_size > maximum_body_len.saturating_sub(body.len()) as u64 {
222 return Err(Error::LimitExceeded);
223 }
224 let decoded_before = body.len();
225 read_bounded_into(&mut decoder, &mut body, maximum_body_len)?;
226 if declared_size != 0 && declared_size != (body.len() - decoded_before) as u64 {
227 return Err(Error::MalformedEnvelope);
228 }
229 if decoder.decoder.get_checksum_from_data().is_some()
230 && decoder.decoder.get_checksum_from_data()
231 != decoder.decoder.get_calculated_checksum()
232 {
233 return Err(Error::MalformedEnvelope);
234 }
235 remaining = decoder.into_inner();
236 decoded_frames += 1;
237 }
238 if decoded_frames == 0 {
239 return Err(Error::MalformedEnvelope);
240 }
241 Ok(body)
242 }
243 }
244}
245
246fn skippable_frame_len(input: &[u8]) -> Result<Option<usize>, Error> {
247 if input.len() < 4 {
248 return Err(Error::MalformedEnvelope);
249 }
250 let magic = u32::from_le_bytes([input[0], input[1], input[2], input[3]]);
251 if !(0x184d_2a50..=0x184d_2a5f).contains(&magic) {
252 return Ok(None);
253 }
254 if input.len() < 8 {
255 return Err(Error::MalformedEnvelope);
256 }
257 let length = u32::from_le_bytes([input[4], input[5], input[6], input[7]]) as usize;
258 let total = length.checked_add(8).ok_or(Error::MalformedEnvelope)?;
259 if total > input.len() {
260 return Err(Error::MalformedEnvelope);
261 }
262 Ok(Some(total))
263}
264
265fn read_bounded(reader: &mut impl Read, maximum_body_len: usize) -> Result<Vec<u8>, Error> {
266 let mut output = Vec::new();
267 read_bounded_into(reader, &mut output, maximum_body_len)?;
268 Ok(output)
269}
270
271fn read_bounded_into(
272 reader: &mut impl Read,
273 output: &mut Vec<u8>,
274 maximum_body_len: usize,
275) -> Result<(), Error> {
276 let mut buffer = [0_u8; 8192];
277 loop {
278 let count = reader
279 .read(&mut buffer)
280 .map_err(|_| Error::MalformedEnvelope)?;
281 if count == 0 {
282 return Ok(());
283 }
284 if count > maximum_body_len.saturating_sub(output.len()) {
285 return Err(Error::LimitExceeded);
286 }
287 output.extend_from_slice(&buffer[..count]);
288 }
289}
290
291#[cfg(test)]
292mod tests {
293 use super::{
294 CompressionCoding, decode_response, decompress, encode_request, encode_response,
295 maybe_compress,
296 };
297 use crate::{Error, Limits, Method, Request, Response, message};
298
299 const GZIP: &[u8] = &[
302 0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xcb, 0x48, 0xcd, 0xc9, 0xc9,
303 0x57, 0x48, 0x2b, 0xca, 0xcf, 0x55, 0x48, 0xcc, 0x53, 0xc8, 0xcc, 0x4b, 0x49, 0x2d, 0x48,
304 0x05, 0x12, 0x79, 0x25, 0x0a, 0xe9, 0x55, 0x99, 0x05, 0x0a, 0xa9, 0x79, 0xc9, 0xf9, 0x29,
305 0xa9, 0x45, 0x0a, 0x19, 0x44, 0x29, 0x03, 0x00, 0x61, 0xc4, 0x34, 0x58, 0x4d, 0x00, 0x00,
306 0x00,
307 ];
308 const ZSTD: &[u8] = &[
309 0x28, 0xb5, 0x2f, 0xfd, 0x04, 0x48, 0x7d, 0x01, 0x00, 0x74, 0x02, 0x68, 0x65, 0x6c, 0x6c,
310 0x6f, 0x20, 0x66, 0x72, 0x6f, 0x6d, 0x20, 0x61, 0x6e, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x70,
311 0x65, 0x6e, 0x64, 0x65, 0x6e, 0x74, 0x20, 0x7a, 0x73, 0x74, 0x64, 0x20, 0x65, 0x6e, 0x63,
312 0x6f, 0x64, 0x65, 0x72, 0x20, 0x01, 0x00, 0xaf, 0x1e, 0xf5, 0x04, 0x2a, 0x5c, 0xa5, 0xa6,
313 ];
314
315 #[test]
316 fn independent_gzip_members_and_strict_trailer() {
317 let expected =
318 b"hello from an independent gzip encoder hello from an independent gzip encoder";
319 assert_eq!(
320 decompress(GZIP, CompressionCoding::Gzip, 1024),
321 Ok(expected.to_vec())
322 );
323 let mut joined = GZIP.to_vec();
324 joined.extend_from_slice(GZIP);
325 assert_eq!(
326 decompress(&joined, CompressionCoding::Gzip, 1024),
327 Ok(expected.repeat(2))
328 );
329 assert_eq!(
330 decompress(&GZIP[..GZIP.len() - 1], CompressionCoding::Gzip, 1024),
331 Err(Error::MalformedEnvelope)
332 );
333 let mut bad_checksum = GZIP.to_vec();
334 bad_checksum[GZIP.len() - 8] ^= 1;
335 assert_eq!(
336 decompress(&bad_checksum, CompressionCoding::Gzip, 1024),
337 Err(Error::MalformedEnvelope)
338 );
339 let mut garbage = GZIP.to_vec();
340 garbage.push(0xff);
341 assert_eq!(
342 decompress(&garbage, CompressionCoding::Gzip, 1024),
343 Err(Error::MalformedEnvelope)
344 );
345 }
346
347 #[test]
348 fn independent_zstd_frames_and_skippable_frame() {
349 let expected =
350 b"hello from an independent zstd encoder hello from an independent zstd encoder";
351 assert_eq!(
352 decompress(ZSTD, CompressionCoding::Zstd, 1024),
353 Ok(expected.to_vec())
354 );
355 let mut joined = ZSTD.to_vec();
356 joined.extend_from_slice(ZSTD);
357 assert_eq!(
358 decompress(&joined, CompressionCoding::Zstd, 1024),
359 Ok(expected.repeat(2))
360 );
361 let mut skipped = vec![0x50, 0x2a, 0x4d, 0x18, 3, 0, 0, 0, b'x', b'y', b'z'];
362 skipped.extend_from_slice(ZSTD);
363 assert_eq!(
364 decompress(&skipped, CompressionCoding::Zstd, 1024),
365 Ok(expected.to_vec())
366 );
367 assert_eq!(
368 decompress(&ZSTD[..ZSTD.len() - 1], CompressionCoding::Zstd, 1024),
369 Err(Error::MalformedEnvelope)
370 );
371 let mut bad_checksum = ZSTD.to_vec();
372 bad_checksum[ZSTD.len() - 1] ^= 1;
373 assert_eq!(
374 decompress(&bad_checksum, CompressionCoding::Zstd, 1024),
375 Err(Error::MalformedEnvelope)
376 );
377 let mut garbage = ZSTD.to_vec();
378 garbage.push(0xff);
379 assert_eq!(
380 decompress(&garbage, CompressionCoding::Zstd, 1024),
381 Err(Error::MalformedEnvelope)
382 );
383 assert_eq!(
384 decompress(&ZSTD.repeat(17), CompressionCoding::Zstd, 4096),
385 Err(Error::LimitExceeded)
386 );
387 let mut oversized_window = ZSTD.to_vec();
388 oversized_window[5] = 0x90;
389 assert_eq!(
390 decompress(&oversized_window, CompressionCoding::Zstd, 64 * 1024 * 1024),
391 Err(Error::MalformedEnvelope)
392 );
393 let mut sized = ZSTD.to_vec();
396 sized[4] = 0x24;
397 sized[5] = 77;
398 assert_eq!(
399 decompress(&sized, CompressionCoding::Zstd, 1024),
400 Ok(expected.to_vec())
401 );
402 sized[5] = 76;
403 assert_eq!(
404 decompress(&sized, CompressionCoding::Zstd, 1024),
405 Err(Error::MalformedEnvelope)
406 );
407 sized[5] = 200;
408 assert_eq!(
409 decompress(&sized, CompressionCoding::Zstd, 100),
410 Err(Error::LimitExceeded)
411 );
412 }
413
414 #[test]
415 fn decoded_length_is_bounded_during_expansion() -> Result<(), Error> {
416 for coding in [CompressionCoding::Gzip, CompressionCoding::Zstd] {
417 let encoded =
418 maybe_compress(&vec![b'a'; 4096], coding)?.ok_or(Error::CompressionFailure)?;
419 assert_eq!(decompress(&encoded, coding, 128), Err(Error::LimitExceeded));
420 }
421 Ok(())
422 }
423
424 #[test]
425 fn body_coding_extension_has_literal_wire_prefixes() -> Result<(), Error> {
426 let limits = Limits::default();
427 let request = Request {
428 method: Method::Post,
429 authority: b"api.example.test".to_vec(),
430 path: b"/compressed".to_vec(),
431 headers: Vec::new(),
432 body: vec![b'a'; 4096],
433 };
434 let response = Response {
435 status: 200,
436 headers: Vec::new(),
437 body: vec![b'b'; 4096],
438 };
439 for (coding, request_prefix, response_prefix) in [
440 (CompressionCoding::Gzip, [0x04, 0x01, 0x01], [0x04, 0x01]),
441 (CompressionCoding::Zstd, [0x04, 0x02, 0x02], [0x04, 0x02]),
442 ] {
443 assert!(encode_request(&request, limits, Some(coding))?.starts_with(&request_prefix));
444 assert!(
445 encode_response(&response, Method::Post, limits, Some(coding))?
446 .starts_with(&response_prefix)
447 );
448
449 let mut small_request = request.clone();
450 small_request.body.clear();
451 assert!(
452 encode_request(&small_request, limits, Some(coding))?.starts_with(&[
453 0x04,
454 0x00,
455 request_prefix[2]
456 ])
457 );
458 }
459 Ok(())
460 }
461
462 #[test]
463 fn compressed_response_requires_the_advertised_coding() -> Result<(), Error> {
464 let response = Response {
465 status: 200,
466 headers: Vec::new(),
467 body: vec![b'x'; 4096],
468 };
469 let coded = maybe_compress(&response.body, CompressionCoding::Gzip)?
470 .ok_or(Error::CompressionFailure)?;
471 let mut wire = vec![0x04, 0x01];
472 wire.extend_from_slice(&message::encode_response_with_body(
473 &response,
474 &coded,
475 Method::Get,
476 Limits::default(),
477 )?);
478 assert_eq!(
479 decode_response(&wire, Method::Get, Limits::default(), None),
480 Err(Error::MalformedEnvelope)
481 );
482 assert_eq!(
483 decode_response(
484 &wire,
485 Method::Get,
486 Limits::default(),
487 Some(CompressionCoding::Zstd),
488 ),
489 Err(Error::MalformedEnvelope)
490 );
491 assert_eq!(
492 decode_response(
493 &wire,
494 Method::Get,
495 Limits::default(),
496 Some(CompressionCoding::Gzip),
497 ),
498 Ok(response)
499 );
500 Ok(())
501 }
502}