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