forme/
image_loader.rs

1//! # Image Loading and Decoding
2//!
3//! Loads images from file paths, data URIs, or raw base64 strings and prepares
4//! them for PDF embedding. JPEG images pass through without re-encoding
5//! (the PDF spec supports DCTDecode natively). PNG images are decoded to RGB
6//! pixels with a separate alpha channel for SMask transparency.
7
8use std::io::Cursor;
9
10/// A fully decoded/loaded image ready for PDF embedding.
11#[derive(Debug, Clone)]
12pub struct LoadedImage {
13    pub pixel_data: ImagePixelData,
14    pub width_px: u32,
15    pub height_px: u32,
16}
17
18/// The pixel data in a format the PDF serializer can consume directly.
19#[derive(Debug, Clone)]
20pub enum ImagePixelData {
21    /// Raw JPEG bytes — embed directly with DCTDecode.
22    Jpeg {
23        data: Vec<u8>,
24        color_space: JpegColorSpace,
25    },
26    /// Decoded RGB pixels + optional alpha channel.
27    Decoded {
28        /// width * height * 3 bytes (RGB)
29        rgb: Vec<u8>,
30        /// width * height bytes (grayscale alpha). None if fully opaque.
31        alpha: Option<Vec<u8>>,
32    },
33}
34
35/// JPEG color space for the PDF /ColorSpace entry.
36#[derive(Debug, Clone, Copy)]
37pub enum JpegColorSpace {
38    DeviceRGB,
39    DeviceGray,
40}
41
42/// Load an image from a source string.
43///
44/// Supported `src` formats:
45/// - `data:image/...;base64,...` — data URI
46/// - File path (absolute or relative) — reads from disk
47/// - Raw base64-encoded image data
48pub fn load_image(src: &str) -> Result<LoadedImage, String> {
49    let raw_bytes = read_source_bytes(src)?;
50    decode_image_bytes(&raw_bytes)
51}
52
53/// Read only the pixel dimensions from an image source without decoding pixels.
54/// Returns (width, height) in pixels. Much cheaper than `load_image`.
55pub fn load_image_dimensions(src: &str) -> Result<(u32, u32), String> {
56    let raw_bytes = read_source_bytes(src)?;
57    let reader = image::io::Reader::new(Cursor::new(&raw_bytes))
58        .with_guessed_format()
59        .map_err(|e| format!("Image format detection error: {}", e))?;
60    reader
61        .into_dimensions()
62        .map_err(|e| format!("Failed to read image dimensions: {}", e))
63}
64
65/// Resolve the source string to raw image bytes.
66fn read_source_bytes(src: &str) -> Result<Vec<u8>, String> {
67    // Data URI: data:image/png;base64,iVBOR...
68    if src.starts_with("data:image/") {
69        let comma_pos = src
70            .find(',')
71            .ok_or_else(|| "Invalid data URI: missing comma".to_string())?;
72        let b64_data = &src[comma_pos + 1..];
73        return base64_decode(b64_data);
74    }
75
76    // File path — try reading from disk (not available in WASM)
77    // Only match explicit path prefixes to avoid treating base64 strings
78    // (which contain '/') as file paths.
79    if src.starts_with('/') || src.starts_with("./") || src.starts_with("../") {
80        #[cfg(not(target_arch = "wasm32"))]
81        {
82            return std::fs::read(src)
83                .map_err(|e| format!("Failed to read image file '{}': {}", src, e));
84        }
85        #[cfg(target_arch = "wasm32")]
86        {
87            return Err(format!(
88                "File path images not supported in WASM: '{}'. Use data URIs or base64.",
89                src
90            ));
91        }
92    }
93
94    // Try raw base64
95    base64_decode(src)
96}
97
98fn base64_decode(input: &str) -> Result<Vec<u8>, String> {
99    use base64::Engine;
100    base64::engine::general_purpose::STANDARD
101        .decode(input)
102        .map_err(|e| format!("Base64 decode error: {}", e))
103}
104
105/// Detect image format from magic bytes and decode accordingly.
106fn decode_image_bytes(data: &[u8]) -> Result<LoadedImage, String> {
107    if data.len() < 4 {
108        return Err("Image data too short".to_string());
109    }
110
111    if is_jpeg(data) {
112        decode_jpeg(data)
113    } else if is_png(data) {
114        decode_png(data)
115    } else if is_webp(data) {
116        decode_webp(data)
117    } else {
118        Err("Unsupported image format (expected JPEG, PNG, or WebP)".to_string())
119    }
120}
121
122fn is_jpeg(data: &[u8]) -> bool {
123    data.len() >= 2 && data[0] == 0xFF && data[1] == 0xD8
124}
125
126fn is_png(data: &[u8]) -> bool {
127    data.len() >= 4 && data[0] == 0x89 && data[1] == 0x50 && data[2] == 0x4E && data[3] == 0x47
128}
129
130fn is_webp(data: &[u8]) -> bool {
131    data.len() >= 12 && &data[0..4] == b"RIFF" && &data[8..12] == b"WEBP"
132}
133
134/// JPEG: read dimensions and color space without decoding pixels.
135/// The raw JPEG bytes are passed through to the PDF (DCTDecode).
136fn decode_jpeg(data: &[u8]) -> Result<LoadedImage, String> {
137    let reader = image::io::Reader::new(Cursor::new(data))
138        .with_guessed_format()
139        .map_err(|e| format!("JPEG format detection error: {}", e))?;
140
141    let (width, height) = reader
142        .into_dimensions()
143        .map_err(|e| format!("Failed to read JPEG dimensions: {}", e))?;
144
145    // Detect color space from JPEG component count.
146    // Re-read to get color info since into_dimensions() consumed the reader.
147    let color_space = detect_jpeg_color_space(data);
148
149    Ok(LoadedImage {
150        pixel_data: ImagePixelData::Jpeg {
151            data: data.to_vec(),
152            color_space,
153        },
154        width_px: width,
155        height_px: height,
156    })
157}
158
159/// Scan JPEG markers to find the SOF (Start of Frame) segment and read
160/// the number of components to determine color space.
161fn detect_jpeg_color_space(data: &[u8]) -> JpegColorSpace {
162    let mut i = 2; // skip SOI marker (FF D8)
163    while i + 1 < data.len() {
164        if data[i] != 0xFF {
165            break;
166        }
167        let marker = data[i + 1];
168        // SOF markers: C0-C3, C5-C7, C9-CB, CD-CF
169        let is_sof = matches!(marker, 0xC0..=0xC3 | 0xC5..=0xC7 | 0xC9..=0xCB | 0xCD..=0xCF);
170        if is_sof {
171            // SOF segment: length(2) + precision(1) + height(2) + width(2) + num_components(1)
172            if i + 9 < data.len() {
173                let num_components = data[i + 9];
174                return if num_components == 1 {
175                    JpegColorSpace::DeviceGray
176                } else {
177                    JpegColorSpace::DeviceRGB
178                };
179            }
180        }
181        // Skip to next marker
182        if i + 3 < data.len() {
183            let seg_len = u16::from_be_bytes([data[i + 2], data[i + 3]]) as usize;
184            i += 2 + seg_len;
185        } else {
186            break;
187        }
188    }
189    // Default to RGB if we can't determine
190    JpegColorSpace::DeviceRGB
191}
192
193/// WebP: decode to RGBA, split into RGB + alpha (same pipeline as PNG).
194fn decode_webp(data: &[u8]) -> Result<LoadedImage, String> {
195    let reader = image::io::Reader::new(Cursor::new(data))
196        .with_guessed_format()
197        .map_err(|e| format!("WebP format detection error: {}", e))?;
198
199    let img = reader
200        .decode()
201        .map_err(|e| format!("Failed to decode WebP: {}", e))?;
202
203    let rgba = img.to_rgba8();
204    let width = rgba.width();
205    let height = rgba.height();
206
207    let pixel_count = (width * height) as usize;
208    let mut rgb = Vec::with_capacity(pixel_count * 3);
209    let mut alpha = Vec::with_capacity(pixel_count);
210    let mut has_transparency = false;
211
212    for pixel in rgba.pixels() {
213        rgb.push(pixel[0]);
214        rgb.push(pixel[1]);
215        rgb.push(pixel[2]);
216        let a = pixel[3];
217        alpha.push(a);
218        if a != 255 {
219            has_transparency = true;
220        }
221    }
222
223    Ok(LoadedImage {
224        pixel_data: ImagePixelData::Decoded {
225            rgb,
226            alpha: if has_transparency { Some(alpha) } else { None },
227        },
228        width_px: width,
229        height_px: height,
230    })
231}
232
233/// PNG: decode to RGBA, split into RGB + alpha.
234fn decode_png(data: &[u8]) -> Result<LoadedImage, String> {
235    let reader = image::io::Reader::new(Cursor::new(data))
236        .with_guessed_format()
237        .map_err(|e| format!("PNG format detection error: {}", e))?;
238
239    let img = reader
240        .decode()
241        .map_err(|e| format!("Failed to decode PNG: {}", e))?;
242
243    let rgba = img.to_rgba8();
244    let width = rgba.width();
245    let height = rgba.height();
246
247    let pixel_count = (width * height) as usize;
248    let mut rgb = Vec::with_capacity(pixel_count * 3);
249    let mut alpha = Vec::with_capacity(pixel_count);
250    let mut has_transparency = false;
251
252    for pixel in rgba.pixels() {
253        rgb.push(pixel[0]);
254        rgb.push(pixel[1]);
255        rgb.push(pixel[2]);
256        let a = pixel[3];
257        alpha.push(a);
258        if a != 255 {
259            has_transparency = true;
260        }
261    }
262
263    Ok(LoadedImage {
264        pixel_data: ImagePixelData::Decoded {
265            rgb,
266            alpha: if has_transparency { Some(alpha) } else { None },
267        },
268        width_px: width,
269        height_px: height,
270    })
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276
277    #[test]
278    fn test_is_jpeg() {
279        assert!(is_jpeg(&[0xFF, 0xD8, 0xFF, 0xE0]));
280        assert!(!is_jpeg(&[0x89, 0x50, 0x4E, 0x47]));
281        assert!(!is_jpeg(&[0xFF]));
282    }
283
284    #[test]
285    fn test_is_png() {
286        assert!(is_png(&[0x89, 0x50, 0x4E, 0x47]));
287        assert!(!is_png(&[0xFF, 0xD8, 0xFF, 0xE0]));
288        assert!(!is_png(&[0x89, 0x50]));
289    }
290
291    #[test]
292    fn test_invalid_data_uri() {
293        let result = load_image("data:image/png;base64");
294        assert!(result.is_err());
295    }
296
297    #[test]
298    fn test_too_short_data() {
299        let result = decode_image_bytes(&[0x00, 0x01]);
300        assert!(result.is_err());
301    }
302
303    #[test]
304    fn test_is_webp() {
305        // Valid RIFF....WEBP header
306        assert!(is_webp(b"RIFF\x00\x00\x00\x00WEBP"));
307        // Too short
308        assert!(!is_webp(b"RIFF\x00\x00\x00\x00WEB"));
309        // Wrong prefix
310        assert!(!is_webp(b"XXXX\x00\x00\x00\x00WEBP"));
311        // Wrong suffix
312        assert!(!is_webp(b"RIFF\x00\x00\x00\x00XXXX"));
313        // JPEG bytes
314        assert!(!is_webp(&[0xFF, 0xD8, 0xFF, 0xE0]));
315    }
316
317    #[test]
318    fn test_decode_webp() {
319        // Create a 2x2 RGBA image, encode as WebP, then decode via our pipeline
320        let img = image::RgbaImage::from_fn(2, 2, |_, _| image::Rgba([255, 0, 0, 255]));
321        let mut buf = Vec::new();
322        let encoder = image::codecs::webp::WebPEncoder::new_lossless(&mut buf);
323        image::ImageEncoder::write_image(encoder, img.as_raw(), 2, 2, image::ColorType::Rgba8)
324            .unwrap();
325
326        let loaded = decode_image_bytes(&buf).unwrap();
327        assert_eq!(loaded.width_px, 2);
328        assert_eq!(loaded.height_px, 2);
329        match &loaded.pixel_data {
330            ImagePixelData::Decoded { rgb, alpha } => {
331                // All pixels should be red
332                for i in 0..4 {
333                    assert_eq!(rgb[i * 3], 255, "R channel");
334                    assert_eq!(rgb[i * 3 + 1], 0, "G channel");
335                    assert_eq!(rgb[i * 3 + 2], 0, "B channel");
336                }
337                assert!(alpha.is_none(), "Fully opaque should have no alpha");
338            }
339            _ => panic!("WebP should decode to Decoded variant"),
340        }
341    }
342
343    #[test]
344    fn test_decode_webp_with_alpha() {
345        let img = image::RgbaImage::from_fn(2, 2, |x, _| {
346            if x == 0 {
347                image::Rgba([0, 255, 0, 128])
348            } else {
349                image::Rgba([0, 255, 0, 255])
350            }
351        });
352        let mut buf = Vec::new();
353        let encoder = image::codecs::webp::WebPEncoder::new_lossless(&mut buf);
354        image::ImageEncoder::write_image(encoder, img.as_raw(), 2, 2, image::ColorType::Rgba8)
355            .unwrap();
356
357        let loaded = decode_image_bytes(&buf).unwrap();
358        match &loaded.pixel_data {
359            ImagePixelData::Decoded { rgb: _, alpha } => {
360                let alpha = alpha.as_ref().expect("Should have alpha channel");
361                assert_eq!(alpha[0], 128);
362                assert_eq!(alpha[1], 255);
363            }
364            _ => panic!("WebP should decode to Decoded variant"),
365        }
366    }
367
368    #[test]
369    fn test_unsupported_format() {
370        let result = decode_image_bytes(&[0x00, 0x01, 0x02, 0x03, 0x04]);
371        assert!(result.is_err());
372    }
373
374    #[test]
375    fn test_decode_minimal_png() {
376        // Create a 1x1 red PNG using the image crate
377        let mut img = image::RgbaImage::new(1, 1);
378        img.put_pixel(0, 0, image::Rgba([255, 0, 0, 255]));
379
380        let mut buf = Vec::new();
381        let encoder = image::codecs::png::PngEncoder::new(&mut buf);
382        image::ImageEncoder::write_image(encoder, img.as_raw(), 1, 1, image::ColorType::Rgba8)
383            .unwrap();
384
385        let loaded = decode_image_bytes(&buf).unwrap();
386        assert_eq!(loaded.width_px, 1);
387        assert_eq!(loaded.height_px, 1);
388        match &loaded.pixel_data {
389            ImagePixelData::Decoded { rgb, alpha } => {
390                assert_eq!(rgb, &[255, 0, 0]);
391                assert!(alpha.is_none(), "Fully opaque should have no alpha");
392            }
393            _ => panic!("PNG should decode to Decoded variant"),
394        }
395    }
396
397    #[test]
398    fn test_decode_png_with_alpha() {
399        let mut img = image::RgbaImage::new(1, 1);
400        img.put_pixel(0, 0, image::Rgba([255, 0, 0, 128]));
401
402        let mut buf = Vec::new();
403        let encoder = image::codecs::png::PngEncoder::new(&mut buf);
404        image::ImageEncoder::write_image(encoder, img.as_raw(), 1, 1, image::ColorType::Rgba8)
405            .unwrap();
406
407        let loaded = decode_image_bytes(&buf).unwrap();
408        match &loaded.pixel_data {
409            ImagePixelData::Decoded { rgb, alpha } => {
410                assert_eq!(rgb, &[255, 0, 0]);
411                assert_eq!(alpha.as_ref().unwrap(), &[128]);
412            }
413            _ => panic!("PNG should decode to Decoded variant"),
414        }
415    }
416
417    #[test]
418    fn test_decode_minimal_jpeg() {
419        // Create a 2x2 JPEG (JPEG requires min 1x1 but some encoders need 2x2)
420        let img = image::RgbImage::from_fn(2, 2, |_, _| image::Rgb([0, 128, 255]));
421
422        let mut buf = Vec::new();
423        let encoder = image::codecs::jpeg::JpegEncoder::new(&mut buf);
424        image::ImageEncoder::write_image(encoder, img.as_raw(), 2, 2, image::ColorType::Rgb8)
425            .unwrap();
426
427        let loaded = decode_image_bytes(&buf).unwrap();
428        assert_eq!(loaded.width_px, 2);
429        assert_eq!(loaded.height_px, 2);
430        match &loaded.pixel_data {
431            ImagePixelData::Jpeg { data, color_space } => {
432                assert!(data.starts_with(&[0xFF, 0xD8]));
433                assert!(matches!(color_space, JpegColorSpace::DeviceRGB));
434            }
435            _ => panic!("JPEG should stay as Jpeg variant"),
436        }
437    }
438
439    #[test]
440    fn test_base64_data_uri() {
441        // Create a small PNG, encode as data URI
442        let mut img = image::RgbaImage::new(1, 1);
443        img.put_pixel(0, 0, image::Rgba([0, 255, 0, 255]));
444
445        let mut buf = Vec::new();
446        let encoder = image::codecs::png::PngEncoder::new(&mut buf);
447        image::ImageEncoder::write_image(encoder, img.as_raw(), 1, 1, image::ColorType::Rgba8)
448            .unwrap();
449
450        use base64::Engine;
451        let b64 = base64::engine::general_purpose::STANDARD.encode(&buf);
452        let data_uri = format!("data:image/png;base64,{}", b64);
453
454        let loaded = load_image(&data_uri).unwrap();
455        assert_eq!(loaded.width_px, 1);
456        assert_eq!(loaded.height_px, 1);
457    }
458}
forme/image_loader.rs

forme/
image_loader.rs
