Repository: algesten/ureq
Pull request #178 (2020-10-04 23:32:54)

Add proxy on agent.


Patch:
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()?;
 
diff --git a/src/unit.rs b/src/unit.rs
index bb080f0..3ec4794 100644
--- a/src/unit.rs
+++ b/src/unit.rs
@@ -145,7 +145,10 @@ pub(crate) fn connect(
 ) -> Result<Response, Error> {
     //
 
-    let host = req.get_host()?;
+    let host = unit
+        .url
+        .host_str()
+        .ok_or(Error::BadUrl("no host".to_string()))?;
     let url = &unit.url;
     let method = &unit.req.method;
     // open socket
@@ -361,6 +364,7 @@ fn send_prelude(unit: &Unit, stream: &mut Stream, redir: bool) -> io::Result<()>
     // finish
     write!(prelude, "\r\n")?;
 
+    debug!("writing prelude: {}", String::from_utf8_lossy(&prelude));
     // write all to the wire
     stream.write_all(&prelude[..])?;
 


Test patch:
diff --git a/src/test/redirect.rs b/src/test/redirect.rs
index 03639e2..5ec1a41 100644
--- a/src/test/redirect.rs
+++ b/src/test/redirect.rs
@@ -1,3 +1,9 @@
+use std::{
+    io::{self, Write},
+    net::TcpStream,
+};
+use test::testserver::{self, TestServer};
+
 use crate::test;
 
 use super::super::*;
@@ -76,6 +82,25 @@ fn redirect_get() {
     assert_eq!(resp.header("x-foo").unwrap(), "bar");
 }
 
+#[test]
+fn redirect_host() {
+    // Set up a redirect to a host that doesn't exist; it should fail.
+    let srv = TestServer::new(|mut stream: TcpStream| -> io::Result<()> {
+        testserver::read_headers(&stream);
+        write!(stream, "HTTP/1.1 302 Found\r\n")?;
+        write!(stream, "Location: http://example.invalid/\r\n")?;
+        write!(stream, "\r\n")?;
+        Ok(())
+    });
+    let url = format!("http://localhost:{}/", srv.port);
+    let resp = crate::get(&url).call();
+    assert!(
+        matches!(resp.synthetic_error(), Some(Error::DnsFailed(_))),
+        "{:?}",
+        resp.synthetic_error()
+    );
+}
+
 #[test]
 fn redirect_post() {
     test::set_handler("/redirect_post1", |_| {