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 self.left_bits -= size;
50 } else if self.look_bits(size).is_ok() && self.left_bits >= size {
51 self.left_bits -= size;
52 } else {
53 self.left_bits = 0;
54 }
55 }
56
57 fn get_bits(&mut self, size: usize) -> Result<usize, Error> {
58 let bits = self.look_bits(size);
59 self.skip_bits(size);
60 bits
61 }
62
63 fn get_signed_bits(&mut self, size: usize) -> Result<isize, Error> {
64 let bits = self.get_bits(size - 1)? as isize;
65 let sign = self.get_bits(1)?;
66 if sign == 1 {
67 Ok(bits)
68 } else {
69 Ok(-bits)
70 }
71 }
72}
73
74pub struct AnimationControl {
76 pub backgroud_color: u32,
78 pub loop_count: u16,
80}
81
82pub struct AnimationFrame {
84 pub frame_x: usize,
86 pub frame_y: usize,
88 pub width: usize,
90 pub height: usize,
92 pub duration: usize,
94 pub alpha_blending: bool,
96 pub disopse: bool,
98 pub frame: Vec<u8>,
100 pub alpha: Option<Vec<u8>>,
102}
103
104pub struct WebpHeader {
106 pub width: usize,
108 pub height: usize,
110 pub canvas_width: usize,
112 pub canvas_height: usize,
114 pub image_chunksize: usize,
116 pub has_icc_profile: bool,
118 pub has_alpha: bool,
120 pub has_exif: bool,
122 pub has_xmp: bool,
124 pub has_animation: bool,
126 pub lossy: bool,
128 pub image: Vec<u8>,
130 pub icc_profile: Option<Vec<u8>>,
132 pub alpha: Option<Vec<u8>>,
134 pub exif: Option<Vec<u8>>,
136 pub xmp: Option<Vec<u8>>,
138 pub animation: Option<AnimationControl>,
140 pub animation_frame: Option<Vec<AnimationFrame>>,
142}
143
144impl WebpHeader {
145 pub fn new() -> Self {
146 Self {
147 width: 0,
148 height: 0,
149 canvas_width: 0,
150 canvas_height: 0,
151 image_chunksize: 0,
152 has_icc_profile: false,
153 has_alpha: false,
154 has_exif: false,
155 has_xmp: false,
156 has_animation: false,
157 lossy: false,
158 image: vec![],
159 icc_profile: None,
160 exif: None,
161 alpha: None,
162 xmp: None,
163 animation: None,
164 animation_frame: None,
165 }
166 }
167}
168
169pub fn read_u24<B: BinaryReader>(reader: &mut B) -> Result<u32, Error> {
171 let mut b = [0_u8; 3];
172 reader.read_exact(&mut b)?;
173 Ok((b[0] as u32) | ((b[1] as u32) << 8) | ((b[2] as u32) << 16))
174}
175
176fn parse_animation_frame_payload(data: &[u8]) -> Result<(Vec<u8>, Option<Vec<u8>>), Error> {
177 let mut reader = BytesReader::from(data.to_vec());
178 let mut frame = None;
179 let mut alpha = None;
180
181 while (reader.offset()? as usize) + 8 <= data.len() {
182 let chunk_id = reader.read_ascii_string(4)?;
183 let size = reader.read_u32_le()? as usize;
184 let chunk = reader.read_bytes_as_vec(size)?;
185 match chunk_id.as_str() {
186 "ALPH" => alpha = Some(chunk),
187 "VP8 " | "VP8L" => {
188 frame = Some(chunk);
189 break;
190 }
191 _ => {}
192 }
193 if size & 1 == 1 && (reader.offset()? as usize) < data.len() {
194 reader.skip_ptr(1)?;
195 }
196 }
197
198 frame
199 .map(|frame| (frame, alpha))
200 .ok_or_else(|| Box::new(std::io::Error::from(std::io::ErrorKind::Other)) as Error)
201}
202
203pub fn read_header<B: BinaryReader>(reader: &mut B) -> Result<WebpHeader, Error> {
205 let riff = reader.read_ascii_string(4)?;
206 if riff != "RIFF" {
207 return Err(Box::new(std::io::Error::from(std::io::ErrorKind::Other)));
208 }
209 let mut cksize = reader.read_u32_le()? as usize;
210 let webp = reader.read_ascii_string(4)?;
211 if webp != "WEBP" {
212 return Err(Box::new(std::io::Error::from(std::io::ErrorKind::Other)));
213 }
214 cksize -= 4;
215 let mut webp_header = WebpHeader::new();
216
217 loop {
218 let vp8 = reader.read_ascii_string(4)?;
219 let size = reader.read_u32_le()? as usize;
220 let padded_size = size + (size & 1);
221 match vp8.as_str() {
222 "VP8 " => {
223 webp_header.lossy = true;
224 webp_header.image_chunksize = size;
225 let buf = reader.read_bytes_as_vec(size)?;
226 let flags = buf[0] as usize | ((buf[1] as usize) << 8) | ((buf[2] as usize) << 8);
227 let key_frame = (flags & 0x0001) == 0;
228
229 let w = buf[6] as usize | ((buf[7] as usize) << 8);
230 webp_header.width = w & 0x3fff;
231 let w = buf[8] as usize | ((buf[9] as usize) << 8);
232 webp_header.height = w & 0x3fff;
233
234 let mut reader = BitReader::new(&buf[10..]);
235 if key_frame {
236 let _ = reader.get_bits(1)?;
237 let _ = reader.get_bits(1)?;
238 }
239
240 let mut quant = [0_isize; NUM_MB_SEGMENTS];
241 let mut filter = [0_isize; NUM_MB_SEGMENTS];
242 let mut seg = [0_usize; MB_FEATURE_TREE_PROBS];
243
244 let segmentation_enabled = reader.get_bits(1)?;
245 if segmentation_enabled == 1 {
246 let update_segment_feature_data = reader.get_bits(1)?;
247 if reader.get_bits(1)? == 1 {
248 for quant_item in quant.iter_mut().take(NUM_MB_SEGMENTS) {
249 *quant_item = if reader.get_bits(1)? == 1 {
250 reader.get_signed_bits(7)?
251 } else {
252 0
253 };
254 }
255 for filter_item in filter.iter_mut().take(NUM_MB_SEGMENTS) {
256 *filter_item = if reader.get_bits(1)? == 1 {
257 reader.get_signed_bits(6)?
258 } else {
259 0
260 };
261 }
262 }
263 if update_segment_feature_data == 1 {
264 for seg_item in seg.iter_mut().take(MB_FEATURE_TREE_PROBS) {
265 *seg_item = if reader.get_bits(1)? == 1 {
266 reader.get_bits(8)?
267 } else {
268 0
269 };
270 }
271 }
272 }
273
274 webp_header.image = buf;
275 }
276 "VP8L" => {
277 webp_header.lossy = false;
278 webp_header.image_chunksize = size;
279 webp_header.image = reader.read_bytes_as_vec(size)?;
280 }
281 "VP8X" => {
282 let flag = reader.read_byte()?;
283 if flag & 0x20 > 0 {
284 webp_header.has_icc_profile = true;
285 }
286 if flag & 0x10 > 0 {
287 webp_header.has_alpha = true;
288 }
289 if flag & 0x08 > 0 {
290 webp_header.has_exif = true;
291 }
292 if flag & 0x04 > 0 {
293 webp_header.has_xmp = true;
294 }
295 if flag & 0x02 > 0 {
296 webp_header.has_animation = true;
297 }
298
299 let _ = read_u24(reader)?;
300 webp_header.canvas_width = read_u24(reader)? as usize + 1;
301 webp_header.canvas_height = read_u24(reader)? as usize + 1;
302 if size > 10 {
303 reader.skip_ptr(size - 10)?;
304 }
305 }
306 "ALPH" => {
307 if webp_header.has_alpha {
308 webp_header.alpha = Some(reader.read_bytes_as_vec(size)?);
309 } else {
310 reader.skip_ptr(size)?;
311 }
312 }
313 "ANIM" => {
314 if webp_header.has_animation {
315 let backgroud_color = reader.read_u32_le()?;
316 let loop_count = reader.read_u16_le()?;
317 if size > 8 {
318 reader.skip_ptr(size - 8)?;
319 }
320 webp_header.animation = Some(AnimationControl {
321 backgroud_color,
322 loop_count,
323 });
324 } else {
325 reader.skip_ptr(size)?;
326 }
327 }
328 "ANMF" | "ANIF" => {
329 if webp_header.has_animation {
330 let frame_x = read_u24(reader)? as usize * 2;
331 let frame_y = read_u24(reader)? as usize * 2;
332 let width = read_u24(reader)? as usize + 1;
333 let height = read_u24(reader)? as usize + 1;
334 let duration = read_u24(reader)? as usize;
335 let flag = reader.read_byte()?;
336 let alpha_blending = (flag & 0x02) == 0;
337 let disopse = (flag & 0x01) != 0;
338
339 let buf = reader.read_bytes_as_vec(size - 16)?;
340 let (frame, alpha) = parse_animation_frame_payload(&buf)?;
341 let animation_frame = AnimationFrame {
342 frame_x,
343 frame_y,
344 width,
345 height,
346 duration,
347 alpha_blending,
348 disopse,
349 frame,
350 alpha,
351 };
352 if let Some(frames) = webp_header.animation_frame.as_mut() {
353 frames.push(animation_frame);
354 } else {
355 webp_header.animation_frame = Some(vec![animation_frame]);
356 }
357 } else {
358 reader.skip_ptr(size)?;
359 }
360 }
361 "EXIF" => {
362 if webp_header.has_exif {
363 webp_header.exif = Some(reader.read_bytes_as_vec(size)?);
364 } else {
365 reader.skip_ptr(size)?;
366 }
367 }
368 "XMP " => {
369 if webp_header.has_xmp {
370 webp_header.xmp = Some(reader.read_bytes_as_vec(size)?);
371 } else {
372 reader.skip_ptr(size)?;
373 }
374 }
375 "ICCP" => {
376 if webp_header.has_icc_profile {
377 webp_header.icc_profile = Some(reader.read_bytes_as_vec(size)?);
378 } else {
379 reader.skip_ptr(size)?;
380 }
381 }
382 _ => {
383 reader.skip_ptr(size)?;
384 }
385 }
386 if size & 1 == 1 {
387 reader.skip_ptr(1)?;
388 }
389 if cksize <= padded_size + 8 {
390 break;
391 }
392 cksize -= padded_size + 8;
393 }
394 Ok(webp_header)
395}