1use bin_rs::reader::{BinaryReader, BytesReader};
4
5type Error = Box<dyn std::error::Error>;
6
7const MB_FEATURE_TREE_PROBS: usize = 3;
8const NUM_MB_SEGMENTS: usize = 4;
9
10pub(crate) struct BitReader {
11 pub buffer: Vec<u8>,
12 ptr: usize,
13 left_bits: usize,
14 last_byte: u32,
15 warning: bool,
16}
17
18impl BitReader {
19 pub fn new(data: &[u8]) -> Self {
20 Self {
21 buffer: data.to_vec(),
22 last_byte: 0,
23 ptr: 0,
24 left_bits: 0,
25 warning: false,
26 }
27 }
28
29 fn look_bits(&mut self, size: usize) -> Result<usize, Error> {
30 while self.left_bits < size {
31 if self.ptr >= self.buffer.len() {
32 self.warning = true;
33 if size >= 12 {
34 return Ok(0x1);
35 }
36 return Ok(0x0);
37 }
38 self.last_byte = (self.last_byte << 8) | (self.buffer[self.ptr] as u32);
39 self.ptr += 1;
40 self.left_bits += 8;
41 }
42
43 let bits = (self.last_byte >> (self.left_bits - size)) & ((1 << size) - 1);
44 Ok(bits as usize)
45 }
46
47 fn skip_bits(&mut self, size: usize) {
48 if self.left_bits <= size {
49 let _ = self.look_bits(size);
50 }
51 self.left_bits = self.left_bits.saturating_sub(size);
52 }
53
54 fn get_bits(&mut self, size: usize) -> Result<usize, Error> {
55 let bits = self.look_bits(size);
56 self.skip_bits(size);
57 bits
58 }
59
60 fn get_signed_bits(&mut self, size: usize) -> Result<isize, Error> {
61 let bits = self.get_bits(size - 1)? as isize;
62 let sign = self.get_bits(1)?;
63 if sign == 1 {
64 Ok(bits)
65 } else {
66 Ok(-bits)
67 }
68 }
69}
70
71pub struct AnimationControl {
73 pub backgroud_color: u32,
75 pub loop_count: u16,
77}
78
79pub struct AnimationFrame {
81 pub frame_x: usize,
83 pub frame_y: usize,
85 pub width: usize,
87 pub height: usize,
89 pub duration: usize,
91 pub alpha_blending: bool,
93 pub disopse: bool,
95 pub frame: Vec<u8>,
97 pub alpha: Option<Vec<u8>>,
99}
100
101pub struct WebpHeader {
103 pub width: usize,
105 pub height: usize,
107 pub canvas_width: usize,
109 pub canvas_height: usize,
111 pub image_chunksize: usize,
113 pub has_icc_profile: bool,
115 pub has_alpha: bool,
117 pub has_exif: bool,
119 pub has_xmp: bool,
121 pub has_animation: bool,
123 pub lossy: bool,
125 pub image: Vec<u8>,
127 pub icc_profile: Option<Vec<u8>>,
129 pub alpha: Option<Vec<u8>>,
131 pub exif: Option<Vec<u8>>,
133 pub xmp: Option<Vec<u8>>,
135 pub animation: Option<AnimationControl>,
137 pub animation_frame: Option<Vec<AnimationFrame>>,
139}
140
141impl WebpHeader {
142 pub fn new() -> Self {
143 Self {
144 width: 0,
145 height: 0,
146 canvas_width: 0,
147 canvas_height: 0,
148 image_chunksize: 0,
149 has_icc_profile: false,
150 has_alpha: false,
151 has_exif: false,
152 has_xmp: false,
153 has_animation: false,
154 lossy: false,
155 image: vec![],
156 icc_profile: None,
157 exif: None,
158 alpha: None,
159 xmp: None,
160 animation: None,
161 animation_frame: None,
162 }
163 }
164}
165
166impl Default for WebpHeader {
167 fn default() -> Self {
168 Self::new()
169 }
170}
171
172pub fn read_u24<B: BinaryReader>(reader: &mut B) -> Result<u32, Error> {
174 let mut b = [0_u8; 3];
175 reader.read_exact(&mut b)?;
176 Ok((b[0] as u32) | ((b[1] as u32) << 8) | ((b[2] as u32) << 16))
177}
178
179fn parse_animation_frame_payload(data: &[u8]) -> Result<(Vec<u8>, Option<Vec<u8>>), Error> {
180 let mut reader = BytesReader::from(data.to_vec());
181 let mut frame = None;
182 let mut alpha = None;
183
184 while (reader.offset()? as usize) + 8 <= data.len() {
185 let chunk_id = reader.read_ascii_string(4)?;
186 let size = reader.read_u32_le()? as usize;
187 let chunk = reader.read_bytes_as_vec(size)?;
188 match chunk_id.as_str() {
189 "ALPH" => alpha = Some(chunk),
190 "VP8 " | "VP8L" => {
191 frame = Some(chunk);
192 break;
193 }
194 _ => {}
195 }
196 if size & 1 == 1 && (reader.offset()? as usize) < data.len() {
197 reader.skip_ptr(1)?;
198 }
199 }
200
201 frame
202 .map(|frame| (frame, alpha))
203 .ok_or_else(|| Box::new(std::io::Error::from(std::io::ErrorKind::Other)) as Error)
204}
205
206#[allow(clippy::collapsible_match)]
208pub fn read_header<B: BinaryReader>(reader: &mut B) -> Result<WebpHeader, Error> {
209 let riff = reader.read_ascii_string(4)?;
210 if riff != "RIFF" {
211 return Err(Box::new(std::io::Error::from(std::io::ErrorKind::Other)));
212 }
213 let mut cksize = reader.read_u32_le()? as usize;
214 let webp = reader.read_ascii_string(4)?;
215 if webp != "WEBP" {
216 return Err(Box::new(std::io::Error::from(std::io::ErrorKind::Other)));
217 }
218 cksize -= 4;
219 let mut webp_header = WebpHeader::new();
220
221 loop {
222 let vp8 = reader.read_ascii_string(4)?;
223 let size = reader.read_u32_le()? as usize;
224 let padded_size = size + (size & 1);
225 match vp8.as_str() {
226 "VP8 " => {
227 webp_header.lossy = true;
228 webp_header.image_chunksize = size;
229 let buf = reader.read_bytes_as_vec(size)?;
230 let flags = buf[0] as usize | ((buf[1] as usize) << 8) | ((buf[2] as usize) << 8);
231 let key_frame = (flags & 0x0001) == 0;
232
233 let w = buf[6] as usize | ((buf[7] as usize) << 8);
234 webp_header.width = w & 0x3fff;
235 let w = buf[8] as usize | ((buf[9] as usize) << 8);
236 webp_header.height = w & 0x3fff;
237
238 let mut reader = BitReader::new(&buf[10..]);
239 if key_frame {
240 let _ = reader.get_bits(1)?;
241 let _ = reader.get_bits(1)?;
242 }
243
244 let mut quant = [0_isize; NUM_MB_SEGMENTS];
245 let mut filter = [0_isize; NUM_MB_SEGMENTS];
246 let mut seg = [0_usize; MB_FEATURE_TREE_PROBS];
247
248 let segmentation_enabled = reader.get_bits(1)?;
249 if segmentation_enabled == 1 {
250 let update_segment_feature_data = reader.get_bits(1)?;
251 if reader.get_bits(1)? == 1 {
252 for quant_item in quant.iter_mut().take(NUM_MB_SEGMENTS) {
253 *quant_item = if reader.get_bits(1)? == 1 {
254 reader.get_signed_bits(7)?
255 } else {
256 0
257 };
258 }
259 for filter_item in filter.iter_mut().take(NUM_MB_SEGMENTS) {
260 *filter_item = if reader.get_bits(1)? == 1 {
261 reader.get_signed_bits(6)?
262 } else {
263 0
264 };
265 }
266 }
267 if update_segment_feature_data == 1 {
268 for seg_item in seg.iter_mut().take(MB_FEATURE_TREE_PROBS) {
269 *seg_item = if reader.get_bits(1)? == 1 {
270 reader.get_bits(8)?
271 } else {
272 0
273 };
274 }
275 }
276 }
277
278 webp_header.image = buf;
279 }
280 "VP8L" => {
281 webp_header.lossy = false;
282 webp_header.image_chunksize = size;
283 webp_header.image = reader.read_bytes_as_vec(size)?;
284 }
285 "VP8X" => {
286 let flag = reader.read_byte()?;
287 if flag & 0x20 > 0 {
288 webp_header.has_icc_profile = true;
289 }
290 if flag & 0x10 > 0 {
291 webp_header.has_alpha = true;
292 }
293 if flag & 0x08 > 0 {
294 webp_header.has_exif = true;
295 }
296 if flag & 0x04 > 0 {
297 webp_header.has_xmp = true;
298 }
299 if flag & 0x02 > 0 {
300 webp_header.has_animation = true;
301 }
302
303 let _ = read_u24(reader)?;
304 webp_header.canvas_width = read_u24(reader)? as usize + 1;
305 webp_header.canvas_height = read_u24(reader)? as usize + 1;
306 if size > 10 {
307 reader.skip_ptr(size - 10)?;
308 }
309 }
310 "ALPH" => {
311 if webp_header.has_alpha {
312 webp_header.alpha = Some(reader.read_bytes_as_vec(size)?);
313 } else {
314 reader.skip_ptr(size)?;
315 }
316 }
317 "ANIM" => {
318 if webp_header.has_animation {
319 let backgroud_color = reader.read_u32_le()?;
320 let loop_count = reader.read_u16_le()?;
321 if size > 8 {
322 reader.skip_ptr(size - 8)?;
323 }
324 webp_header.animation = Some(AnimationControl {
325 backgroud_color,
326 loop_count,
327 });
328 } else {
329 reader.skip_ptr(size)?;
330 }
331 }
332 "ANMF" | "ANIF" => {
333 if webp_header.has_animation {
334 let frame_x = read_u24(reader)? as usize * 2;
335 let frame_y = read_u24(reader)? as usize * 2;
336 let width = read_u24(reader)? as usize + 1;
337 let height = read_u24(reader)? as usize + 1;
338 let duration = read_u24(reader)? as usize;
339 let flag = reader.read_byte()?;
340 let alpha_blending = (flag & 0x02) == 0;
341 let disopse = (flag & 0x01) != 0;
342
343 let buf = reader.read_bytes_as_vec(size - 16)?;
344 let (frame, alpha) = parse_animation_frame_payload(&buf)?;
345 let animation_frame = AnimationFrame {
346 frame_x,
347 frame_y,
348 width,
349 height,
350 duration,
351 alpha_blending,
352 disopse,
353 frame,
354 alpha,
355 };
356 if let Some(frames) = webp_header.animation_frame.as_mut() {
357 frames.push(animation_frame);
358 } else {
359 webp_header.animation_frame = Some(vec![animation_frame]);
360 }
361 } else {
362 reader.skip_ptr(size)?;
363 }
364 }
365 "EXIF" => {
366 if webp_header.has_exif {
367 webp_header.exif = Some(reader.read_bytes_as_vec(size)?);
368 } else {
369 reader.skip_ptr(size)?;
370 }
371 }
372 "XMP " => {
373 if webp_header.has_xmp {
374 webp_header.xmp = Some(reader.read_bytes_as_vec(size)?);
375 } else {
376 reader.skip_ptr(size)?;
377 }
378 }
379 "ICCP" => {
380 if webp_header.has_icc_profile {
381 webp_header.icc_profile = Some(reader.read_bytes_as_vec(size)?);
382 } else {
383 reader.skip_ptr(size)?;
384 }
385 }
386 _ => {
387 reader.skip_ptr(size)?;
388 }
389 }
390 if size & 1 == 1 {
391 reader.skip_ptr(1)?;
392 }
393 if cksize <= padded_size + 8 {
394 break;
395 }
396 cksize -= padded_size + 8;
397 }
398 Ok(webp_header)
399}