Skip to main content

zrip_core/frame/
header.rs

1#![forbid(unsafe_code)]
2
3use crate::error::DecompressError;
4use crate::frame::ZSTD_MAGIC;
5
6#[derive(Debug, Clone)]
7pub struct FrameHeader {
8    pub window_size: u64,
9    pub frame_content_size: Option<u64>,
10    pub dict_id: Option<u32>,
11    pub content_checksum: bool,
12    pub single_segment: bool,
13    pub header_size: usize,
14}
15
16pub fn parse_frame_header(data: &[u8]) -> Result<FrameHeader, DecompressError> {
17    if data.len() < 4 {
18        return Err(DecompressError::InputExhausted);
19    }
20
21    let magic = u32::from_le_bytes([data[0], data[1], data[2], data[3]]);
22    if magic != ZSTD_MAGIC {
23        return Err(DecompressError::BadMagic);
24    }
25
26    parse_frame_header_after_magic(&data[4..], 4)
27}
28
29pub fn parse_frame_header_after_magic(
30    data: &[u8],
31    magic_size: usize,
32) -> Result<FrameHeader, DecompressError> {
33    if data.is_empty() {
34        return Err(DecompressError::BadFrameHeader);
35    }
36
37    let descriptor = data[0];
38    let dict_id_flag = descriptor & 0x03;
39    let content_checksum = (descriptor & 0x04) != 0;
40    let single_segment = (descriptor & 0x20) != 0;
41    let fcs_field_size_flag = (descriptor >> 6) & 0x03;
42
43    let reserved = (descriptor & 0x08) != 0;
44    if reserved {
45        return Err(DecompressError::BadFrameHeader);
46    }
47    let unused = (descriptor & 0x10) != 0;
48    if unused {
49        return Err(DecompressError::BadFrameHeader);
50    }
51
52    let mut offset = 1;
53
54    let window_size = if single_segment {
55        0
56    } else {
57        if data.len() <= offset {
58            return Err(DecompressError::BadFrameHeader);
59        }
60        let window_desc = data[offset];
61        offset += 1;
62        let exponent = (window_desc >> 3) as u64;
63        let mantissa = (window_desc & 0x07) as u64;
64        let window_base = 1u64 << (10 + exponent);
65        let window_add = (window_base >> 3) * mantissa;
66        window_base + window_add
67    };
68
69    let dict_id_size = match dict_id_flag {
70        0 => 0,
71        1 => 1,
72        2 => 2,
73        3 => 4,
74        _ => unreachable!(),
75    };
76
77    let dict_id = if dict_id_size > 0 {
78        if data.len() < offset + dict_id_size {
79            return Err(DecompressError::BadFrameHeader);
80        }
81        let id = match dict_id_size {
82            1 => data[offset] as u32,
83            2 => u16::from_le_bytes([data[offset], data[offset + 1]]) as u32,
84            4 => u32::from_le_bytes([
85                data[offset],
86                data[offset + 1],
87                data[offset + 2],
88                data[offset + 3],
89            ]),
90            _ => unreachable!(),
91        };
92        offset += dict_id_size;
93        Some(id)
94    } else {
95        None
96    };
97
98    let fcs_field_size = match fcs_field_size_flag {
99        0 => {
100            if single_segment {
101                1
102            } else {
103                0
104            }
105        }
106        1 => 2,
107        2 => 4,
108        3 => 8,
109        _ => unreachable!(),
110    };
111
112    let frame_content_size = if fcs_field_size > 0 {
113        if data.len() < offset + fcs_field_size {
114            return Err(DecompressError::BadFrameHeader);
115        }
116        let fcs = match fcs_field_size {
117            1 => data[offset] as u64,
118            2 => u16::from_le_bytes([data[offset], data[offset + 1]]) as u64 + 256,
119            4 => u32::from_le_bytes([
120                data[offset],
121                data[offset + 1],
122                data[offset + 2],
123                data[offset + 3],
124            ]) as u64,
125            8 => u64::from_le_bytes([
126                data[offset],
127                data[offset + 1],
128                data[offset + 2],
129                data[offset + 3],
130                data[offset + 4],
131                data[offset + 5],
132                data[offset + 6],
133                data[offset + 7],
134            ]),
135            _ => unreachable!(),
136        };
137        offset += fcs_field_size;
138        Some(fcs)
139    } else {
140        None
141    };
142
143    let final_window_size = if single_segment {
144        frame_content_size.unwrap_or(0)
145    } else {
146        window_size
147    };
148
149    Ok(FrameHeader {
150        window_size: final_window_size,
151        frame_content_size,
152        dict_id,
153        content_checksum,
154        single_segment,
155        header_size: offset + magic_size,
156    })
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    #[test]
164    fn parse_minimal_header() {
165        let data = [0x28, 0xB5, 0x2F, 0xFD, 0x20, 0x00];
166        let hdr = parse_frame_header(&data).unwrap();
167        assert!(hdr.single_segment);
168        assert!(!hdr.content_checksum);
169        assert_eq!(hdr.dict_id, None);
170        assert_eq!(hdr.frame_content_size, Some(0));
171    }
172
173    #[test]
174    fn bad_magic() {
175        let data = [0x00, 0x00, 0x00, 0x00, 0x00];
176        assert!(matches!(
177            parse_frame_header(&data),
178            Err(DecompressError::BadMagic)
179        ));
180    }
181
182    #[test]
183    fn parse_header_after_magic() {
184        let data = [0x20, 0x00];
185        let hdr = parse_frame_header_after_magic(&data, 0).unwrap();
186        assert!(hdr.single_segment);
187        assert_eq!(hdr.frame_content_size, Some(0));
188        assert_eq!(hdr.header_size, 2);
189    }
190}