From 0ce390f4789e7b62bb4fe27f4d7505f31058c498 Mon Sep 17 00:00:00 2001
From: Ben Richeson <36977340+Benricheson101@users.noreply.github.com>
Date: Tue, 24 Dec 2024 18:03:49 -0500
Subject: [PATCH] add config option for removing query parameters (#1)

---
 config.example.toml |  3 +++
 src/pass.rs         | 57 ++++++++++++++++++++++++++++++++-------------
 tests/passes.rs     | 10 +++++---
 3 files changed, 51 insertions(+), 19 deletions(-)

diff --git a/config.example.toml b/config.example.toml
index 4295a22..d2c2faf 100644
--- a/config.example.toml
+++ b/config.example.toml
@@ -20,6 +20,9 @@ label = "Tweet"
 regex = "https://(?:x|twitter)\\.com"
 # The stem to replace the matched area with.
 stem = "https://vxtwitter.com"
+# The query params to keep in the URL -- empty ([]) to remove query string entirely or
+# omitted to keep entire query string
+keep_query = []
 
 [[pass]]
 label = "Instagram Post"
diff --git a/src/pass.rs b/src/pass.rs
index aedec7b..77a5c0b 100644
--- a/src/pass.rs
+++ b/src/pass.rs
@@ -1,4 +1,4 @@
-use std::fmt::Write;
+use std::{collections::HashMap, fmt::Write};
 
 use regex::Regex;
 use serde::{Deserialize, Deserializer};
@@ -9,6 +9,7 @@ pub struct Pass {
     #[serde(deserialize_with = "pass_regex")]
     pub regex: Regex,
     pub stem: String,
+    pub keep_query: Option<Vec<String>>,
 }
 
 /// An enum representing the spoiler tags on a link.
@@ -25,7 +26,10 @@ pub enum SpoilerTags {
 }
 
 impl Pass {
-    pub fn extract<'a>(&'a self, content: &'a str) -> impl Iterator<Item = (&'a str, SpoilerTags)> {
+    pub fn extract<'a>(
+        &'a self,
+        content: &'a str,
+    ) -> impl Iterator<Item = (&'a str, &'a str, SpoilerTags)> {
         self.regex.captures_iter(content).map(|capture| {
             let (_, [sp_open, path, sp_close]) = capture.extract();
             let spoiler_marker = match (!sp_open.is_empty(), !sp_close.is_empty()) {
@@ -34,28 +38,36 @@ impl Pass {
                 _ => SpoilerTags::Mismatched,
             };
 
-            (path, spoiler_marker)
+            let (path, query) = path.split_once('?').unwrap_or((path, ""));
+
+            (path, query, spoiler_marker)
         })
     }
 
     pub fn apply<'a>(&'a self, content: &'a str) -> Option<String> {
         let Self { label, stem, .. } = self;
 
-        let out = self
-            .extract(content)
-            .fold(String::new(), |mut out, (path, spoiler_tags)| {
-                let spoil = spoiler_tags != SpoilerTags::None;
+        let out =
+            self.extract(content)
+                .fold(String::new(), |mut out, (path, query, spoiler_tags)| {
+                    let spoil = spoiler_tags != SpoilerTags::None;
+
+                    let query_string = match &self.keep_query {
+                        None => format!("?{query}"),
+                        Some(keep) if !keep.is_empty() => filter_query(query, keep),
+                        _ => String::new(),
+                    };
 
-                if spoil {
-                    let _ = write!(&mut out, "||");
-                }
-                let _ = write!(&mut out, "[`{label}`]({stem}{path}) ");
-                if spoil {
-                    let _ = write!(&mut out, "|| ");
-                }
+                    if spoil {
+                        let _ = write!(&mut out, "||");
+                    }
+                    let _ = write!(&mut out, "[`{label}`]({stem}{path}{query_string}) ");
+                    if spoil {
+                        let _ = write!(&mut out, "|| ");
+                    }
 
-                out
-            });
+                    out
+                });
 
         (!out.is_empty()).then_some(out)
     }
@@ -81,3 +93,16 @@ fn pass_regex<'de, D: Deserializer<'de>>(de: D) -> Result<Regex, D::Error> {
     Regex::new(&["(?:^|\\s)", "(\\|\\||)", &core, "(/\\S+)", "(\\s?\\|\\||)"].concat())
         .map_err(D::Error::custom)
 }
+
+/// Removes all query parameters from a query string except those in the provided list
+fn filter_query(qs: &str, keep: &Vec<String>) -> String {
+    let query_map: HashMap<_, _> = qs.split('&').filter_map(|p| p.split_once('=')).collect();
+
+    let params = query_map
+        .iter()
+        .filter_map(|(k, v)| keep.contains(&k.to_string()).then(|| format!("{k}={v}")))
+        .collect::<Vec<_>>()
+        .join("&");
+
+    format!("?{params}")
+}
diff --git a/tests/passes.rs b/tests/passes.rs
index 2f48c6f..b151418 100644
--- a/tests/passes.rs
+++ b/tests/passes.rs
@@ -8,7 +8,7 @@ fn standard_passes() {
         pass.extract(
             "
                 These are just some random test urls.
-                - https://x.com/rustbeltenjoyer/status/1776056709737320578
+                - https://x.com/rustbeltenjoyer/status/1776056709737320578?s=46&t=owouwu
                 - ||https://www.instagram.com/p/C5W2QwZrt-Z/ ||
                 - https://www.tiktok.com/t/ZPRTX3AwH/
             ",
@@ -21,14 +21,18 @@ fn standard_passes() {
         extracted.next(),
         Some((
             "/rustbeltenjoyer/status/1776056709737320578",
+            "s=46&t=owouwu",
             SpoilerTags::None
         ))
     );
     assert_eq!(
         extracted.next(),
-        Some(("/p/C5W2QwZrt-Z/", SpoilerTags::Spoiler))
+        Some(("/p/C5W2QwZrt-Z/", "", SpoilerTags::Spoiler))
+    );
+    assert_eq!(
+        extracted.next(),
+        Some(("/t/ZPRTX3AwH/", "", SpoilerTags::None))
     );
-    assert_eq!(extracted.next(), Some(("/t/ZPRTX3AwH/", SpoilerTags::None)));
 
     assert!(extracted.next().is_none());
 }