Skip to main content

wacore_binary/
util.rs

1use crate::error::{BinaryError, Result};
2use crate::zlib_pool::decompress_zlib_pooled;
3use bytes::{Buf, Bytes, BytesMut};
4use std::borrow::Cow;
5
6/// Protocol frames larger than 16 MiB after decompression are rejected.
7const MAX_DECOMPRESSED_SIZE: u64 = 16 * 1024 * 1024;
8
9fn decompress_zlib(compressed: &[u8]) -> Result<Vec<u8>> {
10    decompress_zlib_pooled(compressed, MAX_DECOMPRESSED_SIZE)
11        .map_err(|e| BinaryError::Zlib(e.to_string()))
12}
13
14pub fn unpack(data: &[u8]) -> Result<Cow<'_, [u8]>> {
15    if data.is_empty() {
16        return Err(BinaryError::EmptyData);
17    }
18    let data_type = data[0];
19    let data = &data[1..];
20
21    if (data_type & 2) > 0 {
22        Ok(Cow::Owned(decompress_zlib(data)?))
23    } else {
24        Ok(Cow::Borrowed(data))
25    }
26}
27
28/// Unpack a network payload into an owned `Bytes` buffer.
29///
30/// Uncompressed payloads reuse the existing `BytesMut` allocation
31/// and freeze it without copying. Compressed payloads allocate a
32/// decompression buffer which is then wrapped as `Bytes`.
33pub fn unpack_bytes(mut data: BytesMut) -> Result<Bytes> {
34    if data.is_empty() {
35        return Err(BinaryError::EmptyData);
36    }
37    let data_type = data[0];
38
39    if (data_type & 2) > 0 {
40        Ok(Bytes::from(decompress_zlib(&data[1..])?))
41    } else {
42        data.advance(1);
43        Ok(data.freeze())
44    }
45}