From 3014f58a28eee69543fea156835aa053a6c91097 Mon Sep 17 00:00:00 2001 From: Jacob Hoffman-Andrews Date: Mon, 22 Jun 2020 23:23:39 -0700 Subject: [PATCH] Add scheme to PoolKey and let port be None. (#84) PoolKey calls unwrap() on an option that can be None. Specifically, the local variable `port` can be None when PoolKey is constructed with a Url whose scheme is unrecognized in url.port_or_known_default(). To fix that, make port an Option. Also, make scheme part of the PoolKey. This prevents, for instance, a stream opened for `https://example.com:9999` being reused on a request for `http://example.com:9999`. Remove the test-only pool.get() accessor. This was used in only one test, agent_pool in range.rs. This seemed like it was testing the agent more than it was testing ranges, so I moved it to agent.rs and edited to remove the range-testing parts. Also, reject unrecognized URLs earlier in connect_socket so they don't reach try_get_connection. --- src/agent.rs | 30 ++++++++++++++++++++++++++++++ src/pool.rs | 35 ++++++++++++----------------------- src/test/range.rs | 38 -------------------------------------- src/unit.rs | 4 ++++ 4 files changed, 46 insertions(+), 61 deletions(-) diff --git a/src/agent.rs b/src/agent.rs index ad8266e..4eaf6d1 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -262,6 +262,7 @@ pub(crate) fn basic_auth(user: &str, pass: &str) -> String { #[cfg(test)] mod tests { use super::*; + use std::io::Read; ///////////////////// AGENT TESTS ////////////////////////////// @@ -273,6 +274,35 @@ mod tests { }); } + #[test] + #[cfg(any(feature = "tls", feature = "native-tls"))] + fn agent_pool() { + let agent = crate::agent(); + let url = "https://ureq.s3.eu-central-1.amazonaws.com/sherlock.txt"; + // req 1 + let resp = agent.get(url).call(); + assert!(resp.ok()); + let mut reader = resp.into_reader(); + let mut buf = vec![]; + // reading the entire content will return the connection to the pool + reader.read_to_end(&mut buf).unwrap(); + + fn poolsize(agent: &Agent) -> usize { + let mut lock = agent.state.lock().unwrap(); + let state = lock.as_mut().unwrap(); + state.pool().len() + } + assert_eq!(poolsize(&agent), 1); + + // req 2 should be done with a reused connection + let resp = agent.get(url).call(); + assert!(resp.ok()); + assert_eq!(poolsize(&agent), 0); + let mut reader = resp.into_reader(); + let mut buf = vec![]; + reader.read_to_end(&mut buf).unwrap(); + } + //////////////////// REQUEST TESTS ///////////////////////////// #[test] diff --git a/src/pool.rs b/src/pool.rs index 0fe8a56..f5db413 100644 --- a/src/pool.rs +++ b/src/pool.rs @@ -33,43 +33,32 @@ impl ConnectionPool { pub fn len(&self) -> usize { self.recycle.len() } - - #[cfg(all(test, any(feature = "tls", feature = "native-tls")))] - pub fn get(&self, hostname: &str, port: u16) -> Option<&Stream> { - let key = PoolKey { - hostname: hostname.into(), - port, - }; - self.recycle.get(&key) - } } #[derive(Debug, PartialEq, Clone, Eq, Hash)] struct PoolKey { + scheme: String, hostname: String, - port: u16, + port: Option, } impl PoolKey { fn new(url: &Url) -> Self { - let port = if cfg!(test) { - if let Some(p) = url.port_or_known_default() { - Some(p) - } else if url.scheme() == "test" { - Some(42) - } else { - None - } - } else { - url.port_or_known_default() - }; + let port = url.port_or_known_default(); PoolKey { - hostname: url.host_str().unwrap_or(DEFAULT_HOST).into(), - port: port.expect("Failed to get port for pool key"), + scheme: url.scheme().to_string(), + hostname: url.host_str().unwrap_or("").to_string(), + port, } } } +#[test] +fn poolkey_new() { + // Test that PoolKey::new() does not panic on unrecognized schemes. + PoolKey::new(&Url::parse("zzz:///example.com").unwrap()); +} + /// Read wrapper that returns the stream to the pool once the /// read is exhausted (reached a 0). /// diff --git a/src/test/range.rs b/src/test/range.rs index 9c4adbc..cdef429 100644 --- a/src/test/range.rs +++ b/src/test/range.rs @@ -20,41 +20,3 @@ fn read_range() { [83, 99, 111, 116, 116, 34, 10, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32] ) } - -#[test] -#[cfg(any(feature = "tls", feature = "native-tls"))] -fn agent_pool() { - let agent = agent(); - - // req 1 - let resp = agent - .get("https://ureq.s3.eu-central-1.amazonaws.com/sherlock.txt") - .set("Range", "bytes=1000-1999") - .call(); - assert_eq!(resp.status(), 206); - let mut reader = resp.into_reader(); - let mut buf = vec![]; - // reading the entire content will return the connection to the pool - let len = reader.read_to_end(&mut buf).unwrap(); - assert_eq!(len, 1000); - - { - let mut lock = agent.state.lock().unwrap(); - let state = lock.as_mut().unwrap(); - let pool = state.pool(); - assert_eq!(pool.len(), 1); - let f = format!("{:?}", pool.get("ureq.s3.eu-central-1.amazonaws.com", 443)); - assert_eq!(f, "Some(Stream[https])"); // not a great way of testing. - } - - // req 2 should be done with a reused connection - let resp = agent - .get("https://ureq.s3.eu-central-1.amazonaws.com/sherlock.txt") - .set("Range", "bytes=5000-6999") - .call(); - assert_eq!(resp.status(), 206); - let mut reader = resp.into_reader(); - let mut buf = vec![]; - let len = reader.read_to_end(&mut buf).unwrap(); - assert_eq!(len, 2000); -} diff --git a/src/unit.rs b/src/unit.rs index cea0e05..b6fd18b 100644 --- a/src/unit.rs +++ b/src/unit.rs @@ -291,6 +291,10 @@ pub(crate) fn combine_query(url: &Url, query: &QString, mix_queries: bool) -> St /// Connect the socket, either by using the pool or grab a new one. fn connect_socket(unit: &Unit, use_pooled: bool) -> Result<(Stream, bool), Error> { + match unit.url.scheme() { + "http" | "https" | "test" => (), + _ => return Err(Error::UnknownScheme(unit.url.scheme().to_string())), + }; if use_pooled { let state = &mut unit.agent.lock().unwrap(); if let Some(agent) = state.as_mut() {