Add scheme to PoolKey and let port be None. (fortanix#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.

Author: jsha

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]

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).
 ///

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);
-}

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() {