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)]
|
#[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]
|
||||||
|
|||||||
35
src/pool.rs
35
src/pool.rs
@@ -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).
|
||||||
///
|
///
|
||||||
|
|||||||
@@ -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);
|
|
||||||
}
|
|
||||||
|
|||||||
@@ -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() {
|
||||||
|
|||||||
Reference in New Issue
Block a user