Edge-case covered: the end of the file

src/main.rs

@@ -13,54 +13,20 @@ THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR I
 use std::io::{Read, Write};
 use std::path::Path;
 
-fn main() {
+/// extracts png images from file that path points to, saving every image to destination directory.
-    let args: Vec<String> = std::env::args().collect();
+/// If an error occurs - returns immediately.
-    if args.len() < 3 {
+fn rip_png(path: &Path, destination: &Path) {
-        println!("USAGE: pngrip [DESTINATION] [FILES]...");
-        std::process::exit(0);
-    }
-
-    // handle DESTINATION argument
-    let mut destination = Path::new(&args[1]);
-    match destination.exists() {
-        true => {
-            if !destination.is_dir() {
-                // destination exists, but is not a directory
-                destination = Path::new(".");
-                eprintln!("[ERROR] Provided destination path is not a directory ! Saving to current directory...");
-            }
-        }
-        false => {
-            // destination does not exist, create it
-            match std::fs::create_dir_all(&destination) {
-                Ok(_) => {}
-                Err(e) => {
-                    destination = Path::new(".");
-                    eprintln!("[ERROR] Could not create destination directory: {}. Saving to current directory...", e);
-                }
-            }
-        }
-    }
-
-
-    // go through all files and try to rip all PNGs
     let png_identifier: [u8; 8] = [137, 80, 78, 71, 13, 10, 26, 10];
     let iend_identifier: [u8; 4] = [73, 69, 78, 68];
-    for file_to_check in &args[2..] {
-        let path = Path::new(file_to_check);
-        if !path.exists() {
-            eprintln!("[ERROR] \"{}\" does not exist", file_to_check);
-            continue;
-        }
 
     let filename;
     match path.file_name() {
         Some(name) => {
-                filename = name.to_string_lossy();
+            filename = String::from(name.to_string_lossy());
         }
         None => {
-                eprintln!("[ERROR] Could not get filename from \"{}\"", file_to_check);
+            eprintln!("[ERROR] Could not get filename from \"{}\"", path.display());
-                continue;
+            return;
         }
     }
 
@@ -74,7 +40,7 @@ fn main() {
 
         Err(e) => {
             eprintln!("[ERROR] On opening \"{}\": {}", filename, e);
-                continue;
+            return;
         }
     }
 
@@ -83,7 +49,7 @@ fn main() {
         Ok(_) => {}
         Err(e) => {
             eprintln!("[ERROR] On reading \"{}\": {}", filename, e);
-                continue;
+            return;
         }
     }
 
@@ -92,9 +58,9 @@ fn main() {
     let mut start_pos: usize = 0;
     let mut last_pos: usize = 0;
     let mut end_pos: usize = 0;
-        for i in 0..file_bytes.len()-8 {
+    for i in 0..file_bytes.len()-iend_identifier.len() {
-
+        if i < file_bytes.len() - png_identifier.len() - iend_identifier.len()
-            if file_bytes[i..i + png_identifier.len()] == png_identifier {
+            && file_bytes[i..i + png_identifier.len()] == png_identifier {
             start_pos = i;
         }
 
@@ -115,7 +81,7 @@ fn main() {
                 }
                 Err(e) => {
                     eprintln!("[ERROR] On creating \"{}\": {}", &ripped_image_filename, e);
-                        continue;
+                    return;
                 }
             }
 
@@ -123,14 +89,55 @@ fn main() {
                 Ok(_) => {}
                 Err(e) => {
                     eprintln!("[ERROR] On writing to \"{}\": {}", ripped_image_filename, e);
-                        continue;
+                    return;
+                }
+            }
         }
     }
+
+
+    println!("[INFO] Ripped {} images from \"{}\" in total", image_count, filename);
 }
+
+fn main() {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() < 3 {
+        println!("USAGE: pngrip [DESTINATION] [FILES]...");
+        std::process::exit(0);
     }
 
+    // handle DESTINATION argument
+    let mut destination = Path::new(&args[1]);
+    match destination.exists() {
+        true => {
+            if !destination.is_dir() {
+                // destination exists, but is not a directory
+                destination = Path::new(".");
+                eprintln!("[ERROR] Provided destination path is not a directory ! Saving to current directory...");
+            }
+        }
+        false => {
+            // destination does not exist, create it
+            match std::fs::create_dir_all(&destination) {
+                Ok(_) => {}
+                Err(e) => {
+                    destination = Path::new(".");
+                    eprintln!("[ERROR] Could not create destination directory: {}. Saving to current directory...", e);
+                }
+            }
+        }
+    }
+
+
+    // go through all files and try to rip all PNGs
+    for file_to_check in &args[2..] {
+        let path = Path::new(file_to_check);
+        if !path.exists() {
+            eprintln!("[ERROR] \"{}\" does not exist", file_to_check);
+            continue;
+        }
 
-        println!("[INFO] Ripped {} images in total", image_count);
+        rip_png(path, destination);
     }
 
 }