From c530bba261a56842d0094358a4c0111f77c0b3f2 Mon Sep 17 00:00:00 2001
From: Kornel <kornel@geekhood.net>
Date: Wed, 18 Jul 2018 04:55:21 +0100
Subject: [PATCH] --keep option (#121)

---
 src/headers.rs |  7 +++++--
 src/lib.rs     |  7 ++++++-
 src/main.rs    | 15 ++++++++++++++-
 tests/flags.rs |  2 +-
 4 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/src/headers.rs b/src/headers.rs
index d08c4a11..7979e1ba 100644
--- a/src/headers.rs
+++ b/src/headers.rs
@@ -1,3 +1,4 @@
+use std::collections::HashSet;
 use byteorder::{BigEndian, ReadBytesExt};
 use colors::{BitDepth, ColorType};
 use crc::crc32;
@@ -28,10 +29,12 @@ pub struct IhdrData {
 pub enum Headers {
     /// None
     None,
-    /// Some, with a list of 4-character chunk codes
-    Some(Vec<String>),
+    /// Remove specific chunks
+    Strip(Vec<String>),
     /// Headers that won't affect rendering (all but cHRM, gAMA, iCCP, sBIT, sRGB, bKGD, hIST, pHYs, sPLT)
     Safe,
+    /// Remove all non-critical chunks except these
+    Keep(HashSet<String>),
     /// All non-critical headers
     All,
 }
diff --git a/src/lib.rs b/src/lib.rs
index dd8a7068..07eb25e7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -808,7 +808,12 @@ fn perform_strip(png: &mut PngData, opts: &Options) {
     match opts.strip {
         // Strip headers
         Headers::None => (),
-        Headers::Some(ref hdrs) => for hdr in hdrs {
+        Headers::Keep(ref hdrs) => {
+            png.aux_headers.retain(|chunk, _| {
+                hdrs.contains(chunk)
+            });
+        },
+        Headers::Strip(ref hdrs) => for hdr in hdrs {
             png.aux_headers.remove(hdr);
         },
         Headers::Safe => {
diff --git a/src/main.rs b/src/main.rs
index 7a175ce4..960139fc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -94,6 +94,13 @@ fn main() {
             .takes_value(true)
             .value_name("mode")
             .conflicts_with("strip-safe"))
+        .arg(Arg::with_name("keep")
+            .help("Strip all optional metadata except objects in the comma-separated list")
+            .long("keep")
+            .takes_value(true)
+            .value_name("list")
+            .conflicts_with("strip")
+            .conflicts_with("strip-safe"))
         .arg(Arg::with_name("alpha")
             .help("Perform additional alpha optimizations")
             .short("a")
@@ -411,6 +418,12 @@ fn parse_opts_into_struct(matches: &ArgMatches) -> Result<(OutFile, Option<PathB
         opts.idat_recoding = false;
     }
 
+    if let Some(hdrs) = matches.value_of("keep") {
+        opts.strip = Headers::Keep(hdrs.split(',')
+            .map(|x| x.trim().to_owned())
+            .collect())
+    }
+
     if let Some(hdrs) = matches.value_of("strip") {
         let hdrs = hdrs.split(',')
             .map(|x| x.trim().to_owned())
@@ -433,7 +446,7 @@ fn parse_opts_into_struct(matches: &ArgMatches) -> Result<(OutFile, Option<PathB
                     return Err(format!("{} chunk is not allowed to be stripped", i));
                 }
             }
-            opts.strip = Headers::Some(hdrs);
+            opts.strip = Headers::Strip(hdrs);
         }
     }
 
diff --git a/tests/flags.rs b/tests/flags.rs
index 430d0abd..e99f3ffd 100644
--- a/tests/flags.rs
+++ b/tests/flags.rs
@@ -74,7 +74,7 @@ fn verbose_mode() {
 fn strip_headers_list() {
     let input = PathBuf::from("tests/files/strip_headers_list.png");
     let (output, mut opts) = get_opts(&input);
-    opts.strip = Headers::Some(vec!["iCCP".to_owned(), "tEXt".to_owned()]);
+    opts.strip = Headers::Strip(vec!["iCCP".to_owned(), "tEXt".to_owned()]);
 
     let png = PngData::new(&input, opts.fix_errors).unwrap();