Skip to main content

bao_browser/
screenshot.rs

1// @trace REQ-BRW-002  REQ-CDP-007: Screenshot capture via servo rendering pipeline
2use std::io::Cursor;
3
4use image::{ImageFormat, RgbaImage};
5
6use crate::error::BrowserError;
7
8pub enum ScreenshotFormat {
9    Png,
10    Jpeg,
11    WebP,
12}
13
14pub fn encode_image(image: &RgbaImage, format: ScreenshotFormat) -> Result<Vec<u8>, BrowserError> {
15    let mut buf = Cursor::new(Vec::with_capacity(
16        image.width() as usize * image.height() as usize,
17    ));
18    match format {
19        ScreenshotFormat::Png => image
20            .write_to(&mut buf, ImageFormat::Png)
21            .map_err(|e| BrowserError::Rendering(format!("PNG encode failed: {e}")))?,
22        ScreenshotFormat::Jpeg => {
23            let rgb = image::DynamicImage::ImageRgba8(image.clone()).to_rgb8();
24            rgb.write_to(&mut buf, ImageFormat::Jpeg)
25                .map_err(|e| BrowserError::Rendering(format!("JPEG encode failed: {e}")))?;
26        }
27        ScreenshotFormat::WebP => image
28            .write_to(&mut buf, ImageFormat::WebP)
29            .map_err(|e| BrowserError::Rendering(format!("WebP encode failed: {e}")))?,
30    }
31    Ok(buf.into_inner())
32}
33
34#[cfg(test)]
35mod tests {
36    // @trace REQ-BRW-002 [req:REQ-BRW-002,REQ-CDP-007] [level:unit]
37    use super::*;
38    use image::Rgba;
39
40    fn red_image(w: u32, h: u32) -> RgbaImage {
41        RgbaImage::from_pixel(w, h, Rgba([255, 0, 0, 255]))
42    }
43
44    #[test]
45    fn encode_1x1_png_returns_nonempty_with_magic() {
46        let img = red_image(1, 1);
47        let out = encode_image(&img, ScreenshotFormat::Png).unwrap();
48        assert!(!out.is_empty());
49        assert_eq!(&out[0..4], &[0x89, 0x50, 0x4E, 0x47]);
50    }
51
52    #[test]
53    fn encode_1x1_jpeg_returns_nonempty_with_magic() {
54        let img = red_image(1, 1);
55        let out = encode_image(&img, ScreenshotFormat::Jpeg).unwrap();
56        assert!(!out.is_empty());
57        assert_eq!(&out[0..2], &[0xFF, 0xD8]);
58    }
59
60    #[test]
61    fn encode_100x100_png_produces_valid_data() {
62        let img = red_image(100, 100);
63        let out = encode_image(&img, ScreenshotFormat::Png).unwrap();
64        assert!(!out.is_empty());
65        assert_eq!(&out[0..4], &[0x89, 0x50, 0x4E, 0x47]);
66    }
67
68    #[test]
69    fn encode_100x100_jpeg_produces_valid_data() {
70        let img = red_image(100, 100);
71        let out = encode_image(&img, ScreenshotFormat::Jpeg).unwrap();
72        assert!(!out.is_empty());
73        assert_eq!(&out[0..2], &[0xFF, 0xD8]);
74    }
75
76    #[test]
77    fn encode_1x1_webp_returns_nonempty_with_riff_magic() {
78        let img = red_image(1, 1);
79        let out = encode_image(&img, ScreenshotFormat::WebP).unwrap();
80        assert!(!out.is_empty());
81        // WebP files start with "RIFF" header
82        assert_eq!(&out[0..4], b"RIFF");
83    }
84
85    #[test]
86    fn encode_100x100_webp_produces_valid_data() {
87        let img = red_image(100, 100);
88        let out = encode_image(&img, ScreenshotFormat::WebP).unwrap();
89        assert!(!out.is_empty());
90        assert_eq!(&out[0..4], b"RIFF");
91    }
92
93    #[test]
94    fn encode_empty_image_no_panic() {
95        let img = red_image(0, 0);
96        let _ = encode_image(&img, ScreenshotFormat::Png);
97        let _ = encode_image(&img, ScreenshotFormat::Jpeg);
98        let _ = encode_image(&img, ScreenshotFormat::WebP);
99    }
100
101    #[test]
102    fn all_formats_produce_nontrivial_output() {
103        let img = red_image(200, 200);
104        let png = encode_image(&img, ScreenshotFormat::Png).unwrap();
105        let jpeg = encode_image(&img, ScreenshotFormat::Jpeg).unwrap();
106        let webp = encode_image(&img, ScreenshotFormat::WebP).unwrap();
107        assert!(
108            png.len() > 100,
109            "PNG should be nontrivial: {} bytes",
110            png.len()
111        );
112        assert!(
113            jpeg.len() > 100,
114            "JPEG should be nontrivial: {} bytes",
115            jpeg.len()
116        );
117        assert!(
118            webp.len() > 100,
119            "WebP should be nontrivial: {} bytes",
120            webp.len()
121        );
122    }
123}