1use bytes::{Buf, BufMut, Bytes, BytesMut};
17#[cfg(not(feature = "libdeflate"))]
18use flate2::{Compress, Compression, FlushCompress, Status};
19use flate2::{Decompress, FlushDecompress};
20
21use crate::codec::{McWrite, Reader, read_varint_prefix, varint_len};
22use crate::crypto::{Cfb8Decryptor, Cfb8Encryptor};
23use crate::error::{DecodeError, EncodeError, FrameError};
24use crate::packet::{Packet, PacketFilter, PacketSet, RawPacket};
25
26pub const MAX_FRAME_LEN: usize = (1 << 21) - 1;
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub struct FrameLimits {
32 pub max_frame_len: usize,
34 pub max_uncompressed_len: usize,
37}
38
39impl Default for FrameLimits {
40 fn default() -> Self {
41 FrameLimits {
42 max_frame_len: MAX_FRAME_LEN,
43 max_uncompressed_len: 8 << 20,
44 }
45 }
46}
47
48#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
50pub struct DecoderStats {
51 pub frames: u64,
53 pub skipped: u64,
55 pub skipped_compressed: u64,
57 pub inflated_bytes: u64,
59}
60
61pub struct FrameDecoder {
63 buf: BytesMut,
64 decryptor: Option<Cfb8Decryptor>,
65 compression: bool,
66 inflater: Decompress,
69 #[cfg(feature = "libdeflate")]
70 libdeflate: libdeflater::Decompressor,
71 filter: PacketFilter,
72 limits: FrameLimits,
73 stats: DecoderStats,
74}
75
76impl std::fmt::Debug for FrameDecoder {
77 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78 f.debug_struct("FrameDecoder")
79 .field("buffered", &self.buf.len())
80 .field("encrypted", &self.decryptor.is_some())
81 .field("compression", &self.compression)
82 .field("stats", &self.stats)
83 .finish_non_exhaustive()
84 }
85}
86
87impl Default for FrameDecoder {
88 fn default() -> Self {
89 Self::new()
90 }
91}
92
93impl FrameDecoder {
94 pub fn new() -> Self {
96 FrameDecoder {
97 buf: BytesMut::with_capacity(16 * 1024),
98 decryptor: None,
99 compression: false,
100 inflater: Decompress::new(true),
101 #[cfg(feature = "libdeflate")]
102 libdeflate: libdeflater::Decompressor::new(),
103 filter: PacketFilter::all(),
104 limits: FrameLimits::default(),
105 stats: DecoderStats::default(),
106 }
107 }
108
109 pub fn set_limits(&mut self, limits: FrameLimits) {
111 self.limits = FrameLimits {
112 max_frame_len: limits.max_frame_len.min(MAX_FRAME_LEN),
113 ..limits
114 };
115 }
116
117 pub fn set_compression(&mut self, threshold: i32) {
120 self.compression = threshold >= 0;
121 }
122
123 pub fn compression(&self) -> bool {
125 self.compression
126 }
127
128 pub fn enable_encryption(&mut self, secret: &[u8; 16]) {
132 let mut d = Cfb8Decryptor::new(secret);
133 d.decrypt(&mut self.buf);
134 self.decryptor = Some(d);
135 }
136
137 pub fn is_encrypted(&self) -> bool {
139 self.decryptor.is_some()
140 }
141
142 pub fn set_filter(&mut self, filter: PacketFilter) {
144 self.filter = filter;
145 }
146
147 pub fn filter(&self) -> &PacketFilter {
149 &self.filter
150 }
151
152 pub fn stats(&self) -> DecoderStats {
154 self.stats
155 }
156
157 pub fn buffered(&self) -> usize {
159 self.buf.len()
160 }
161
162 pub fn push(&mut self, data: &[u8]) {
164 let start = self.buf.len();
165 self.buf.extend_from_slice(data);
166 self.decrypt_from(start);
167 }
168
169 pub fn buffer_mut(&mut self) -> &mut BytesMut {
172 &mut self.buf
173 }
174
175 pub fn commit(&mut self, old_len: usize) {
177 self.decrypt_from(old_len);
178 }
179
180 fn decrypt_from(&mut self, start: usize) {
181 if let Some(d) = &mut self.decryptor
182 && start < self.buf.len()
183 {
184 d.decrypt(&mut self.buf[start..]);
185 }
186 }
187
188 pub fn needed(&self) -> Option<usize> {
190 match read_varint_prefix(&self.buf[..self.buf.len().min(3)]) {
191 Ok(Some((len, n))) if len >= 0 => (n + len as usize).checked_sub(self.buf.len()),
192 _ => None,
193 }
194 }
195
196 pub fn next_packet(&mut self) -> Result<Option<RawPacket>, FrameError> {
201 loop {
202 let Some(frame) = self.next_frame()? else {
203 return Ok(None);
204 };
205 self.stats.frames += 1;
206 if let Some(p) = self.unpack(frame)? {
207 return Ok(Some(p));
208 }
209 self.stats.skipped += 1;
210 }
211 }
212
213 fn next_frame(&mut self) -> Result<Option<Bytes>, FrameError> {
215 let head = &self.buf[..self.buf.len().min(3)];
216 let (len, n) = match read_varint_prefix(head) {
217 Ok(Some(v)) => v,
218 Ok(None) if head.len() < 3 => return Ok(None),
219 _ => return Err(FrameError::BadLength),
220 };
221 if len <= 0 {
222 return Err(FrameError::BadLength);
223 }
224 let len = len as usize;
225 if len > self.limits.max_frame_len {
226 return Err(FrameError::FrameTooLarge {
227 len,
228 max: self.limits.max_frame_len,
229 });
230 }
231 if self.buf.len() < n + len {
232 self.buf.reserve(n + len - self.buf.len());
234 return Ok(None);
235 }
236 self.buf.advance(n);
237 Ok(Some(self.buf.split_to(len).freeze()))
238 }
239
240 fn unpack(&mut self, frame: Bytes) -> Result<Option<RawPacket>, FrameError> {
242 if !self.compression {
243 return self.split_id(frame);
244 }
245 let mut r = Reader::new(&frame);
246 let data_len = r.read_varint()?;
247 if data_len < 0 {
248 return Err(FrameError::Decode(DecodeError::NegativeLength(i64::from(
249 data_len,
250 ))));
251 }
252 if data_len == 0 {
253 let rest = frame.slice(r.position()..);
254 return self.split_id(rest);
255 }
256 let data_len = data_len as usize;
257 if data_len > self.limits.max_uncompressed_len {
258 return Err(FrameError::DecompressedTooLarge {
259 len: data_len,
260 max: self.limits.max_uncompressed_len,
261 });
262 }
263 let input = &frame[r.position()..];
264 if !self.filter.is_all() {
265 self.inflater.reset(true);
267 let mut head = [0u8; 5];
268 let want = data_len.min(5);
269 self.inflater
270 .decompress(input, &mut head[..want], FlushDecompress::None)
271 .map_err(|_| FrameError::BadCompression("corrupt zlib stream"))?;
272 let got = self.inflater.total_out() as usize;
273 let id = match read_varint_prefix(&head[..got])? {
274 Some((id, _)) => id,
275 None if got < want => return Err(FrameError::BadCompression("stream ended early")),
276 None => return Err(FrameError::Decode(DecodeError::VarIntTooLong)),
277 };
278 self.stats.inflated_bytes += got as u64;
279 if !self.filter.contains(id) {
280 self.stats.skipped_compressed += 1;
281 return Ok(None);
282 }
283 }
284 let out = self.inflate(input, data_len)?;
285 self.stats.inflated_bytes += data_len as u64;
286 let data = Bytes::from(out);
287 let mut r = Reader::new(&data);
288 let id = r.read_varint()?;
289 let body = data.slice(r.position()..);
290 Ok(Some(RawPacket::new(id, body)))
291 }
292
293 #[cfg(feature = "libdeflate")]
295 fn inflate(&mut self, input: &[u8], data_len: usize) -> Result<Vec<u8>, FrameError> {
296 let mut out = vec![0u8; data_len];
297 match self.libdeflate.zlib_decompress(input, &mut out) {
298 Ok(n) if n == data_len => Ok(out),
299 Ok(_) | Err(libdeflater::DecompressionError::InsufficientSpace) => Err(
300 FrameError::BadCompression("size differs from declared length"),
301 ),
302 Err(_) => Err(FrameError::BadCompression("corrupt zlib stream")),
303 }
304 }
305
306 #[cfg(not(feature = "libdeflate"))]
308 fn inflate(&mut self, input: &[u8], data_len: usize) -> Result<Vec<u8>, FrameError> {
309 self.inflater.reset(true);
310 let mut out: Vec<u8> = Vec::with_capacity(data_len);
311 let status = self
312 .inflater
313 .decompress_vec(input, &mut out, FlushDecompress::Finish)
314 .map_err(|_| FrameError::BadCompression("corrupt zlib stream"))?;
315 if out.len() != data_len {
316 return Err(FrameError::BadCompression(
317 "size differs from declared length",
318 ));
319 }
320 if status != Status::StreamEnd {
321 let consumed = self.inflater.total_in() as usize;
323 let mut extra = [0u8; 1];
324 let st = self
325 .inflater
326 .decompress(
327 &input[consumed.min(input.len())..],
328 &mut extra,
329 FlushDecompress::Finish,
330 )
331 .map_err(|_| FrameError::BadCompression("corrupt zlib stream"))?;
332 if self.inflater.total_out() as usize != data_len || st != Status::StreamEnd {
333 return Err(FrameError::BadCompression(
334 "size differs from declared length",
335 ));
336 }
337 }
338 Ok(out)
339 }
340
341 fn split_id(&self, data: Bytes) -> Result<Option<RawPacket>, FrameError> {
342 let mut r = Reader::new(&data);
343 let id = r.read_varint()?;
344 if !self.filter.contains(id) {
345 return Ok(None);
346 }
347 let body = data.slice(r.position()..);
348 Ok(Some(RawPacket::new(id, body)))
349 }
350}
351
352pub struct FrameEncoder {
356 out: BytesMut,
357 scratch: BytesMut,
358 compressed: Vec<u8>,
359 threshold: Option<usize>,
360 #[cfg(not(feature = "libdeflate"))]
361 deflater: Compress,
362 #[cfg(feature = "libdeflate")]
363 deflater: libdeflater::Compressor,
364 encryptor: Option<Cfb8Encryptor>,
365}
366
367impl std::fmt::Debug for FrameEncoder {
368 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
369 f.debug_struct("FrameEncoder")
370 .field("pending", &self.out.len())
371 .field("threshold", &self.threshold)
372 .field("encrypted", &self.encryptor.is_some())
373 .finish_non_exhaustive()
374 }
375}
376
377impl Default for FrameEncoder {
378 fn default() -> Self {
379 Self::new()
380 }
381}
382
383impl FrameEncoder {
384 pub fn new() -> Self {
386 FrameEncoder {
387 out: BytesMut::with_capacity(4096),
388 scratch: BytesMut::with_capacity(1024),
389 compressed: Vec::new(),
390 threshold: None,
391 #[cfg(not(feature = "libdeflate"))]
393 deflater: Compress::new(Compression::default(), true),
394 #[cfg(feature = "libdeflate")]
395 deflater: libdeflater::Compressor::new(libdeflater::CompressionLvl::default()),
396 encryptor: None,
397 }
398 }
399
400 pub fn set_compression(&mut self, threshold: i32) {
402 self.threshold = usize::try_from(threshold).ok();
403 }
404
405 pub fn enable_encryption(&mut self, secret: &[u8; 16]) {
407 self.encryptor = Some(Cfb8Encryptor::new(secret));
408 }
409
410 pub fn write_packet<P: Packet>(&mut self, p: &P) -> Result<(), EncodeError> {
412 self.scratch.clear();
413 self.scratch.put_varint(P::ID);
414 p.encode(&mut self.scratch)?;
415 self.frame_scratch()
416 }
417
418 pub fn write_set<S: PacketSet>(&mut self, p: &S) -> Result<(), EncodeError> {
420 self.scratch.clear();
421 self.scratch.put_varint(p.id());
422 p.encode_body(&mut self.scratch)?;
423 self.frame_scratch()
424 }
425
426 pub fn write_raw(&mut self, p: &RawPacket) -> Result<(), EncodeError> {
428 self.scratch.clear();
429 self.scratch.put_varint(p.id);
430 self.scratch.put_slice(&p.body);
431 self.frame_scratch()
432 }
433
434 fn frame_scratch(&mut self) -> Result<(), EncodeError> {
435 let start = self.out.len();
436 let len = self.scratch.len();
437 let too_big = |len| EncodeError::LengthOverflow {
438 len,
439 max: MAX_FRAME_LEN,
440 };
441 match self.threshold {
442 None => {
443 if len > MAX_FRAME_LEN {
444 return Err(too_big(len));
445 }
446 self.out.put_varint(len as i32);
447 self.out.put_slice(&self.scratch);
448 }
449 Some(t) if len < t => {
450 if len + 1 > MAX_FRAME_LEN {
451 return Err(too_big(len + 1));
452 }
453 self.out.put_varint(len as i32 + 1);
454 self.out.put_u8(0);
455 self.out.put_slice(&self.scratch);
456 }
457 Some(_) => {
458 let len_i = i32::try_from(len).map_err(|_| too_big(len))?;
459 self.deflate()?;
460 let frame_len = varint_len(len_i) + self.compressed.len();
461 if frame_len > MAX_FRAME_LEN {
462 return Err(too_big(frame_len));
463 }
464 self.out.put_varint(frame_len as i32);
465 self.out.put_varint(len_i);
466 self.out.put_slice(&self.compressed);
467 }
468 }
469 if let Some(e) = &mut self.encryptor {
470 e.encrypt(&mut self.out[start..]);
471 }
472 Ok(())
473 }
474
475 #[cfg(feature = "libdeflate")]
477 fn deflate(&mut self) -> Result<(), EncodeError> {
478 let bound = self.deflater.zlib_compress_bound(self.scratch.len());
479 self.compressed.clear();
480 self.compressed.resize(bound, 0);
481 let n = self
482 .deflater
483 .zlib_compress(&self.scratch, &mut self.compressed)
484 .map_err(|_| EncodeError::Invalid("deflate failed"))?;
485 self.compressed.truncate(n);
486 Ok(())
487 }
488
489 #[cfg(not(feature = "libdeflate"))]
491 fn deflate(&mut self) -> Result<(), EncodeError> {
492 self.deflater.reset();
493 self.compressed.clear();
494 self.compressed
495 .reserve(self.scratch.len() + self.scratch.len() / 1000 + 64);
496 loop {
497 let consumed = self.deflater.total_in() as usize;
498 let st = self
499 .deflater
500 .compress_vec(
501 &self.scratch[consumed..],
502 &mut self.compressed,
503 FlushCompress::Finish,
504 )
505 .map_err(|_| EncodeError::Invalid("deflate failed"))?;
506 if st == Status::StreamEnd {
507 return Ok(());
508 }
509 self.compressed.reserve(self.compressed.capacity().max(64));
510 }
511 }
512
513 pub fn pending(&self) -> &[u8] {
515 &self.out
516 }
517
518 pub fn take(&mut self) -> BytesMut {
520 self.out.split()
521 }
522
523 pub fn out_mut(&mut self) -> &mut BytesMut {
525 &mut self.out
526 }
527}
528
529#[cfg(test)]
530mod tests {
531 use super::*;
532 use flate2::{Compress, Compression, FlushCompress};
533
534 fn roundtrip(threshold: Option<i32>, secret: Option<[u8; 16]>, filter: PacketFilter) {
535 let mut enc = FrameEncoder::new();
536 let mut dec = FrameDecoder::new();
537 if let Some(t) = threshold {
538 enc.set_compression(t);
539 dec.set_compression(t);
540 }
541 if let Some(s) = secret {
542 enc.enable_encryption(&s);
543 dec.enable_encryption(&s);
544 }
545 dec.set_filter(filter);
546 let packets: Vec<RawPacket> = (0..50)
547 .map(|i| {
548 let body: Vec<u8> = (0..(i * 37) % 700).map(|j| (j * 7 + i) as u8).collect();
549 RawPacket::new(i % 9, Bytes::from(body))
550 })
551 .collect();
552 for p in &packets {
553 enc.write_raw(p).unwrap();
554 }
555 let wire = enc.take();
556 let mut got = Vec::new();
558 for chunk in wire.chunks(13) {
559 dec.push(chunk);
560 while let Some(p) = dec.next_packet().unwrap() {
561 got.push(p);
562 }
563 }
564 let want: Vec<RawPacket> = packets
565 .into_iter()
566 .filter(|p| filter.contains(p.id))
567 .collect();
568 assert_eq!(got, want);
569 assert_eq!(dec.buffered(), 0);
570 }
571
572 #[test]
573 fn roundtrips() {
574 let key = [3u8; 16];
575 let f = PacketFilter::none().with(1).with(4);
576 for (t, s, filt) in [
577 (None, None, PacketFilter::all()),
578 (Some(256), None, PacketFilter::all()),
579 (Some(0), None, PacketFilter::all()),
580 (Some(64), Some(key), PacketFilter::all()),
581 (None, Some(key), f),
582 (Some(256), None, f),
583 (Some(1), Some(key), f),
584 ] {
585 roundtrip(t, s, filt);
586 }
587 }
588
589 #[test]
590 fn known_answer_frames() {
591 let mut enc = FrameEncoder::new();
593 enc.write_raw(&RawPacket::new(0, Bytes::from_static(&[1, 2])))
594 .unwrap();
595 assert_eq!(&enc.take()[..], &[0x03, 0x00, 0x01, 0x02]);
596 enc.set_compression(256);
598 enc.write_raw(&RawPacket::new(0x21, Bytes::from_static(&[9])))
599 .unwrap();
600 assert_eq!(&enc.take()[..], &[0x03, 0x00, 0x21, 0x09]);
601 let mut body = vec![0x26];
604 body.extend_from_slice(&[0u8; 300]);
605 let mut z = Compress::new(Compression::default(), true);
606 let mut c = Vec::with_capacity(400);
607 z.compress_vec(&body, &mut c, FlushCompress::Finish)
608 .unwrap();
609 assert_eq!(&c[..2], &[0x78, 0x9c], "zlib header for the default level");
610 let mut frame = BytesMut::new();
611 frame.put_varint((varint_len(301) + c.len()) as i32);
612 frame.put_varint(301);
613 frame.put_slice(&c);
614 let mut dec = FrameDecoder::new();
615 dec.set_compression(256);
616 dec.push(&frame);
617 let p = dec.next_packet().unwrap().unwrap();
618 assert_eq!(p.id, 0x26);
619 assert_eq!(&p.body[..], &[0u8; 300][..]);
620 }
621
622 #[test]
623 fn rejects_malformed() {
624 let mut dec = FrameDecoder::new();
625 dec.push(&[0xff, 0xff, 0xff, 0x01]);
626 assert!(matches!(dec.next_packet(), Err(FrameError::BadLength)));
627 let mut dec = FrameDecoder::new();
628 dec.push(&[0x00]);
629 assert!(matches!(dec.next_packet(), Err(FrameError::BadLength)));
630 let mut dec = FrameDecoder::new();
632 dec.set_compression(0);
633 dec.push(&[0x05, 0x80, 0x80, 0x80, 0x10, 0x00]);
634 assert!(matches!(
635 dec.next_packet(),
636 Err(FrameError::DecompressedTooLarge { .. })
637 ));
638 let mut dec = FrameDecoder::new();
640 dec.set_compression(0);
641 dec.push(&[0x04, 0x05, 0x01, 0x02, 0x03]);
642 assert!(matches!(
643 dec.next_packet(),
644 Err(FrameError::BadCompression(_))
645 ));
646 let mut z = Compress::new(Compression::default(), true);
648 let mut c = Vec::with_capacity(64);
649 z.compress_vec(&[1, 2, 3], &mut c, FlushCompress::Finish)
650 .unwrap();
651 let mut frame = BytesMut::new();
652 frame.put_varint(1 + c.len() as i32);
653 frame.put_varint(4);
654 frame.put_slice(&c);
655 let mut dec = FrameDecoder::new();
656 dec.set_compression(0);
657 dec.push(&frame);
658 assert!(matches!(
659 dec.next_packet(),
660 Err(FrameError::BadCompression(_))
661 ));
662 let mut frame = BytesMut::new();
664 frame.put_varint(1 + c.len() as i32);
665 frame.put_varint(2);
666 frame.put_slice(&c);
667 let mut dec = FrameDecoder::new();
668 dec.set_compression(0);
669 dec.push(&frame);
670 assert!(matches!(
671 dec.next_packet(),
672 Err(FrameError::BadCompression(_))
673 ));
674 }
675
676 #[test]
677 fn partial_inflation_skips_unwanted() {
678 let mut enc = FrameEncoder::new();
679 enc.set_compression(16);
680 let big = RawPacket::new(0x38, Bytes::from(vec![7u8; 100_000]));
681 enc.write_raw(&big).unwrap();
682 enc.write_raw(&RawPacket::new(0x00, Bytes::from(vec![1u8; 64])))
683 .unwrap();
684 let mut dec = FrameDecoder::new();
685 dec.set_compression(16);
686 dec.set_filter(PacketFilter::all().without(0x38));
687 dec.push(&enc.take());
688 let p = dec.next_packet().unwrap().unwrap();
689 assert_eq!(p.id, 0);
690 let s = dec.stats();
691 assert_eq!((s.frames, s.skipped, s.skipped_compressed), (2, 1, 1));
692 assert!(
693 s.inflated_bytes < 100,
694 "inflated {} bytes",
695 s.inflated_bytes
696 );
697 }
698
699 #[test]
700 fn encryption_enabled_mid_buffer() {
701 let key = [9u8; 16];
704 let mut plain = FrameEncoder::new();
705 plain
706 .write_raw(&RawPacket::new(1, Bytes::from_static(b"plain")))
707 .unwrap();
708 let mut crypt = FrameEncoder::new();
709 crypt.enable_encryption(&key);
710 crypt
711 .write_raw(&RawPacket::new(2, Bytes::from_static(b"secret")))
712 .unwrap();
713 let mut dec = FrameDecoder::new();
714 dec.push(&plain.take());
715 dec.push(&crypt.take());
716 assert_eq!(dec.next_packet().unwrap().unwrap().id, 1);
717 dec.enable_encryption(&key);
718 let p = dec.next_packet().unwrap().unwrap();
719 assert_eq!((p.id, &p.body[..]), (2, &b"secret"[..]));
720 }
721}