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.
This commit is contained in:
Jacob Hoffman-Andrews
2020-06-22 23:23:39 -07:00
committed by GitHub
parent 00461fb5bd
commit 3014f58a28
4 changed files with 46 additions and 61 deletions

View File

@@ -262,6 +262,7 @@ pub(crate) fn basic_auth(user: &str, pass: &str) -> String {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use std::io::Read;
///////////////////// AGENT TESTS ////////////////////////////// ///////////////////// 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 ///////////////////////////// //////////////////// REQUEST TESTS /////////////////////////////
#[test] #[test]

View File

@@ -33,43 +33,32 @@ impl ConnectionPool {
pub fn len(&self) -> usize { pub fn len(&self) -> usize {
self.recycle.len() 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)] #[derive(Debug, PartialEq, Clone, Eq, Hash)]
struct PoolKey { struct PoolKey {
scheme: String,
hostname: String, hostname: String,
port: u16, port: Option<u16>,
} }
impl PoolKey { impl PoolKey {
fn new(url: &Url) -> Self { fn new(url: &Url) -> Self {
let port = if cfg!(test) { let port = url.port_or_known_default();
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()
};
PoolKey { PoolKey {
hostname: url.host_str().unwrap_or(DEFAULT_HOST).into(), scheme: url.scheme().to_string(),
port: port.expect("Failed to get port for pool key"), 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 wrapper that returns the stream to the pool once the
/// read is exhausted (reached a 0). /// read is exhausted (reached a 0).
/// ///

View File

@@ -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] [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);
}

View File

@@ -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. /// 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> { 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 { if use_pooled {
let state = &mut unit.agent.lock().unwrap(); let state = &mut unit.agent.lock().unwrap();
if let Some(agent) = state.as_mut() { if let Some(agent) = state.as_mut() {