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:
committed by
GitHub
parent
00461fb5bd
commit
3014f58a28
30
src/agent.rs
30
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]
|
||||
|
||||
35
src/pool.rs
35
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<u16>,
|
||||
}
|
||||
|
||||
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).
|
||||
///
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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() {
|
||||
|
||||
Reference in New Issue
Block a user