Repository: algesten/ureq
Pull request #138 (2020-08-31 06:12:08)

Put the crate version in the User-Agent header.
https://tools.ietf.org/html/rfc7231#section-5.5.3
Each product identifier consists of a name and optional version.

     product         = token ["/" product-version]
     product-version = token


Patch:
diff --git a/src/agent.rs b/src/agent.rs
index 4eaf6d1..b7a43ee 100644
--- a/src/agent.rs
+++ b/src/agent.rs
@@ -166,6 +166,38 @@ impl Agent {
         Request::new(&self, method.into(), path.into())
     }
 
+    /// Sets the maximum number of connections allowed in the connection pool.
+    /// By default, this is set to 100. Setting this to zero would disable
+    /// connection pooling.
+    ///
+    /// ```
+    /// let agent = ureq::agent();
+    /// agent.set_max_pool_connections(200);
+    /// ```
+    pub fn set_max_pool_connections(&self, max_connections: usize) {
+        let mut optional_state = self.state.lock().unwrap();
+        if let Some(state) = optional_state.as_mut() {
+            state.pool.set_max_idle_connections(max_connections);
+        }
+    }
+
+    /// Sets the maximum number of connections per host to keep in the
+    /// connection pool. By default, this is set to 1. Setting this to zero
+    /// would disable connection pooling.
+    ///
+    /// ```
+    /// let agent = ureq::agent();
+    /// agent.set_max_pool_connections_per_host(10);
+    /// ```
+    pub fn set_max_pool_connections_per_host(&self, max_connections: usize) {
+        let mut optional_state = self.state.lock().unwrap();
+        if let Some(state) = optional_state.as_mut() {
+            state
+                .pool
+                .set_max_idle_connections_per_host(max_connections);
+        }
+    }
+
     /// Gets a cookie in this agent by name. Cookies are available
     /// either by setting it in the agent, or by making requests
     /// that `Set-Cookie` in the agent.
@@ -256,20 +288,21 @@ pub(crate) fn basic_auth(user: &str, pass: &str) -> String {
         Some(idx) => &user[..idx],
         None => user,
     };
-    ::base64::encode(&format!("{}:{}", safe, pass))
+    base64::encode(&format!("{}:{}", safe, pass))
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use std::io::Read;
+    use std::thread;
 
     ///////////////////// AGENT TESTS //////////////////////////////
 
     #[test]
     fn agent_implements_send() {
         let mut agent = Agent::new();
-        ::std::thread::spawn(move || {
+        thread::spawn(move || {
             agent.set("Foo", "Bar");
         });
     }
@@ -309,7 +342,7 @@ mod tests {
     fn request_implements_send() {
         let agent = Agent::new();
         let mut request = Request::new(&agent, "GET".to_string(), "/foo".to_string());
-        ::std::thread::spawn(move || {
+        thread::spawn(move || {
             request.set("Foo", "Bar");
         });
     }
diff --git a/src/body.rs b/src/body.rs
index f87763e..b6b188f 100644
--- a/src/body.rs
+++ b/src/body.rs
@@ -1,4 +1,5 @@
 use crate::stream::Stream;
+use std::fmt;
 use std::io::{copy, empty, Cursor, Read, Result as IoResult, Write};
 
 #[cfg(feature = "charset")]
@@ -23,8 +24,8 @@ pub(crate) enum Payload {
     Bytes(Vec<u8>),
 }
 
-impl ::std::fmt::Debug for Payload {
-    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::result::Result<(), ::std::fmt::Error> {
+impl fmt::Debug for Payload {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             Payload::Empty => write!(f, "Empty"),
             Payload::Text(t, _) => write!(f, "{}", t),
@@ -50,8 +51,8 @@ pub(crate) struct SizedReader {
     pub reader: Box<dyn Read + 'static>,
 }
 
-impl ::std::fmt::Debug for SizedReader {
-    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::result::Result<(), ::std::fmt::Error> {
+impl fmt::Debug for SizedReader {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(f, "SizedReader[size={:?},reader]", self.size)
     }
 }
diff --git a/src/header.rs b/src/header.rs
index b788c39..bf3dfed 100644
--- a/src/header.rs
+++ b/src/header.rs
@@ -1,4 +1,5 @@
 use crate::error::Error;
+use std::fmt;
 use std::str::FromStr;
 
 #[derive(Clone, PartialEq)]
@@ -14,8 +15,8 @@ pub struct Header {
     index: usize,
 }
 
-impl ::std::fmt::Debug for Header {
-    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::result::Result<(), ::std::fmt::Error> {
+impl fmt::Debug for Header {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(f, "{}", self.line)
     }
 }
diff --git a/src/pool.rs b/src/pool.rs
index d9ca195..ae2dd3f 100644
--- a/src/pool.rs
+++ b/src/pool.rs
@@ -1,3 +1,4 @@
+use std::collections::hash_map::Entry;
 use std::collections::{HashMap, VecDeque};
 use std::io::{Read, Result as IoResult};
 
@@ -8,7 +9,8 @@ use crate::Proxy;
 use url::Url;
 
 pub const DEFAULT_HOST: &str = "localhost";
-const MAX_IDLE_CONNECTIONS: usize = 100;
+const DEFAULT_MAX_IDLE_CONNECTIONS: usize = 100;
+const DEFAULT_MAX_IDLE_CONNECTIONS_PER_HOST: usize = 1;
 
 /// Holder of recycled connections.
 ///
@@ -20,21 +22,71 @@ const MAX_IDLE_CONNECTIONS: usize = 100;
 #[derive(Default, Debug)]
 pub(crate) struct ConnectionPool {
     // the actual pooled connection. however only one per hostname:port.
-    recycle: HashMap<PoolKey, Stream>,
+    recycle: HashMap<PoolKey, VecDeque<Stream>>,
     // This is used to keep track of which streams to expire when the
     // pool reaches MAX_IDLE_CONNECTIONS. The corresponding PoolKeys for
     // recently used Streams are added to the back of the queue;
     // old streams are removed from the front.
     lru: VecDeque<PoolKey>,
+    max_idle_connections: usize,
+    max_idle_connections_per_host: usize,
 }
 
 impl ConnectionPool {
     pub fn new() -> Self {
         ConnectionPool {
+            max_idle_connections: DEFAULT_MAX_IDLE_CONNECTIONS,
+            max_idle_connections_per_host: DEFAULT_MAX_IDLE_CONNECTIONS_PER_HOST,
             ..Default::default()
         }
     }
 
+    pub fn set_max_idle_connections(&mut self, max_connections: usize) {
+        if self.max_idle_connections == max_connections {
+            return;
+        }
+        self.max_idle_connections = max_connections;
+
+        if max_connections == 0 {
+            // Clear the connection pool, caching is disabled.
+            self.lru.clear();
+            self.recycle.clear();
+            return;
+        }
+
+        // Remove any extra connections if the number was decreased.
+        while self.lru.len() > max_connections {
+            self.remove_oldest();
+        }
+    }
+
+    pub fn set_max_idle_connections_per_host(&mut self, max_connections: usize) {
+        if self.max_idle_connections_per_host == max_connections {
+            return;
+        }
+        self.max_idle_connections_per_host = max_connections;
+
+        if max_connections == 0 {
+            // Clear the connection pool, caching is disabled.
+            self.lru.clear();
+            self.recycle.clear();
+            return;
+        }
+
+        // Remove any extra streams if the number was decreased.
+        for (key, val) in self.recycle.iter_mut() {
+            while val.len() > max_connections {
+                val.pop_front();
+                let index = self
+                    .lru
+                    .iter()
+                    .position(|x| x == key)
+                    .expect("PoolKey not found in lru");
+                self.lru.remove(index);
+            }
+        }
+    }
+
     /// How the unit::connect tries to get a pooled connection.
     pub fn try_get_connection(&mut self, url: &Url, proxy: &Option<Proxy>) -> Option<Stream> {
         let key = PoolKey::new(url, proxy);
@@ -42,36 +94,85 @@ impl ConnectionPool {
     }
 
     fn remove(&mut self, key: &PoolKey) -> Option<Stream> {
-        if !self.recycle.contains_key(&key) {
-            return None;
+        match self.recycle.entry(key.clone()) {
+            Entry::Occupied(mut occupied_entry) => {
+                let streams = occupied_entry.get_mut();
+                // Take the newest stream.
+                let stream = streams.pop_back();
+                assert!(
+                    stream.is_some(),
+                    "key existed in recycle but no streams available"
+                );
+
+                if streams.len() == 0 {
+                    occupied_entry.remove();
+                }
+
+                // Remove the oldest matching PoolKey from self.lru.
+                // since this PoolKey was most recently used, removing the oldest
+                // PoolKey would delay other streams with this address from
+                // being removed.
+                self.remove_from_lru(key);
+
+                stream
+            }
+            Entry::Vacant(_) => None,
         }
-        let index = self.lru.iter().position(|k| k == key);
-        assert!(
-            index.is_some(),
-            "invariant failed: key existed in recycle but not lru"
-        );
-        self.lru.remove(index.unwrap());
-        self.recycle.remove(&key)
+    }
+
+    fn remove_from_lru(&mut self, key: &PoolKey) {
+        let index = self
+            .lru
+            .iter()
+            .position(|x| x == key)
+            .expect("PoolKey not found in lru");
+        self.lru.remove(index);
     }
 
     fn add(&mut self, key: PoolKey, stream: Stream) {
-        // If an entry with the same key already exists, remove it.
-        // The more recently used stream is likely to live longer.
-        self.remove(&key);
-        if self.recycle.len() + 1 > MAX_IDLE_CONNECTIONS {
-            self.remove_oldest();
+        if self.max_idle_connections == 0 || self.max_idle_connections_per_host == 0 {
+            return;
+        }
+
+        match self.recycle.entry(key.clone()) {
+            Entry::Occupied(mut occupied_entry) => {
+                let streams = occupied_entry.get_mut();
+                streams.push_back(stream);
+                if streams.len() > self.max_idle_connections_per_host {
+                    streams.pop_front();
+                    self.remove_from_lru(&key);
+                }
+            }
+            Entry::Vacant(vacant_entry) => {
+                let mut new_deque = VecDeque::new();
+                new_deque.push_back(stream);
+                vacant_entry.insert(new_deque);
+            }
+        }
+        self.lru.push_back(key);
+        if self.lru.len() > self.max_idle_connections {
+            self.remove_oldest()
         }
-        self.lru.push_back(key.clone());
-        self.recycle.insert(key, stream);
     }
 
     fn remove_oldest(&mut self) {
         if let Some(key) = self.lru.pop_front() {
-            let removed = self.recycle.remove(&key);
-            assert!(
-                removed.is_some(),
-                "invariant failed: key existed in lru but not in recycle"
-            );
+            match self.recycle.entry(key) {
+                Entry::Occupied(mut occupied_entry) => {
+                    let streams = occupied_entry.get_mut();
+                    let removed_stream = streams.pop_front();
+                    assert!(
+                        removed_stream.is_some(),
+                        "key existed in recycle but no streams available"
+                    );
+                    if streams.len() == 0 {
+                        occupied_entry.remove();
+                    }
+                }
+                Entry::Vacant(_) => {
+                    panic!("invariant failed: key existed in lru but not in recycle")
+                }
+            }
         } else {
             panic!("tried to remove oldest but no entries found!");
         }
@@ -79,7 +180,7 @@ impl ConnectionPool {
 
     #[cfg(test)]
     pub fn len(&self) -> usize {
-        self.recycle.len()
+        self.lru.len()
     }
 }
 
@@ -123,10 +224,12 @@ fn poolkey_new() {
 }
 
 #[test]
-fn pool_size_limit() {
-    assert_eq!(MAX_IDLE_CONNECTIONS, 100);
+fn pool_connections_limit() {
+    // Test inserting connections with different keys into the pool,
+    // filling and draining it. The pool should evict earlier connections
+    // when the connection limit is reached.
     let mut pool = ConnectionPool::new();
-    let hostnames = (0..200).map(|i| format!("{}.example", i));
+    let hostnames = (0..DEFAULT_MAX_IDLE_CONNECTIONS * 2).map(|i| format!("{}.example", i));
     let poolkeys = hostnames.map(|hostname| PoolKey {
         scheme: "https".to_string(),
         hostname,
@@ -136,22 +239,49 @@ fn pool_size_limit() {
     for key in poolkeys.clone() {
         pool.add(key, Stream::Cursor(std::io::Cursor::new(vec![])));
     }
-    assert_eq!(pool.len(), 100);
+    assert_eq!(pool.len(), DEFAULT_MAX_IDLE_CONNECTIONS);
 
-    for key in poolkeys.skip(100) {
+    for key in poolkeys.skip(DEFAULT_MAX_IDLE_CONNECTIONS) {
         let result = pool.remove(&key);
         assert!(result.is_some(), "expected key was not in pool");
     }
+    assert_eq!(pool.len(), 0)
 }
 
 #[test]
-fn pool_duplicates_limit() {
-    // Test inserting duplicates into the pool, and subsequently
-    // filling and draining it. The duplicates should evict earlier
-    // entries with the same key.
-    assert_eq!(MAX_IDLE_CONNECTIONS, 100);
+fn pool_per_host_connections_limit() {
+    // Test inserting connections with the same key into the pool,
+    // filling and draining it. The pool should evict earlier connections
+    // when the per-host connection limit is reached.
     let mut pool = ConnectionPool::new();
-    let hostnames = (0..100).map(|i| format!("{}.example", i));
+    let poolkey = PoolKey {
+        scheme: "https".to_string(),
+        hostname: "example.com".to_string(),
+        port: Some(999),
+        proxy: None,
+    };
+
+    for _ in 0..pool.max_idle_connections_per_host * 2 {
+        pool.add(
+            poolkey.clone(),
+            Stream::Cursor(std::io::Cursor::new(vec![])),
+        );
+    }
+    assert_eq!(pool.len(), DEFAULT_MAX_IDLE_CONNECTIONS_PER_HOST);
+
+    for _ in 0..DEFAULT_MAX_IDLE_CONNECTIONS_PER_HOST {
+        let result = pool.remove(&poolkey);
+        assert!(result.is_some(), "expected key was not in pool");
+    }
+    assert_eq!(pool.len(), 0);
+}
+
+#[test]
+fn pool_update_connection_limit() {
+    let mut pool = ConnectionPool::new();
+    pool.set_max_idle_connections(50);
+
+    let hostnames = (0..pool.max_idle_connections).map(|i| format!("{}.example", i));
     let poolkeys = hostnames.map(|hostname| PoolKey {
         scheme: "https".to_string(),
         hostname,
@@ -159,15 +289,35 @@ fn pool_duplicates_limit() {
         proxy: None,
     });
     for key in poolkeys.clone() {
-        pool.add(key.clone(), Stream::Cursor(std::io::Cursor::new(vec![])));
         pool.add(key, Stream::Cursor(std::io::Cursor::new(vec![])));
     }
-    assert_eq!(pool.len(), 100);
+    assert_eq!(pool.len(), 50);
+    pool.set_max_idle_connections(25);
+    assert_eq!(pool.len(), 25);
+}
 
-    for key in poolkeys {
-        let result = pool.remove(&key);
-        assert!(result.is_some(), "expected key was not in pool");
+#[test]
+fn pool_update_per_host_connection_limit() {
+    let mut pool = ConnectionPool::new();
+    pool.set_max_idle_connections(50);
+    pool.set_max_idle_connections_per_host(50);
+
+    let poolkey = PoolKey {
+        scheme: "https".to_string(),
+        hostname: "example.com".to_string(),
+        port: Some(999),
+        proxy: None,
+    };
+
+    for _ in 0..50 {
+        pool.add(
+            poolkey.clone(),
+            Stream::Cursor(std::io::Cursor::new(vec![])),
+        );
     }
+    assert_eq!(pool.len(), 50);
+    pool.set_max_idle_connections_per_host(25);
+    assert_eq!(pool.len(), 25);
 }
 
 #[test]
diff --git a/src/request.rs b/src/request.rs
index 76b67d7..8a8b44e 100644
--- a/src/request.rs
+++ b/src/request.rs
@@ -1,16 +1,12 @@
+use std::fmt;
 use std::io::Read;
+use std::result::Result;
 use std::sync::{Arc, Mutex};
 use std::time;
 
 use qstring::QString;
 use url::{form_urlencoded, Url};
 
-#[cfg(feature = "tls")]
-use std::fmt;
-
-#[cfg(all(feature = "native-tls", not(feature = "tls")))]
-use std::fmt;
-
 use crate::agent::{self, Agent, AgentState};
 use crate::body::{Payload, SizedReader};
 use crate::error::Error;
@@ -54,8 +50,8 @@ pub struct Request {
     pub(crate) tls_connector: Option<TLSConnector>,
 }
 
-impl ::std::fmt::Debug for Request {
-    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::result::Result<(), ::std::fmt::Error> {
+impl fmt::Debug for Request {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let (path, query) = self
             .to_url()
             .map(|u| {
@@ -462,32 +458,6 @@ impl Request {
         self
     }
 
-    // pub fn retry(&self, times: u16) -> Request {
-    //     unimplemented!()
-    // }
-    // pub fn sortQuery(&self) -> Request {
-    //     unimplemented!()
-    // }
-    // pub fn sortQueryBy(&self, by: Box<Fn(&str, &str) -> usize>) -> Request {
-    //     unimplemented!()
-    // }
-    // pub fn ca<S>(&self, accept: S) -> Request
-    //     where S: Into<String> {
-    //     unimplemented!()
-    // }
-    // pub fn cert<S>(&self, accept: S) -> Request
-    //     where S: Into<String> {
-    //     unimplemented!()
-    // }
-    // pub fn key<S>(&self, accept: S) -> Request
-    //     where S: Into<String> {
-    //     unimplemented!()
-    // }
-    // pub fn pfx<S>(&self, accept: S) -> Request // TODO what type? u8?
-    //     where S: Into<String> {
-    //     unimplemented!()
-    // }
-
     /// Get the method this request is using.
     ///
     /// Example:
@@ -608,7 +578,7 @@ impl Request {
     ///
     /// Example:
     /// ```
-    /// let tls_connector = std::sync::Arc::new(native_tls::TlsConnector::new());
+    /// let tls_connector = std::sync::Arc::new(native_tls::TlsConnector::new().unwrap());
     /// let req = ureq::post("https://cool.server")
     ///     .set_tls_connector(tls_connector.clone());
     /// ```
diff --git a/src/response.rs b/src/response.rs
index 2b384d3..4e2d1c2 100644
--- a/src/response.rs
+++ b/src/response.rs
@@ -1,3 +1,4 @@
+use std::fmt;
 use std::io::{Cursor, Error as IoError, ErrorKind, Read, Result as IoResult};
 use std::str::FromStr;
 use std::time::Instant;
@@ -65,8 +66,8 @@ struct ResponseStatusIndex {
     response_code: usize,
 }
 
-impl ::std::fmt::Debug for Response {
-    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::result::Result<(), ::std::fmt::Error> {
+impl fmt::Debug for Response {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(
             f,
             "Response[status: {}, status_text: {}]",
@@ -593,14 +594,16 @@ pub(crate) fn set_stream(resp: &mut Response, url: String, unit: Option<Unit>, s
 fn read_next_line<R: Read>(reader: &mut R) -> IoResult<String> {
     let mut buf = Vec::new();
     let mut prev_byte_was_cr = false;
+    let mut one = [0_u8];
 
     loop {
-        let byte = reader.bytes().next();
+        let amt = reader.read(&mut one[..])?;
 
-        let byte = match byte {
-            Some(b) => b?,
-            None => return Err(IoError::new(ErrorKind::ConnectionAborted, "Unexpected EOF")),
-        };
+        if amt == 0 {
+            return Err(IoError::new(ErrorKind::ConnectionAborted, "Unexpected EOF"));
+        }
+
+        let byte = one[0];
 
         if byte == b'\n' && prev_byte_was_cr {
             buf.pop(); // removing the '\r'
diff --git a/src/stream.rs b/src/stream.rs
index 34a685f..b0f8309 100644
--- a/src/stream.rs
+++ b/src/stream.rs
@@ -1,4 +1,7 @@
-use std::io::{Cursor, Error as IoError, ErrorKind, Read, Result as IoResult, Write};
+use std::fmt;
+use std::io::{
+    BufRead, BufReader, Cursor, Error as IoError, ErrorKind, Read, Result as IoResult, Write,
+};
 use std::net::SocketAddr;
 use std::net::TcpStream;
 use std::net::ToSocketAddrs;
@@ -25,14 +28,14 @@ use crate::unit::Unit;
 
 #[allow(clippy::large_enum_variant)]
 pub enum Stream {
-    Http(TcpStream),
+    Http(BufReader<TcpStream>),
     #[cfg(all(feature = "tls", not(feature = "native-tls")))]
-    Https(rustls::StreamOwned<rustls::ClientSession, TcpStream>),
+    Https(BufReader<rustls::StreamOwned<rustls::ClientSession, TcpStream>>),
     #[cfg(all(feature = "native-tls", not(feature = "tls")))]
-    Https(TlsStream<TcpStream>),
+    Https(BufReader<TlsStream<TcpStream>>),
     Cursor(Cursor<Vec<u8>>),
     #[cfg(test)]
-    Test(Box<dyn Read + Send>, Vec<u8>),
+    Test(Box<dyn BufRead + Send>, Vec<u8>),
 }
 
 // DeadlineStream wraps a stream such that read() will return an error
@@ -100,8 +103,8 @@ pub(crate) fn io_err_timeout(error: String) -> IoError {
     IoError::new(ErrorKind::TimedOut, error)
 }
 
-impl ::std::fmt::Debug for Stream {
-    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::result::Result<(), ::std::fmt::Error> {
+impl fmt::Debug for Stream {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(
             f,
             "Stream[{}]",
@@ -161,9 +164,9 @@ impl Stream {
 
     pub(crate) fn socket(&self) -> Option<&TcpStream> {
         match self {
-            Stream::Http(tcpstream) => Some(tcpstream),
+            Stream::Http(b) => Some(b.get_ref()),
             #[cfg(feature = "tls")]
-            Stream::Https(rustls_stream) => Some(&rustls_stream.sock),
+            Stream::Https(b) => Some(&b.get_ref().sock),
             _ => None,
         }
     }
@@ -193,6 +196,36 @@ impl Read for Stream {
     }
 }
 
+impl BufRead for Stream {
+    fn fill_buf(&mut self) -> IoResult<&[u8]> {
+        match self {
+            Stream::Http(r) => r.fill_buf(),
+            #[cfg(any(
+                all(feature = "tls", not(feature = "native-tls")),
+                all(feature = "native-tls", not(feature = "tls")),
+            ))]
+            Stream::Https(r) => r.fill_buf(),
+            Stream::Cursor(r) => r.fill_buf(),
+            #[cfg(test)]
+            Stream::Test(r, _) => r.fill_buf(),
+        }
+    }
+
+    fn consume(&mut self, amt: usize) {
+        match self {
+            Stream::Http(r) => r.consume(amt),
+            #[cfg(any(
+                all(feature = "tls", not(feature = "native-tls")),
+                all(feature = "native-tls", not(feature = "tls")),
+            ))]
+            Stream::Https(r) => r.consume(amt),
+            Stream::Cursor(r) => r.consume(amt),
+            #[cfg(test)]
+            Stream::Test(r, _) => r.consume(amt),
+        }
+    }
+}
+
 impl<R: Read> From<ChunkDecoder<R>> for Stream
 where
     R: Read,
@@ -205,7 +238,7 @@ where
 
 #[cfg(all(feature = "tls", not(feature = "native-tls")))]
 fn read_https(
-    stream: &mut StreamOwned<ClientSession, TcpStream>,
+    stream: &mut BufReader<StreamOwned<ClientSession, TcpStream>>,
     buf: &mut [u8],
 ) -> IoResult<usize> {
     match stream.read(buf) {
@@ -216,7 +249,7 @@ fn read_https(
 }
 
 #[cfg(all(feature = "native-tls", not(feature = "tls")))]
-fn read_https(stream: &mut TlsStream<TcpStream>, buf: &mut [u8]) -> IoResult<usize> {
+fn read_https(stream: &mut BufReader<TlsStream<TcpStream>>, buf: &mut [u8]) -> IoResult<usize> {
     match stream.read(buf) {
         Ok(size) => Ok(size),
         Err(ref e) if is_close_notify(e) => Ok(0),
@@ -243,12 +276,12 @@ fn is_close_notify(e: &std::io::Error) -> bool {
 impl Write for Stream {
     fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
         match self {
-            Stream::Http(sock) => sock.write(buf),
+            Stream::Http(sock) => sock.get_mut().write(buf),
             #[cfg(any(
                 all(feature = "tls", not(feature = "native-tls")),
                 all(feature = "native-tls", not(feature = "tls")),
             ))]
-            Stream::Https(stream) => stream.write(buf),
+            Stream::Https(stream) => stream.get_mut().write(buf),
             Stream::Cursor(_) => panic!("Write to read only stream"),
             #[cfg(test)]
             Stream::Test(_, writer) => writer.write(buf),
@@ -256,12 +289,12 @@ impl Write for Stream {
     }
     fn flush(&mut self) -> IoResult<()> {
         match self {
-            Stream::Http(sock) => sock.flush(),
+            Stream::Http(sock) => sock.get_mut().flush(),
             #[cfg(any(
                 all(feature = "tls", not(feature = "native-tls")),
                 all(feature = "native-tls", not(feature = "tls")),
             ))]
-            Stream::Https(stream) => stream.flush(),
+            Stream::Https(stream) => stream.get_mut().flush(),
             Stream::Cursor(_) => panic!("Flush read only stream"),
             #[cfg(test)]
             Stream::Test(_, writer) => writer.flush(),
@@ -274,7 +307,9 @@ pub(crate) fn connect_http(unit: &Unit) -> Result<Stream, Error> {
     let hostname = unit.url.host_str().unwrap();
     let port = unit.url.port().unwrap_or(80);
 
-    connect_host(unit, hostname, port).map(Stream::Http)
+    connect_host(unit, hostname, port)
+        .map(BufReader::new)
+        .map(Stream::Http)
 }
 
 #[cfg(all(feature = "tls", feature = "native-certs"))]
@@ -316,7 +351,7 @@ pub(crate) fn connect_https(unit: &Unit) -> Result<Stream, Error> {
 
     let stream = rustls::StreamOwned::new(sess, sock);
 
-    Ok(Stream::Https(stream))
+    Ok(Stream::Https(BufReader::new(stream)))
 }
 
 #[cfg(all(feature = "native-tls", not(feature = "tls")))]
@@ -338,7 +373,7 @@ pub(crate) fn connect_https(unit: &Unit) -> Result<Stream, Error> {
             _ => Error::BadStatusRead,
         })?;
 
-    Ok(Stream::Https(stream))
+    Ok(Stream::Https(BufReader::new(stream)))
 }
 
 pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<TcpStream, Error> {
diff --git a/src/unit.rs b/src/unit.rs
index 37530e8..a22ddd9 100644
--- a/src/unit.rs
+++ b/src/unit.rs
@@ -369,7 +369,7 @@ fn send_prelude(unit: &Unit, stream: &mut Stream, redir: bool) -> IoResult<()> {
         }
     }
     if !header::has_header(&unit.headers, "user-agent") {
-        write!(prelude, "User-Agent: ureq\r\n")?;
+        write!(prelude, "User-Agent: ureq/{}\r\n", env!("CARGO_PKG_VERSION"))?;
     }
     if !header::has_header(&unit.headers, "accept") {
         write!(prelude, "Accept: */*\r\n")?;


Test patch:
diff --git a/src/test/agent_test.rs b/src/test/agent_test.rs
index 66356cd..58bdd31 100644
--- a/src/test/agent_test.rs
+++ b/src/test/agent_test.rs
@@ -99,3 +99,50 @@ fn connection_reuse() {
     }
     assert_eq!(resp.status(), 200);
 }
+
+#[cfg(test)]
+fn cookie_and_redirect(mut stream: TcpStream) -> io::Result<()> {
+    let headers = read_headers(&stream);
+    match headers.path() {
+        "/first" => {
+            stream.write_all(b"HTTP/1.1 302 Found\r\n")?;
+            stream.write_all(b"Location: /second\r\n")?;
+            stream.write_all(b"Set-Cookie: first=true\r\n")?;
+            stream.write_all(b"Content-Length: 0\r\n\r\n")?;
+        },
+        "/second" => {
+            if headers.headers().iter().find(|&x| x.contains("Cookie: first")).is_none() {
+                panic!("request did not contain cookie 'first'");
+            }
+            stream.write_all(b"HTTP/1.1 302 Found\r\n")?;
+            stream.write_all(b"Location: /third\r\n")?;
+            stream.write_all(b"Set-Cookie: second=true\r\n")?;
+            stream.write_all(b"Content-Length: 0\r\n\r\n")?;
+        },
+        "/third" => {
+            if headers.headers().iter().find(|&x| x.contains("Cookie: first")).is_none() {
+                panic!("request did not contain cookie 'second'");
+            }
+            stream.write_all(b"HTTP/1.1 200 OK\r\n")?;
+            stream.write_all(b"Set-Cookie: third=true\r\n")?;
+            stream.write_all(b"Content-Length: 0\r\n\r\n")?;
+        },
+        _ => {},
+    }
+    Ok(())
+}
+
+#[cfg(feature = "cookie")]
+#[test]
+fn test_cookies_on_redirect() {
+    let testserver = TestServer::new(cookie_and_redirect);
+    let url = format!("http://localhost:{}/first", testserver.port);
+    let agent = Agent::default().build();
+    let resp = agent.post(&url).call();
+    if resp.error() {
+        panic!("error: {} {}", resp.status(), resp.into_string().unwrap());
+    }
+    assert!(agent.cookie("first").is_some());
+    assert!(agent.cookie("second").is_some());
+    assert!(agent.cookie("third").is_some());
+}
diff --git a/src/test/testserver.rs b/src/test/testserver.rs
index 7c563a1..45b70a8 100644
--- a/src/test/testserver.rs
+++ b/src/test/testserver.rs
@@ -9,14 +9,35 @@ pub struct TestServer {
     pub done: Arc<AtomicBool>,
 }
 
+pub struct TestHeaders(Vec<String>);
+
+impl TestHeaders {
+    // Return the path for a request, e.g. /foo from "GET /foo HTTP/1.1"
+    pub fn path(&self) -> &str {
+        if self.0.len() == 0 {
+            ""
+        } else {
+            &self.0[0].split(" ").nth(1).unwrap()
+        }
+    }
+
+    pub fn headers(&self) -> &[String] {
+        &self.0[1..]
+    }
+}
+
 // Read a stream until reaching a blank line, in order to consume
 // request headers.
-pub fn read_headers(stream: &TcpStream) {
+pub fn read_headers(stream: &TcpStream) -> TestHeaders {
+    let mut results = vec![];
     for line in BufReader::new(stream).lines() {
-        if line.unwrap() == "" {
-            break;
-        }
+        match line {
+            Err(e) => panic!(e),
+            Ok(line) if line == "" => break,
+            Ok(line) => results.push(line),
+        };
     }
+    TestHeaders(results)
 }
 
 impl TestServer {