From 653f79163850f39b0b27c8b887ed4f3ca08ba718 Mon Sep 17 00:00:00 2001
From: llde <lorenzofersteam@live.it>
Date: Fri, 10 Dec 2021 17:56:44 +0100
Subject: [PATCH] Create new configuration option for redirect preserving
 authorization header in Agent. Handle new option in Unit

---
 src/agent.rs | 21 +++++++++++++++++++++
 src/unit.rs  | 33 +++++++++++++++------------------
 2 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/src/agent.rs b/src/agent.rs
index eb03c29..16c27db 100644
--- a/src/agent.rs
+++ b/src/agent.rs
@@ -16,6 +16,16 @@ use {
     cookie_store::CookieStore,
 };
 
+/// Specify the strategy for propagation of authorization headers during Redirects
+///
+/// Never is the default strategy and never send authorization headers in Redirects
+/// SameHost send the authorization header in Redirects only if the host of the Redirect is the same of the previous request, and both use the Https scheme
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum RedirectAuthHeaders {
+    Never,
+    SameHost,
+}
+
 /// Accumulates options towards building an [Agent].
 #[derive(Debug)]
 pub struct AgentBuilder {
@@ -38,6 +48,7 @@ pub(crate) struct AgentConfig {
     pub timeout_write: Option<Duration>,
     pub timeout: Option<Duration>,
     pub redirects: u32,
+    pub redirect_auth_headers: RedirectAuthHeaders,
     pub user_agent: String,
     pub tls_config: Arc<dyn TlsConnector>,
 }
@@ -221,6 +232,7 @@ impl AgentBuilder {
                 timeout_write: None,
                 timeout: None,
                 redirects: 5,
+                redirect_auth_headers: RedirectAuthHeaders::Never,
                 user_agent: format!("ureq/{}", env!("CARGO_PKG_VERSION")),
                 tls_config: crate::default_tls_config(),
             },
@@ -445,6 +457,15 @@ impl AgentBuilder {
         self
     }
 
+    /// Set the strategy for propagation of authorization headers in redirects.
+    ///
+    /// Defaults to RedirectAuthHeaders::Never.
+    ///
+    pub fn set_redirect_auth_headers(mut self, v: RedirectAuthHeaders) -> Self {
+        self.config.redirect_auth_headers = v;
+        self
+    }
+
     /// The user-agent header to associate with all requests from this agent by default.
     ///
     /// Defaults to `ureq/[VERSION]`. You can override the user-agent on an individual request by
diff --git a/src/unit.rs b/src/unit.rs
index 82229c4..7bddac2 100644
--- a/src/unit.rs
+++ b/src/unit.rs
@@ -9,6 +9,7 @@ use url::Url;
 #[cfg(feature = "cookies")]
 use cookie::Cookie;
 
+use crate::agent::RedirectAuthHeaders;
 use crate::body::{self, BodySize, Payload, SizedReader};
 use crate::error::{Error, ErrorKind};
 use crate::header;
@@ -353,6 +354,17 @@ fn connect_socket(unit: &Unit, hostname: &str, use_pooled: bool) -> Result<(Stre
     Ok((stream?, false))
 }
 
+fn can_propagate_authorization_on_redirect(unit: &Unit, previous: &[String]) -> bool {
+    if let RedirectAuthHeaders::SameHost = unit.agent.config.redirect_auth_headers {
+        let host_s = unit.url.host_str().unwrap();
+        let prev_url = Url::parse(&previous[0]).unwrap();
+        let prev_host = prev_url.host_str().unwrap();
+        host_s == prev_host && prev_url.scheme() == "https" && unit.url.scheme() == "https"
+    } else {
+        false
+    }
+}
+
 /// Send request line + headers (all up until the body).
 #[allow(clippy::write_with_newline)]
 fn send_prelude(unit: &Unit, stream: &mut Stream, previous: &[String]) -> io::Result<()> {
@@ -393,28 +405,13 @@ fn send_prelude(unit: &Unit, stream: &mut Stream, previous: &[String]) -> io::Re
     if !header::has_header(&unit.headers, "accept") {
         prelude.write_header("Accept", "*/*")?;
     }
-    //Preserve Authorization in Same Host redirections
-    let host_s = unit.url.host_str().unwrap();
+
     let preserve_auth = if !previous.is_empty() {
-        let prev_url = Url::parse(&previous[0]).unwrap();
-        if let Some(prev_host) = prev_url.host_str() {
-            if host_s == prev_host {
-                if prev_url.scheme() == unit.url.scheme() {
-                    let port = unit.url.port().unwrap_or(scheme_default);
-                    let prev_port = prev_url.port().unwrap_or(scheme_default);
-                    port == prev_port
-                } else {
-                    false
-                }
-            } else {
-                false
-            }
-        } else {
-            false
-        }
+        can_propagate_authorization_on_redirect(unit, previous)
     } else {
         true //Not in redirection
     };
+
     // other headers
     for header in &unit.headers {
         if preserve_auth || !header.is_name("Authorization") {