From 6258fa28e5c3ce461bc8489d30d7a6247517b6e4 Mon Sep 17 00:00:00 2001
From: Vladislav Kalinin <unkorunk@gmail.com>
Date: Sun, 29 Oct 2023 15:26:57 +0100
Subject: [PATCH] Add the ability to customize behavior for filtering headers

---
 src/async_impl/client.rs |  8 +++--
 src/redirect.rs          | 75 +++++++++++++++++++++++++++++++---------
 2 files changed, 64 insertions(+), 19 deletions(-)

diff --git a/src/async_impl/client.rs b/src/async_impl/client.rs
index b6b515b89..7f6065b6a 100644
--- a/src/async_impl/client.rs
+++ b/src/async_impl/client.rs
@@ -38,7 +38,7 @@ use crate::dns::trust_dns::TrustDnsResolver;
 use crate::dns::{gai::GaiResolver, DnsResolverWithOverrides, DynResolver, Resolve};
 use crate::error;
 use crate::into_url::{expect_uri, try_uri};
-use crate::redirect::{self, remove_sensitive_headers};
+use crate::redirect::{self};
 #[cfg(feature = "__tls")]
 use crate::tls::{self, TlsBackend};
 #[cfg(feature = "__tls")]
@@ -2378,7 +2378,11 @@ impl Future for PendingRequest {
                             let mut headers =
                                 std::mem::replace(self.as_mut().headers(), HeaderMap::new());
 
-                            remove_sensitive_headers(&mut headers, &self.url, &self.urls);
+                            self.client.redirect_policy.handle_sensitive_headers(
+                                &mut headers,
+                                &self.url,
+                                &self.urls,
+                            );
                             let uri = expect_uri(&self.url);
                             let body = match self.body {
                                 Some(Some(ref body)) => Body::reusable(body.clone()),
diff --git a/src/redirect.rs b/src/redirect.rs
index eabaea37b..1425a4c86 100644
--- a/src/redirect.rs
+++ b/src/redirect.rs
@@ -23,6 +23,41 @@ use crate::Url;
 /// - `custom` can be used to create a customized policy.
 pub struct Policy {
     inner: PolicyKind,
+    filter: Box<dyn Filter + Send + Sync>,
+}
+
+/// A trait that controls the handling of sensitive headers.
+pub trait Filter {
+    /// Handle sensitive headers.
+    fn handle_sensitive_headers(&self, headers: &mut HeaderMap, next: &Url, previous: &[Url]);
+}
+
+/// Default implementation for the handling of sensitive headers.
+///
+/// This implementation will exclude the authorization and cookie headers.
+#[derive(Debug)]
+pub struct DefaultFilter {}
+
+impl Default for DefaultFilter {
+    fn default() -> Self {
+        Self {}
+    }
+}
+
+impl Filter for DefaultFilter {
+    fn handle_sensitive_headers(&self, headers: &mut HeaderMap, next: &Url, previous: &[Url]) {
+        if let Some(previous) = previous.last() {
+            let cross_host = next.host_str() != previous.host_str()
+                || next.port_or_known_default() != previous.port_or_known_default();
+            if cross_host {
+                headers.remove(AUTHORIZATION);
+                headers.remove(COOKIE);
+                headers.remove("cookie2");
+                headers.remove(PROXY_AUTHORIZATION);
+                headers.remove(WWW_AUTHENTICATE);
+            }
+        }
+    }
 }
 
 /// A type that holds information on the next request and previous requests
@@ -47,6 +82,7 @@ impl Policy {
     pub fn limited(max: usize) -> Self {
         Self {
             inner: PolicyKind::Limit(max),
+            filter: Box::new(DefaultFilter::default()),
         }
     }
 
@@ -54,6 +90,7 @@ impl Policy {
     pub fn none() -> Self {
         Self {
             inner: PolicyKind::None,
+            filter: Box::new(DefaultFilter::default()),
         }
     }
 
@@ -101,6 +138,7 @@ impl Policy {
     {
         Self {
             inner: PolicyKind::Custom(Box::new(policy)),
+            filter: Box::new(DefaultFilter::default()),
         }
     }
 
@@ -150,6 +188,21 @@ impl Policy {
     pub(crate) fn is_default(&self) -> bool {
         matches!(self.inner, PolicyKind::Limit(10))
     }
+
+    /// Set filter for the handling of sensitive headers.
+    pub fn set_filter(&mut self, filter: Box<dyn Filter + Send + Sync>) {
+        self.filter = filter;
+    }
+
+    pub(crate) fn handle_sensitive_headers(
+        &self,
+        headers: &mut HeaderMap,
+        next: &Url,
+        previous: &[Url],
+    ) {
+        self.filter
+            .handle_sensitive_headers(headers, next, previous);
+    }
 }
 
 impl Default for Policy {
@@ -231,20 +284,6 @@ pub(crate) enum ActionKind {
     Error(Box<dyn StdError + Send + Sync>),
 }
 
-pub(crate) fn remove_sensitive_headers(headers: &mut HeaderMap, next: &Url, previous: &[Url]) {
-    if let Some(previous) = previous.last() {
-        let cross_host = next.host_str() != previous.host_str()
-            || next.port_or_known_default() != previous.port_or_known_default();
-        if cross_host {
-            headers.remove(AUTHORIZATION);
-            headers.remove(COOKIE);
-            headers.remove("cookie2");
-            headers.remove(PROXY_AUTHORIZATION);
-            headers.remove(WWW_AUTHENTICATE);
-        }
-    }
-}
-
 #[derive(Debug)]
 struct TooManyRedirects;
 
@@ -313,7 +352,7 @@ fn test_redirect_policy_custom() {
 }
 
 #[test]
-fn test_remove_sensitive_headers() {
+fn test_handle_sensitive_headers() {
     use hyper::header::{HeaderValue, ACCEPT, AUTHORIZATION, COOKIE};
 
     let mut headers = HeaderMap::new();
@@ -325,13 +364,15 @@ fn test_remove_sensitive_headers() {
     let mut prev = vec![Url::parse("http://initial-domain.com/new_path").unwrap()];
     let mut filtered_headers = headers.clone();
 
-    remove_sensitive_headers(&mut headers, &next, &prev);
+    let filter = DefaultFilter::default();
+
+    filter.handle_sensitive_headers(&mut headers, &next, &prev);
     assert_eq!(headers, filtered_headers);
 
     prev.push(Url::parse("http://new-domain.com/path").unwrap());
     filtered_headers.remove(AUTHORIZATION);
     filtered_headers.remove(COOKIE);
 
-    remove_sensitive_headers(&mut headers, &next, &prev);
+    filter.handle_sensitive_headers(&mut headers, &next, &prev);
     assert_eq!(headers, filtered_headers);
 }