From e3138b0aceaa0b87a62b09d0700b5a3a08972192 Mon Sep 17 00:00:00 2001
From: Jacob Hoffman-Andrews <github@hoffman-andrews.com>
Date: Tue, 6 Oct 2020 00:12:26 -0700
Subject: [PATCH] Add proxy on agent. (#178)

---
 src/agent.rs   | 18 ++++++++++++++++++
 src/request.rs | 17 ++++++++++++++---
 src/stream.rs  |  9 +++++----
 3 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/src/agent.rs b/src/agent.rs
index a2d139a..18e9f4f 100644
--- a/src/agent.rs
+++ b/src/agent.rs
@@ -9,6 +9,7 @@ use url::Url;
 
 use crate::header::{self, Header};
 use crate::pool::ConnectionPool;
+use crate::proxy::Proxy;
 use crate::request::Request;
 use crate::resolve::ArcResolver;
 
@@ -55,6 +56,7 @@ pub struct Agent {
 pub(crate) struct AgentState {
     /// Reused connections between requests.
     pub(crate) pool: ConnectionPool,
+    pub(crate) proxy: Option<Proxy>,
     /// Cookies saved between requests.
     /// Invariant: All cookies must have a nonempty domain and path.
     #[cfg(feature = "cookie")]
@@ -220,6 +222,22 @@ impl Agent {
         self
     }
 
+    /// Set the proxy server to use for all connections from this Agent.
+    ///
+    /// Example:
+    /// ```
+    /// let proxy = ureq::Proxy::new("user:password@cool.proxy:9090").unwrap();
+    /// let agent = ureq::agent()
+    ///     .set_proxy(proxy)
+    ///     .build();
+    /// ```
+    pub fn set_proxy(&mut self, proxy: Proxy) -> &mut Agent {
+        let mut state = self.state.lock().unwrap();
+        state.proxy = Some(proxy);
+        drop(state);
+        self
+    }
+
     /// Gets a cookie in this agent by name. Cookies are available
     /// either by setting it in the agent, or by making requests
     /// that `Set-Cookie` in the agent.
diff --git a/src/request.rs b/src/request.rs
index 85bce0c..daee345 100644
--- a/src/request.rs
+++ b/src/request.rs
@@ -12,6 +12,7 @@ use crate::body::BodySize;
 use crate::body::{Payload, SizedReader};
 use crate::error::Error;
 use crate::header::{self, Header};
+use crate::proxy::Proxy;
 use crate::unit::{self, Unit};
 use crate::Response;
 
@@ -43,7 +44,7 @@ pub struct Request {
     pub(crate) timeout_write: u64,
     pub(crate) timeout: Option<time::Duration>,
     pub(crate) redirects: u32,
-    pub(crate) proxy: Option<crate::proxy::Proxy>,
+    pub(crate) proxy: Option<Proxy>,
     #[cfg(feature = "tls")]
     pub(crate) tls_config: Option<TLSClientConfig>,
     #[cfg(all(feature = "native-tls", not(feature = "tls")))]
@@ -559,12 +560,22 @@ impl Request {
     ///     .set_proxy(proxy)
     ///     .build();
     /// ```
-    pub fn set_proxy(&mut self, proxy: crate::proxy::Proxy) -> &mut Request {
+    pub fn set_proxy(&mut self, proxy: Proxy) -> &mut Request {
         self.proxy = Some(proxy);
         self
     }
 
-    /// Set the TLS client config to use for the connection.
+    pub(crate) fn proxy(&self) -> Option<Proxy> {
+        if let Some(proxy) = &self.proxy {
+            Some(proxy.clone())
+        } else if let Some(proxy) = &self.agent.lock().unwrap().proxy {
+            Some(proxy.clone())
+        } else {
+            None
+        }
+    }
+
+    /// Set the TLS client config to use for the connection. See [`ClientConfig`](https://docs.rs/rustls/latest/rustls/struct.ClientConfig.html).
     ///
     /// See [`ClientConfig`](https://docs.rs/rustls/latest/rustls/struct.ClientConfig.html).
     ///
diff --git a/src/stream.rs b/src/stream.rs
index 84ac940..f14a3ae 100644
--- a/src/stream.rs
+++ b/src/stream.rs
@@ -384,7 +384,8 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
     } else {
         unit.deadline
     };
-    let netloc = match unit.req.proxy {
+    let proxy: Option<Proxy> = unit.req.proxy();
+    let netloc = match proxy {
         Some(ref proxy) => format!("{}:{}", proxy.server, proxy.port),
         None => format!("{}:{}", hostname, port),
     };
@@ -399,7 +400,7 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
         return Err(Error::DnsFailed(format!("No ip address for {}", hostname)));
     }
 
-    let proto = if let Some(ref proxy) = unit.req.proxy {
+    let proto = if let Some(ref proxy) = proxy {
         Some(proxy.proto)
     } else {
         None
@@ -419,7 +420,7 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
         let stream = if Some(Proto::SOCKS5) == proto {
             connect_socks5(
                 &unit,
-                unit.req.proxy.to_owned().unwrap(),
+                proxy.clone().unwrap(),
                 deadline,
                 sock_addr,
                 hostname,
@@ -473,7 +474,7 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
     }
 
     if proto == Some(Proto::HTTPConnect) {
-        if let Some(ref proxy) = unit.req.proxy {
+        if let Some(ref proxy) = proxy {
             write!(stream, "{}", proxy.connect(hostname, port)).unwrap();
             stream.flush()?;