Add Response::remote_addr()

src/response.rs

@@ -1,4 +1,5 @@
 use std::io::{self, Read};
+use std::net::SocketAddr;
 use std::str::FromStr;
 use std::sync::Mutex;
 use std::{fmt, io::BufRead};
@@ -70,6 +71,8 @@ pub struct Response {
     unit: Box<Unit>,
     // Boxed to avoid taking up too much size.
     stream: Mutex<Box<Stream>>,
+    /// The socket address of the server that sent the response.
+    remote_addr: SocketAddr,
     /// The redirect history of this response, if any. The history starts with
     /// the first response received and ends with the response immediately
     /// previous to this one.
@@ -224,6 +227,11 @@ impl Response {
         charset_from_content_type(self.header("content-type"))
     }
 
+    /// The socket address of the server that sent the response.
+    pub fn remote_addr(&self) -> SocketAddr {
+        self.remote_addr
+    }
+
     /// Turn this response into a `impl Read` of the body.
     ///
     /// 1. If `Transfer-Encoding: chunked`, the returned reader will unchunk it
@@ -470,6 +478,7 @@ impl Response {
     ///
     /// assert_eq!(resp.status(), 401);
     pub(crate) fn do_from_stream(stream: Stream, unit: Unit) -> Result<Response, Error> {
+        let remote_addr = stream.remote_addr;
         //
         // HTTP/1.1 200 OK\r\n
         let mut stream = stream::DeadlineStream::new(stream, unit.deadline);
@@ -515,6 +524,7 @@ impl Response {
             headers,
             unit: Box::new(unit),
             stream: Mutex::new(Box::new(stream.into())),
+            remote_addr,
             history: vec![],
             length,
             compression,

src/stream.rs

@@ -48,6 +48,8 @@ pub trait TlsConnector: Send + Sync {
 
 pub(crate) struct Stream {
     inner: BufReader<Box<dyn ReadWrite>>,
+    /// The remote address the stream is connected to.
+    pub(crate) remote_addr: SocketAddr,
 }
 
 impl<T: ReadWrite + ?Sized> ReadWrite for Box<T> {
@@ -197,9 +199,10 @@ impl fmt::Debug for Stream {
 }
 
 impl Stream {
-    fn new(t: impl ReadWrite) -> Stream {
+    fn new(t: impl ReadWrite, remote_addr: SocketAddr) -> Stream {
         Stream::logged_create(Stream {
             inner: BufReader::new(Box::new(t)),
+            remote_addr,
         })
     }
 
@@ -215,19 +218,15 @@ impl Stream {
                 vec![],
                 false,
             ))),
+            remote_addr: remote_addr_for_test(),
         })
     }
 
     #[cfg(test)]
     pub(crate) fn from_vec_poolable(v: Vec<u8>) -> Stream {
         Stream::logged_create(Stream {
             inner: BufReader::new(Box::new(TestStream(Box::new(Cursor::new(v)), vec![], true))),
-        })
-    }
-
-    fn from_tcp_stream(t: TcpStream) -> Stream {
-        Stream::logged_create(Stream {
-            inner: BufReader::new(Box::new(t)),
+            remote_addr: remote_addr_for_test(),
         })
     }
 
@@ -348,20 +347,25 @@ pub(crate) fn connect_http(unit: &Unit, hostname: &str) -> Result<Stream, Error>
     //
     let port = unit.url.port().unwrap_or(80);
 
-    connect_host(unit, hostname, port).map(Stream::from_tcp_stream)
+    connect_host(unit, hostname, port).map(|(t, r)| Stream::new(t, r))
 }
 
 pub(crate) fn connect_https(unit: &Unit, hostname: &str) -> Result<Stream, Error> {
     let port = unit.url.port().unwrap_or(443);
 
-    let sock = connect_host(unit, hostname, port)?;
+    let (sock, remote_addr) = connect_host(unit, hostname, port)?;
 
     let tls_conf = &unit.agent.config.tls_config;
     let https_stream = tls_conf.connect(hostname, Box::new(sock))?;
-    Ok(Stream::new(https_stream))
+    Ok(Stream::new(https_stream, remote_addr))
 }
 
-pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<TcpStream, Error> {
+/// If successful, returns a `TcpStream` and the remote address it is connected to.
+pub(crate) fn connect_host(
+    unit: &Unit,
+    hostname: &str,
+    port: u16,
+) -> Result<(TcpStream, SocketAddr), Error> {
     let connect_deadline: Option<Instant> =
         if let Some(timeout_connect) = unit.agent.config.timeout_connect {
             Instant::now().checked_add(timeout_connect)
@@ -388,7 +392,7 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
     let proto = proxy.as_ref().map(|proxy| proxy.proto);
 
     let mut any_err = None;
-    let mut any_stream = None;
+    let mut any_stream_and_addr = None;
     // Find the first sock_addr that accepts a connection
     for sock_addr in sock_addrs {
         // ensure connect timeout or overall timeout aren't yet hit.
@@ -417,15 +421,15 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
         };
 
         if let Ok(stream) = stream {
-            any_stream = Some(stream);
+            any_stream_and_addr = Some((stream, sock_addr));
             break;
         } else if let Err(err) = stream {
             any_err = Some(err);
         }
     }
 
-    let mut stream = if let Some(stream) = any_stream {
-        stream
+    let (mut stream, remote_addr) = if let Some(stream_and_addr) = any_stream_and_addr {
+        stream_and_addr
     } else if let Some(e) = any_err {
         return Err(ErrorKind::ConnectionFailed.msg("Connect error").src(e));
     } else {
@@ -466,7 +470,7 @@ pub(crate) fn connect_host(unit: &Unit, hostname: &str, port: u16) -> Result<Tcp
         }
     }
 
-    Ok(stream)
+    Ok((stream, remote_addr))
 }
 
 #[cfg(feature = "socks-proxy")]
@@ -652,6 +656,11 @@ pub(crate) fn connect_test(unit: &Unit) -> Result<Stream, Error> {
     Err(ErrorKind::UnknownScheme.msg(format!("unknown scheme '{}'", unit.url.scheme())))
 }
 
+pub(crate) fn remote_addr_for_test() -> SocketAddr {
+    use std::net::{Ipv4Addr, SocketAddrV4};
+    SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0).into()
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -707,7 +716,7 @@ mod tests {
         let recorder = ReadRecorder {
             reads: reads.clone(),
         };
-        let stream = Stream::new(recorder);
+        let stream = Stream::new(recorder, remote_addr_for_test());
         let mut deadline_stream = DeadlineStream::new(stream, None);
         let mut buf = [0u8; 1];
         for _ in 0..8193 {