Skip to main content

img_squeeze/
processing.rs

1use crate::constants::{
2    DEFAULT_QUALITY, LIBDEFLATER_HIGH_LEVEL, LIBDEFLATER_LOW_LEVEL, MAX_FILE_SIZE,
3    MAX_IMAGE_DIMENSION, MAX_QUALITY, MIN_QUALITY, ZOPFLI_ITERATIONS,
4};
5use crate::error::{CompressionError, Result};
6use image::{DynamicImage, GenericImageView, ImageEncoder, ImageFormat, ImageReader};
7use indicatif::{ProgressBar, ProgressStyle};
8use oxipng::{Deflaters, InFile, Options, OutFile};
9use std::fs;
10use std::num::NonZeroU8;
11use std::path::{Path, PathBuf};
12
13#[derive(Debug, Clone)]
14pub struct CompressionOptions {
15    pub quality: u8,
16    pub width: Option<u32>,
17    pub height: Option<u32>,
18    pub format: Option<String>,
19}
20
21impl CompressionOptions {
22    pub fn new(
23        quality: Option<u8>,
24        width: Option<u32>,
25        height: Option<u32>,
26        format: Option<String>,
27    ) -> Result<Self> {
28        let quality = quality.unwrap_or(DEFAULT_QUALITY);
29        if !(MIN_QUALITY..=MAX_QUALITY).contains(&quality) {
30            return Err(CompressionError::InvalidQuality(quality));
31        }
32
33        Ok(Self {
34            quality,
35            width,
36            height,
37            format,
38        })
39    }
40}
41
42/// Validates that a file exists at the given path.
43///
44/// # Arguments
45/// * `path` - The path to check for existence
46///
47/// # Returns
48/// * `Ok(())` if the file exists
49/// * `Err(CompressionError::FileNotFound)` if the file does not exist
50///
51/// # Example
52/// ```
53/// use std::path::Path;
54/// use img_squeeze::validate_file_exists;
55///
56/// let result = validate_file_exists(Path::new("nonexistent.jpg"));
57/// assert!(result.is_err());
58/// ```
59pub fn validate_file_exists(path: &Path) -> Result<()> {
60    if !path.exists() {
61        return Err(CompressionError::FileNotFound(path.to_path_buf()));
62    }
63    Ok(())
64}
65
66/// Core image processing pipeline that handles the common workflow:
67/// load -> resize -> process -> save
68///
69/// # Arguments
70/// * `input_path` - Path to the input image file
71/// * `output_path` - Path where the processed image will be saved
72/// * `options` - Compression and processing options
73///
74/// # Returns
75/// * `Ok((original_size, compressed_size))` - Tuple of file sizes in bytes
76/// * `Err(CompressionError)` - If any processing step fails
77///
78/// # Security
79/// - Validates file existence and canonical paths to prevent directory traversal
80/// - Enforces maximum file size and image dimension limits
81/// - Uses secure temporary file handling
82pub fn process_image_pipeline(
83    input_path: &Path,
84    output_path: &Path,
85    options: &CompressionOptions,
86) -> Result<(u64, u64)> {
87    // Load and validate image
88    let (mut img, original_size) = load_image_with_metadata(input_path)?;
89
90    // Resize if needed
91    resize_image(&mut img, options);
92
93    // Process and save
94    let compressed_size = process_and_save_image(&img, output_path, options)?;
95
96    Ok((original_size, compressed_size))
97}
98
99/// Loads an image file and returns it along with file metadata.
100///
101/// # Arguments
102/// * `input_path` - Path to the image file to load
103///
104/// # Returns
105/// * `Ok((image, file_size))` - The loaded image and its file size in bytes
106/// * `Err(CompressionError)` - If loading fails or security limits are exceeded
107///
108/// # Security Features
109/// - Validates file existence and canonical paths to prevent directory traversal
110/// - Enforces maximum file size limit to prevent DoS attacks
111/// - Validates image dimensions to prevent memory exhaustion
112/// - Checks file size before attempting to load the image
113pub fn load_image_with_metadata(input_path: &Path) -> Result<(DynamicImage, u64)> {
114    validate_file_exists(input_path)?;
115
116    // Check for unsupported input formats and provide helpful guidance
117    if let Some(ext) = input_path.extension().and_then(|s| s.to_str()) {
118        match ext.to_ascii_lowercase().as_str() {
119            "heic" | "heif" => {
120                return Err(CompressionError::UnsupportedFormat(
121                    "HEIC/HEIF format is not yet supported in this version. Use AVIF for modern compression with similar quality and efficiency".to_string()
122                ));
123            }
124            "jxl" | "jpegxl" => {
125                return Err(CompressionError::UnsupportedFormat(
126                    "JPEG XL format is not yet supported in this version. Use AVIF for modern compression with similar quality and efficiency".to_string()
127                ));
128            }
129            _ => {} // Continue with supported formats
130        }
131    }
132
133    // Security: Validate path to prevent directory traversal attacks
134    let canonical_path = input_path
135        .canonicalize()
136        .map_err(|_| CompressionError::FileNotFound(input_path.to_path_buf()))?;
137
138    // Check file size before loading to prevent DoS attacks
139    let file_size = fs::metadata(&canonical_path)?.len();
140    if file_size > MAX_FILE_SIZE {
141        return Err(CompressionError::FileTooLarge(file_size, MAX_FILE_SIZE));
142    }
143
144    let img = ImageReader::open(&canonical_path)?.decode()?;
145
146    // Security: Validate image dimensions to prevent DoS attacks
147    let (width, height) = img.dimensions();
148    if width > MAX_IMAGE_DIMENSION || height > MAX_IMAGE_DIMENSION {
149        return Err(CompressionError::InvalidDimensions(
150            width,
151            height,
152            MAX_IMAGE_DIMENSION,
153        ));
154    }
155
156    Ok((img, file_size))
157}
158
159pub fn resize_image(img: &mut DynamicImage, options: &CompressionOptions) {
160    if let Some(w) = options.width.filter(|&w| w > 0 && w != img.width()) {
161        println!("🔄 Resizing width...");
162        *img = img.resize_exact(w, img.height(), image::imageops::FilterType::Lanczos3);
163        println!("✅ Resized to width: {}", w);
164    }
165
166    if let Some(h) = options.height.filter(|&h| h > 0 && h != img.height()) {
167        println!("🔄 Resizing height...");
168        *img = img.resize_exact(img.width(), h, image::imageops::FilterType::Lanczos3);
169        println!("✅ Resized to height: {}", h);
170    }
171}
172
173pub fn process_and_save_image(
174    img: &DynamicImage,
175    output_path: &Path,
176    options: &CompressionOptions,
177) -> Result<u64> {
178    let output_buf = output_path.to_path_buf();
179    let output_format = determine_output_format(output_path, &options.format)?;
180    save_image(img, &output_buf, output_format, options)?;
181
182    let compressed_size = fs::metadata(output_path)?.len();
183    Ok(compressed_size)
184}
185
186pub fn compress_image(input: PathBuf, output: PathBuf, options: CompressionOptions) -> Result<()> {
187    println!("🗜️  Compressing image: {:?}", input);
188    println!("📁 Output: {:?}", output);
189
190    let pb = ProgressBar::new_spinner();
191    pb.set_style(
192        ProgressStyle::default_spinner()
193            .template("{spinner:.green} {msg}")
194            .unwrap(),
195    );
196    pb.set_message("Loading image...");
197
198    let (mut img, original_size) = load_image_with_metadata(&input)?;
199    pb.finish_with_message("✅ Image loaded");
200
201    println!(
202        "📊 Original size: {} bytes ({}x{})",
203        original_size,
204        img.width(),
205        img.height()
206    );
207
208    // Resize if needed
209    resize_image(&mut img, &options);
210
211    pb.set_message("Saving compressed image...");
212    let compressed_size = process_and_save_image(&img, &output, &options)?;
213    pb.finish_with_message("✅ Compression complete");
214    let compression_ratio =
215        ((original_size as f64 - compressed_size as f64) / original_size as f64) * 100.0;
216
217    println!("📈 Compressed size: {} bytes", compressed_size);
218    println!("🎯 Compression ratio: {:.1}%", compression_ratio);
219
220    if compression_ratio > 0.0 {
221        println!(
222            "✅ Successfully reduced file size by {:.1}%",
223            compression_ratio
224        );
225    } else {
226        println!("⚠️  File size increased by {:.1}%", compression_ratio.abs());
227    }
228
229    Ok(())
230}
231
232pub fn determine_output_format(output: &Path, format: &Option<String>) -> Result<ImageFormat> {
233    if let Some(fmt) = format {
234        match fmt.to_lowercase().as_str() {
235            "jpeg" | "jpg" => Ok(ImageFormat::Jpeg),
236            "png" => Ok(ImageFormat::Png),
237            "webp" => Ok(ImageFormat::WebP),
238            "avif" => Ok(ImageFormat::Avif),
239            "heic" | "heif" => Err(CompressionError::UnsupportedFormat(
240                format!("{} format is not yet supported in this version. Use AVIF for modern compression", fmt)
241            )),
242            "jxl" | "jpegxl" => Err(CompressionError::UnsupportedFormat(
243                format!("{} format is not yet supported in this version. Use AVIF for modern compression", fmt)
244            )),
245            _ => Err(CompressionError::UnsupportedFormat(fmt.clone())),
246        }
247    } else if let Some(ext) = output.extension().and_then(|ext| ext.to_str()) {
248        let ext = ext.to_ascii_lowercase();
249        match ext.as_str() {
250            "jpg" | "jpeg" => Ok(ImageFormat::Jpeg),
251            "png" => Ok(ImageFormat::Png),
252            "webp" => Ok(ImageFormat::WebP),
253            "avif" => Ok(ImageFormat::Avif),
254            "heic" | "heif" => Err(CompressionError::UnsupportedFormat(
255                format!("{} format is not yet supported in this version. Use AVIF for modern compression", ext)
256            )),
257            "jxl" | "jpegxl" => Err(CompressionError::UnsupportedFormat(
258                format!("{} format is not yet supported in this version. Use AVIF for modern compression", ext)
259            )),
260            _ => Ok(ImageFormat::Jpeg),
261        }
262    } else {
263        Ok(ImageFormat::Jpeg)
264    }
265}
266
267pub fn save_image(
268    img: &DynamicImage,
269    output: &PathBuf,
270    format: ImageFormat,
271    options: &CompressionOptions,
272) -> Result<()> {
273    if let Some(parent) = output.parent() {
274        fs::create_dir_all(parent)
275            .map_err(|_| CompressionError::DirectoryCreationFailed(parent.to_path_buf()))?;
276    }
277
278    match format {
279        ImageFormat::Jpeg => {
280            img.save_with_format(output, image::ImageFormat::Jpeg)?;
281        }
282        ImageFormat::Png => {
283            // 使用 oxipng 进行 PNG 优化
284            let (_width, _height) = img.dimensions();
285
286            // Performance: Use secure temp file with proper cleanup
287            let temp_path = output.with_extension("temp.png");
288            img.save_with_format(&temp_path, image::ImageFormat::Png)?;
289
290            // Performance: Ensure cleanup on any error using RAII pattern
291            struct TempFileGuard(PathBuf);
292            impl Drop for TempFileGuard {
293                fn drop(&mut self) {
294                    let _ = fs::remove_file(&self.0);
295                }
296            }
297            let _guard = TempFileGuard(temp_path.clone());
298
299            // 配置 oxipng 选项
300            let mut oxipng_options = Options::from_preset(4); // 使用预设 4 (最高压缩)
301            oxipng_options.force = true; // 强制覆盖
302
303            // 根据质量设置调整压缩级别
304            if options.quality >= 90 {
305                oxipng_options.deflate = Deflaters::Zopfli {
306                    iterations: NonZeroU8::new(ZOPFLI_ITERATIONS).unwrap(),
307                };
308            } else if options.quality >= 70 {
309                oxipng_options.deflate = Deflaters::Libdeflater {
310                    compression: LIBDEFLATER_HIGH_LEVEL,
311                };
312            } else {
313                oxipng_options.deflate = Deflaters::Libdeflater {
314                    compression: LIBDEFLATER_LOW_LEVEL,
315                };
316            }
317
318            // 使用 oxipng 优化文件
319            let input = InFile::Path(temp_path.clone());
320            let out = OutFile::Path {
321                path: Some(output.clone()),
322                preserve_attrs: false,
323            };
324            oxipng::optimize(&input, &out, &oxipng_options)
325                .map_err(|e| CompressionError::PngOptimization(e.to_string()))?;
326
327            // Temp file automatically cleaned up by guard
328        }
329        ImageFormat::WebP => {
330            img.save_with_format(output, image::ImageFormat::WebP)?;
331        }
332        ImageFormat::Avif => {
333            // Honor quality and enable parallel encoding (when "image/rayon" is enabled).
334            use image::codecs::avif::AvifEncoder;
335            let rgba = img.to_rgba8();
336            let (w, h) = (rgba.width(), rgba.height());
337            let mut file = std::fs::File::create(output)?;
338            // Heuristic speed; consider surfacing as an option later.
339            let speed: u8 = 6;
340            let enc = AvifEncoder::new_with_speed_quality(&mut file, speed, options.quality);
341            // If you later thread this via CLI, call: enc = enc.with_num_threads(Some(n));
342            enc.write_image(
343                rgba.as_raw(),
344                w,
345                h,
346                image::ExtendedColorType::Rgba8
347            )?;
348        }
349        _ => {
350            return Err(CompressionError::UnsupportedFormat(format!("{:?}", format)));
351        }
352    }
353
354    Ok(())
355}
356
357#[cfg(test)]
358mod tests {
359    use super::*;
360
361    #[test]
362    fn test_compression_options_creation() {
363        let options =
364            CompressionOptions::new(Some(85), Some(800), Some(600), Some("webp".to_string()))
365                .unwrap();
366        assert_eq!(options.quality, 85);
367        assert_eq!(options.width, Some(800));
368        assert_eq!(options.height, Some(600));
369        assert_eq!(options.format, Some("webp".to_string()));
370    }
371
372    #[test]
373    fn test_compression_options_default() {
374        let options = CompressionOptions::new(None, None, None, None).unwrap();
375        assert_eq!(options.quality, 80);
376        assert_eq!(options.width, None);
377        assert_eq!(options.height, None);
378        assert_eq!(options.format, None);
379    }
380
381    #[test]
382    fn test_compression_options_invalid_quality() {
383        let result = CompressionOptions::new(Some(0), None, None, None);
384        assert!(matches!(result, Err(CompressionError::InvalidQuality(0))));
385
386        let result = CompressionOptions::new(Some(101), None, None, None);
387        assert!(matches!(result, Err(CompressionError::InvalidQuality(101))));
388    }
389
390    #[test]
391    fn test_determine_output_format() {
392        let path = Path::new("test.jpg");
393        let format = determine_output_format(path, &None).unwrap();
394        assert_eq!(format, ImageFormat::Jpeg);
395
396        let path = Path::new("test.png");
397        let format = determine_output_format(path, &None).unwrap();
398        assert_eq!(format, ImageFormat::Png);
399
400        let path = Path::new("test.webp");
401        let format = determine_output_format(path, &None).unwrap();
402        assert_eq!(format, ImageFormat::WebP);
403
404        let path = Path::new("test.avif");
405        let format = determine_output_format(path, &None).unwrap();
406        assert_eq!(format, ImageFormat::Avif);
407
408        let path = Path::new("test.unknown");
409        let format = determine_output_format(path, &None).unwrap();
410        assert_eq!(format, ImageFormat::Jpeg);
411    }
412
413    #[test]
414    fn test_determine_output_format_with_override() {
415        let path = Path::new("test.jpg");
416        let format = determine_output_format(path, &Some("png".to_string())).unwrap();
417        assert_eq!(format, ImageFormat::Png);
418
419        let path = Path::new("test.png");
420        let format = determine_output_format(path, &Some("avif".to_string())).unwrap();
421        assert_eq!(format, ImageFormat::Avif);
422    }
423
424    #[test]
425    fn test_determine_output_format_unsupported() {
426        let path = Path::new("test.jpg");
427        let result = determine_output_format(path, &Some("unsupported".to_string()));
428        assert!(matches!(
429            result,
430            Err(CompressionError::UnsupportedFormat(_))
431        ));
432
433        // Test HEIC/HEIF recognition with helpful error message
434        let result = determine_output_format(path, &Some("heic".to_string()));
435        assert!(matches!(
436            result,
437            Err(CompressionError::UnsupportedFormat(_))
438        ));
439        if let Err(CompressionError::UnsupportedFormat(msg)) = result {
440            assert!(msg.contains("not yet supported"));
441            assert!(msg.contains("AVIF"));
442        }
443
444        // Test JPEG XL recognition with helpful error message  
445        let result = determine_output_format(path, &Some("jxl".to_string()));
446        assert!(matches!(
447            result,
448            Err(CompressionError::UnsupportedFormat(_))
449        ));
450        if let Err(CompressionError::UnsupportedFormat(msg)) = result {
451            assert!(msg.contains("not yet supported"));
452            assert!(msg.contains("AVIF"));
453        }
454
455        // Alias: jpegxl should behave like jxl
456        let result = determine_output_format(path, &Some("jpegxl".to_string()));
457        assert!(matches!(result, Err(CompressionError::UnsupportedFormat(_))));
458        if let Err(CompressionError::UnsupportedFormat(msg)) = result {
459            assert!(msg.contains("not yet supported"));
460            assert!(msg.contains("AVIF"));
461        }
462    }
463
464    #[test]
465    fn test_determine_output_format_case_insensitive_exts() {
466        // After normalizing extension cases, uppercase should work.
467        let path = Path::new("IMAGE.AVIF");
468        let format = determine_output_format(path, &None).unwrap();
469        assert_eq!(format, ImageFormat::Avif);
470
471        // Unsupported alias via extension
472        let path = Path::new("photo.jpegxl");
473        let result = determine_output_format(path, &None);
474        assert!(matches!(result, Err(CompressionError::UnsupportedFormat(_))));
475    }
476
477    #[test]
478    fn test_resize_image_dimensions() {
479        let mut img = DynamicImage::new_rgb8(2000, 1500);
480        let options = CompressionOptions::new(Some(80), Some(1000), None, None).unwrap();
481
482        resize_image(&mut img, &options);
483
484        assert_eq!(img.dimensions(), (1000, 1500));
485    }
486
487    #[test]
488    fn test_resize_image_height_only() {
489        let mut img = DynamicImage::new_rgb8(2000, 1500);
490        let options = CompressionOptions::new(Some(80), None, Some(750), None).unwrap();
491
492        resize_image(&mut img, &options);
493
494        assert_eq!(img.dimensions(), (2000, 750));
495    }
496
497    #[test]
498    fn test_resize_image_both_dimensions() {
499        let mut img = DynamicImage::new_rgb8(2000, 1500);
500        let options = CompressionOptions::new(Some(80), Some(800), Some(600), None).unwrap();
501
502        resize_image(&mut img, &options);
503
504        assert_eq!(img.dimensions(), (800, 600));
505    }
506
507    #[test]
508    fn test_resize_image_no_dimensions() {
509        let mut img = DynamicImage::new_rgb8(2000, 1500);
510        let options = CompressionOptions::new(Some(80), None, None, None).unwrap();
511
512        resize_image(&mut img, &options);
513
514        assert_eq!(img.dimensions(), (2000, 1500));
515    }
516
517    #[test]
518    fn test_resize_image_same_dimensions() {
519        let mut img = DynamicImage::new_rgb8(2000, 1500);
520        let options = CompressionOptions::new(Some(80), Some(2000), Some(1500), None).unwrap();
521
522        resize_image(&mut img, &options);
523
524        assert_eq!(img.dimensions(), (2000, 1500));
525    }
526
527    #[test]
528    fn test_load_image_with_metadata_not_found() {
529        let path = Path::new("nonexistent.jpg");
530        let result = load_image_with_metadata(path);
531        assert!(matches!(result, Err(CompressionError::FileNotFound(_))));
532    }
533}