Repository: algesten/ureq
Pull request #157 (2020-09-25 07:13:37)

Fix set_max_idle_connections(0).
This broke during some recent refactorings. The special case branch for
max_connections == 0 wasn't actually setting the max_idle_connections
field.

This change simply deletes that branch, since the existing code that
whittles down the pool size one element at a time suffices. In the
common case this will be set on an empty pool anyhow.


Patch:
diff --git a/src/agent.rs b/src/agent.rs
index b5b79da..1aa5929 100644
--- a/src/agent.rs
+++ b/src/agent.rs
@@ -6,6 +6,7 @@ use std::sync::Mutex;
 use crate::header::{self, Header};
 use crate::pool::ConnectionPool;
 use crate::request::Request;
+use crate::resolve::ArcResolver;
 
 /// Agents keep state between requests.
 ///
@@ -53,15 +54,12 @@ pub(crate) struct AgentState {
     /// Cookies saved between requests.
     #[cfg(feature = "cookie")]
     pub(crate) jar: CookieJar,
+    pub(crate) resolver: ArcResolver,
 }
 
 impl AgentState {
     fn new() -> Self {
-        AgentState {
-            pool: ConnectionPool::new(),
-            #[cfg(feature = "cookie")]
-            jar: CookieJar::new(),
-        }
+        Self::default()
     }
     pub fn pool(&mut self) -> &mut ConnectionPool {
         &mut self.pool
@@ -194,6 +192,29 @@ impl Agent {
             .set_max_idle_connections_per_host(max_connections);
     }
 
+    /// Configures a custom resolver to be used by this agent. By default,
+    /// address-resolution is done by std::net::ToSocketAddrs. This allows you
+    /// to override that resolution with your own alternative. Useful for
+    /// testing and special-cases like DNS-based load balancing.
+    ///
+    /// A `Fn(&str) -> io::Result<Vec<SocketAddr>>` is a valid resolver,
+    /// passing a closure is a simple way to override. Note that you might need
+    /// explicit type `&str` on the closure argument for type inference to
+    /// succeed.
+    /// ```
+    /// use std::net::ToSocketAddrs;
+    ///
+    /// let mut agent = ureq::agent();
+    /// agent.set_resolver(|addr: &str| match addr {
+    ///    "example.com" => Ok(vec![([127,0,0,1], 8096).into()]),
+    ///    addr => addr.to_socket_addrs().map(Iterator::collect),
+    /// });
+    /// ```
+    pub fn set_resolver(&mut self, resolver: impl crate::Resolver + 'static) -> &mut Self {
+        self.state.lock().unwrap().resolver = resolver.into();
+        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/lib.rs b/src/lib.rs
index df3e1f2..de2129f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -125,6 +125,7 @@ mod header;
 mod pool;
 mod proxy;
 mod request;
+mod resolve;
 mod response;
 mod stream;
 mod unit;
@@ -140,6 +141,7 @@ pub use crate::error::Error;
 pub use crate::header::Header;
 pub use crate::proxy::Proxy;
 pub use crate::request::Request;
+pub use crate::resolve::Resolver;
 pub use crate::response::Response;
 
 // re-export
diff --git a/src/pool.rs b/src/pool.rs
index 4663647..8d66d40 100644
--- a/src/pool.rs
+++ b/src/pool.rs
@@ -74,22 +74,11 @@ impl Default for ConnectionPool {
 }
 
 impl ConnectionPool {
-    pub fn new() -> Self {
-        Self::default()
-    }
-
     pub fn set_max_idle_connections(&mut self, max_connections: usize) {
         if self.max_idle_connections == max_connections {
             return;
         }
 
-        if max_connections == 0 {
-            // Clear the connection pool, caching is disabled.
-            self.lru.clear();
-            self.recycle.clear();
-            return;
-        }
-
         // Remove any extra connections if the number was decreased.
         while self.lru.len() > max_connections {
             self.remove_oldest();
@@ -251,7 +240,7 @@ fn pool_connections_limit() {
     // Test inserting connections with different keys into the pool,
     // filling and draining it. The pool should evict earlier connections
     // when the connection limit is reached.
-    let mut pool = ConnectionPool::new();
+    let mut pool = ConnectionPool::default();
     let hostnames = (0..DEFAULT_MAX_IDLE_CONNECTIONS * 2).map(|i| format!("{}.example", i));
     let poolkeys = hostnames.map(|hostname| PoolKey {
         scheme: "https".to_string(),
@@ -276,7 +265,7 @@ fn pool_per_host_connections_limit() {
     // Test inserting connections with the same key into the pool,
     // filling and draining it. The pool should evict earlier connections
     // when the per-host connection limit is reached.
-    let mut pool = ConnectionPool::new();
+    let mut pool = ConnectionPool::default();
     let poolkey = PoolKey {
         scheme: "https".to_string(),
         hostname: "example.com".to_string(),
@@ -301,7 +290,7 @@ fn pool_per_host_connections_limit() {
 
 #[test]
 fn pool_update_connection_limit() {
-    let mut pool = ConnectionPool::new();
+    let mut pool = ConnectionPool::default();
     pool.set_max_idle_connections(50);
 
     let hostnames = (0..pool.max_idle_connections).map(|i| format!("{}.example", i));
@@ -321,7 +310,7 @@ fn pool_update_connection_limit() {
 
 #[test]
 fn pool_update_per_host_connection_limit() {
-    let mut pool = ConnectionPool::new();
+    let mut pool = ConnectionPool::default();
     pool.set_max_idle_connections(50);
     pool.set_max_idle_connections_per_host(50);
 
@@ -347,7 +336,7 @@ fn pool_update_per_host_connection_limit() {
 fn pool_checks_proxy() {
     // Test inserting different poolkeys with same address but different proxies.
     // Each insertion should result in an additional entry in the pool.
-    let mut pool = ConnectionPool::new();
+    let mut pool = ConnectionPool::default();
     let url = Url::parse("zzz:///example.com").unwrap();
 
     pool.add(
@@ -394,7 +383,7 @@ impl<R: Read + Sized + Into<Stream>> PoolReturnRead<R> {
     fn return_connection(&mut self) {
         // guard we only do this once.
         if let (Some(unit), Some(reader)) = (self.unit.take(), self.reader.take()) {
-            let state = &mut unit.agent.lock().unwrap();
+            let state = &mut unit.req.agent.lock().unwrap();
             // bring back stream here to either go into pool or dealloc
             let stream = reader.into();
             if !stream.is_poolable() {
@@ -402,7 +391,7 @@ impl<R: Read + Sized + Into<Stream>> PoolReturnRead<R> {
                 return;
             }
             // insert back into pool
-            let key = PoolKey::new(&unit.url, &unit.proxy);
+            let key = PoolKey::new(&unit.url, &unit.req.proxy);
             state.pool().add(key, stream);
         }
     }
diff --git a/src/resolve.rs b/src/resolve.rs
new file mode 100644
index 0000000..ca6497d
--- /dev/null
+++ b/src/resolve.rs
@@ -0,0 +1,59 @@
+use std::fmt;
+use std::io::Result as IoResult;
+use std::net::{SocketAddr, ToSocketAddrs};
+use std::sync::Arc;
+
+pub trait Resolver: Send + Sync {
+    fn resolve(&self, netloc: &str) -> IoResult<Vec<SocketAddr>>;
+}
+
+#[derive(Debug)]
+pub(crate) struct StdResolver;
+
+impl Resolver for StdResolver {
+    fn resolve(&self, netloc: &str) -> IoResult<Vec<SocketAddr>> {
+        ToSocketAddrs::to_socket_addrs(netloc).map(|iter| iter.collect())
+    }
+}
+
+impl<F> Resolver for F
+where
+    F: Fn(&str) -> IoResult<Vec<SocketAddr>>,
+    F: Send + Sync,
+{
+    fn resolve(&self, netloc: &str) -> IoResult<Vec<SocketAddr>> {
+        self(netloc)
+    }
+}
+
+#[derive(Clone)]
+pub(crate) struct ArcResolver(Arc<dyn Resolver>);
+
+impl<R> From<R> for ArcResolver
+where
+    R: Resolver + 'static,
+{
+    fn from(r: R) -> Self {
+        Self(Arc::new(r))
+    }
+}
+
+impl fmt::Debug for ArcResolver {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "ArcResolver(...)")
+    }
+}
+
+impl std::ops::Deref for ArcResolver {
+    type Target = dyn Resolver;
+
+    fn deref(&self) -> &Self::Target {
+        self.0.as_ref()
+    }
+}
+
+impl Default for ArcResolver {
+    fn default() -> Self {
+        StdResolver.into()
+    }
+}
diff --git a/src/stream.rs b/src/stream.rs
index b0f8309..c51520f 100644
--- a/src/stream.rs
+++ b/src/stream.rs
@@ -4,7 +4,6 @@ use std::io::{
 };
 use std::net::SocketAddr;
 use std::net::TcpStream;
-use std::net::ToSocketAddrs;
 use std::time::Duration;
 use std::time::Instant;
 
@@ -343,8 +342,12 @@ pub(crate) fn connect_https(unit: &Unit) -> Result<Stream, Error> {
 
     let sni = webpki::DNSNameRef::try_from_ascii_str(hostname)
         .map_err(|err| Error::DnsFailed(err.to_string()))?;
-    let tls_conf: &Arc<rustls::ClientConfig> =
-        unit.tls_config.as_ref().map(|c| &c.0).unwrap_or(&*TLS_CONF);
+    let tls_conf: &Arc<rustls::ClientConfig> = unit
+        .req
+        .tls_config
+        .as_ref()
+        .map(|c| &c.0)
+        .unwrap_or(&*TLS_CONF);
     let sess = rustls::ClientSession::new(&tls_conf, sni);
 
     let sock = connect_host(unit, hostname, port)?;
@@ -362,7 +365,7 @@ pub(crate) fn connect_https(unit: &Unit) -> Result<Stream, Error> {
     let port = unit.url.port().unwrap_or(443);
     let sock = connect_host(unit, hostname, port)?;
 
-    let tls_connector: Arc<native_tls::TlsConnector> = match &unit.tls_connector {
+    let tls_connector: Arc<native_tls::TlsConnector> = match &unit.req.tls_connector {
         Some(connector) => connector.0.clone(),
         None => Arc::new(native_tls::TlsConnector::new().map_err(|e| Error::TlsError(e))?),
     };
@@ -377,26 +380,28 @@ pub(crate) fn connect_https(unit: &Unit) -> Result<Stream, Error> {
 }
 
 pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<TcpStream, Error> {
-    let deadline: Option<Instant> = if unit.timeout_connect > 0 {
-        Instant::now().checked_add(Duration::from_millis(unit.timeout_connect))
+    let deadline: Option<Instant> = if unit.req.timeout_connect > 0 {
+        Instant::now().checked_add(Duration::from_millis(unit.req.timeout_connect))
     } else {
         unit.deadline
     };
-
-    // TODO: Find a way to apply deadline to DNS lookup.
-    let sock_addrs: Vec<SocketAddr> = match unit.proxy {
+  
+    let netloc = match unit.req.proxy {
         Some(ref proxy) => format!("{}:{}", proxy.server, proxy.port),
         None => format!("{}:{}", hostname, port),
-    }
-    .to_socket_addrs()
-    .map_err(|e| Error::DnsFailed(format!("{}", e)))?
-    .collect();
+    };
+
+    // TODO: Find a way to apply deadline to DNS lookup.
+    let sock_addrs = unit
+        .resolver()
+        .resolve(&netloc)
+        .map_err(|e| Error::DnsFailed(format!("{}", e)))?;
 
     if sock_addrs.is_empty() {
         return Err(Error::DnsFailed(format!("No ip address for {}", hostname)));
     }
 
-    let proto = if let Some(ref proxy) = unit.proxy {
+    let proto = if let Some(ref proxy) = unit.req.proxy {
         Some(proxy.proto)
     } else {
         None
@@ -415,7 +420,8 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
         // connect with a configured timeout.
         let stream = if Some(Proto::SOCKS5) == proto {
             connect_socks5(
-                unit.proxy.to_owned().unwrap(),
+                &unit,
+                unit.req.proxy.to_owned().unwrap(),
                 deadline,
                 sock_addr,
                 hostname,
@@ -448,9 +454,9 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
         stream
             .set_read_timeout(Some(time_until_deadline(deadline)?))
             .ok();
-    } else if unit.timeout_read > 0 {
+    } else if unit.req.timeout_read > 0 {
         stream
-            .set_read_timeout(Some(Duration::from_millis(unit.timeout_read as u64)))
+            .set_read_timeout(Some(Duration::from_millis(unit.req.timeout_read as u64)))
             .ok();
     } else {
         stream.set_read_timeout(None).ok();
@@ -460,16 +466,16 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
         stream
             .set_write_timeout(Some(time_until_deadline(deadline)?))
             .ok();
-    } else if unit.timeout_write > 0 {
+    } else if unit.req.timeout_write > 0 {
         stream
-            .set_write_timeout(Some(Duration::from_millis(unit.timeout_write as u64)))
+            .set_write_timeout(Some(Duration::from_millis(unit.req.timeout_write as u64)))
             .ok();
     } else {
         stream.set_write_timeout(None).ok();
     }
 
     if proto == Some(Proto::HTTPConnect) {
-        if let Some(ref proxy) = unit.proxy {
+        if let Some(ref proxy) = unit.req.proxy {
             write!(stream, "{}", proxy.connect(hostname, port)).unwrap();
             stream.flush()?;
 
@@ -492,11 +498,15 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
 }
 
 #[cfg(feature = "socks-proxy")]
-fn socks5_local_nslookup(hostname: &str, port: u16) -> Result<TargetAddr, std::io::Error> {
-    let addrs: Vec<SocketAddr> = format!("{}:{}", hostname, port)
-        .to_socket_addrs()
-        .map_err(|e| std::io::Error::new(ErrorKind::NotFound, format!("DNS failure: {}.", e)))?
-        .collect();
+fn socks5_local_nslookup(
+    unit: &Unit,
+    hostname: &str,
+    port: u16,
+) -> Result<TargetAddr, std::io::Error> {
+    let addrs: Vec<SocketAddr> = unit
+        .resolver()
+        .resolve(&format!("{}:{}", hostname, port))
+        .map_err(|e| std::io::Error::new(ErrorKind::NotFound, format!("DNS failure: {}.", e)))?;
 
     if addrs.is_empty() {
         return Err(std::io::Error::new(
@@ -518,6 +528,7 @@ fn socks5_local_nslookup(hostname: &str, port: u16) -> Result<TargetAddr, std::i
 
 #[cfg(feature = "socks-proxy")]
 fn connect_socks5(
+    unit: &Unit,
     proxy: Proxy,
     deadline: Option<Instant>,
     proxy_addr: SocketAddr,
@@ -529,7 +540,7 @@ fn connect_socks5(
     use std::str::FromStr;
 
     let host_addr = if Ipv4Addr::from_str(host).is_ok() || Ipv6Addr::from_str(host).is_ok() {
-        match socks5_local_nslookup(host, port) {
+        match socks5_local_nslookup(unit, host, port) {
             Ok(addr) => addr,
             Err(err) => return Err(err),
         }
@@ -621,6 +632,7 @@ fn get_socks5_stream(
 
 #[cfg(not(feature = "socks-proxy"))]
 fn connect_socks5(
+    _unit: &Unit,
     _proxy: Proxy,
     _deadline: Option<Instant>,
     _proxy_addr: SocketAddr,
diff --git a/src/unit.rs b/src/unit.rs
index 1edc634..f3822f2 100644
--- a/src/unit.rs
+++ b/src/unit.rs
@@ -1,5 +1,4 @@
 use std::io::{Result as IoResult, Write};
-use std::sync::{Arc, Mutex};
 use std::time;
 
 use qstring::QString;
@@ -11,16 +10,10 @@ use cookie::{Cookie, CookieJar};
 use crate::agent::AgentState;
 use crate::body::{self, Payload, SizedReader};
 use crate::header;
+use crate::resolve::ArcResolver;
 use crate::stream::{self, connect_test, Stream};
-use crate::Proxy;
 use crate::{Error, Header, Request, Response};
 
-#[cfg(feature = "tls")]
-use crate::request::TLSClientConfig;
-
-#[cfg(all(feature = "native-tls", not(feature = "tls")))]
-use crate::request::TLSConnector;
-
 #[cfg(feature = "cookie")]
 use crate::pool::DEFAULT_HOST;
 
@@ -28,21 +21,12 @@ use crate::pool::DEFAULT_HOST;
 ///
 /// *Internal API*
 pub(crate) struct Unit {
-    pub agent: Arc<Mutex<AgentState>>,
+    pub req: Request,
     pub url: Url,
     pub is_chunked: bool,
     pub query_string: String,
     pub headers: Vec<Header>,
-    pub timeout_connect: u64,
-    pub timeout_read: u64,
-    pub timeout_write: u64,
     pub deadline: Option<time::Instant>,
-    pub method: String,
-    pub proxy: Option<Proxy>,
-    #[cfg(feature = "tls")]
-    pub tls_config: Option<TLSClientConfig>,
-    #[cfg(all(feature = "native-tls", not(feature = "tls")))]
-    pub tls_connector: Option<TLSConnector>,
 }
 
 impl Unit {
@@ -99,26 +83,21 @@ impl Unit {
         };
 
         Unit {
-            agent: Arc::clone(&req.agent),
+            req: req.clone(),
             url: url.clone(),
             is_chunked,
             query_string,
             headers,
-            timeout_connect: req.timeout_connect,
-            timeout_read: req.timeout_read,
-            timeout_write: req.timeout_write,
             deadline,
-            method: req.method.clone(),
-            proxy: req.proxy.clone(),
-            #[cfg(feature = "tls")]
-            tls_config: req.tls_config.clone(),
-            #[cfg(all(feature = "native-tls", not(feature = "tls")))]
-            tls_connector: req.tls_connector.clone(),
         }
     }
 
     pub fn is_head(&self) -> bool {
-        self.method.eq_ignore_ascii_case("head")
+        self.req.method.eq_ignore_ascii_case("head")
+    }
+
+    pub fn resolver(&self) -> ArcResolver {
+        self.req.agent.lock().unwrap().resolver.clone()
     }
 
     #[cfg(test)]
@@ -217,8 +196,8 @@ pub(crate) fn connect(
                     let mut new_unit = Unit::new(req, &new_url, false, &empty);
                     // this is to follow how curl does it. POST, PUT etc change
                     // to GET on a redirect.
-                    new_unit.method = match &unit.method[..] {
-                        "GET" | "HEAD" => unit.method,
+                    new_unit.req.method = match &unit.req.method[..] {
+                        "GET" | "HEAD" => unit.req.method,
                         _ => "GET".into(),
                     };
                     return connect(req, new_unit, use_pooled, redirect_count + 1, empty, true);
@@ -305,11 +284,11 @@ fn connect_socket(unit: &Unit, use_pooled: bool) -> Result<(Stream, bool), Error
         _ => return Err(Error::UnknownScheme(unit.url.scheme().to_string())),
     };
     if use_pooled {
-        let state = &mut unit.agent.lock().unwrap();
+        let state = &mut unit.req.agent.lock().unwrap();
         // The connection may have been closed by the server
         // due to idle timeout while it was sitting in the pool.
         // Loop until we find one that is still good or run out of connections.
-        while let Some(stream) = state.pool.try_get_connection(&unit.url, &unit.proxy) {
+        while let Some(stream) = state.pool.try_get_connection(&unit.url, &unit.req.proxy) {
             let server_closed = stream.server_closed()?;
             if !server_closed {
                 return Ok((stream, true));
@@ -337,7 +316,7 @@ fn send_prelude(unit: &Unit, stream: &mut Stream, redir: bool) -> IoResult<()> {
     write!(
         prelude,
         "{} {}{} HTTP/1.1\r\n",
-        unit.method,
+        unit.req.method,
         unit.url.path(),
         &unit.query_string
     )?;
@@ -404,7 +383,7 @@ fn save_cookies(unit: &Unit, resp: &Response) {
     }
 
     // only lock if we know there is something to process
-    let state = &mut unit.agent.lock().unwrap();
+    let state = &mut unit.req.agent.lock().unwrap();
     for raw_cookie in cookies.iter() {
         let to_parse = if raw_cookie.to_lowercase().contains("domain=") {
             (*raw_cookie).to_string()


Test patch:
diff --git a/src/test/agent_test.rs b/src/test/agent_test.rs
index f080352..ea107ec 100644
--- a/src/test/agent_test.rs
+++ b/src/test/agent_test.rs
@@ -101,6 +101,31 @@ fn connection_reuse() {
     assert_eq!(resp.status(), 200);
 }
 
+#[test]
+fn custom_resolver() {
+    use std::io::Read;
+    use std::net::TcpListener;
+
+    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+
+    let local_addr = listener.local_addr().unwrap();
+
+    let server = std::thread::spawn(move || {
+        let (mut client, _) = listener.accept().unwrap();
+        let mut buf = vec![0u8; 16];
+        let read = client.read(&mut buf).unwrap();
+        buf.truncate(read);
+        buf
+    });
+
+    crate::agent()
+        .set_resolver(move |_: &str| Ok(vec![local_addr]))
+        .get("http://cool.server/")
+        .call();
+
+    assert_eq!(&server.join().unwrap(), b"GET / HTTP/1.1\r\n");
+}
+
 #[cfg(feature = "cookie")]
 #[cfg(test)]
 fn cookie_and_redirect(mut stream: TcpStream) -> io::Result<()> {
diff --git a/src/test/redirect.rs b/src/test/redirect.rs
index 705c672..03639e2 100644
--- a/src/test/redirect.rs
+++ b/src/test/redirect.rs
@@ -46,7 +46,7 @@ fn redirect_head() {
         test::make_response(302, "Go here", vec!["Location: /redirect_head2"], vec![])
     });
     test::set_handler("/redirect_head2", |unit| {
-        assert_eq!(unit.method, "HEAD");
+        assert_eq!(unit.req.method, "HEAD");
         test::make_response(200, "OK", vec!["x-foo: bar"], vec![])
     });
     let resp = head("test://host/redirect_head1").call();
@@ -62,7 +62,7 @@ fn redirect_get() {
         test::make_response(302, "Go here", vec!["Location: /redirect_get2"], vec![])
     });
     test::set_handler("/redirect_get2", |unit| {
-        assert_eq!(unit.method, "GET");
+        assert_eq!(unit.req.method, "GET");
         assert!(unit.has("Range"));
         assert_eq!(unit.header("Range").unwrap(), "bytes=10-50");
         test::make_response(200, "OK", vec!["x-foo: bar"], vec![])
@@ -82,7 +82,7 @@ fn redirect_post() {
         test::make_response(302, "Go here", vec!["Location: /redirect_post2"], vec![])
     });
     test::set_handler("/redirect_post2", |unit| {
-        assert_eq!(unit.method, "GET");
+        assert_eq!(unit.req.method, "GET");
         test::make_response(200, "OK", vec!["x-foo: bar"], vec![])
     });
     let resp = post("test://host/redirect_post1").call();