fix: improve test coverage

Author: conblem

src/api/proxy.rs

@@ -5,7 +5,6 @@ use ppp::model::{Addresses, Header};
 use std::future::Future;
 use std::io::IoSlice;
 use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6};
-use std::ops::{Deref, DerefMut};
 use std::pin::Pin;
 use std::task::{Context, Poll};
 use tokio::io::{AsyncRead, AsyncWrite, Error as IoError, ErrorKind, ReadBuf, Result as IoResult};
@@ -27,6 +26,12 @@ impl RemoteAddr for TcpStream {
     }
 }
 
+impl<T: RemoteAddr> RemoteAddr for ProxyStream<T> {
+    fn remote_addr(&self) -> IoResult<SocketAddr> {
+        self.stream.remote_addr()
+    }
+}
+
 pub(super) fn wrap(
     listener: TcpListener,
     proxy: ProxyProtocol,
@@ -144,54 +149,40 @@ where
     }
 }
 
-impl<T> Deref for ProxyStream<T> {
-    type Target = T;
-
-    fn deref(&self) -> &Self::Target {
-        &self.stream
-    }
-}
-
-impl<T> DerefMut for ProxyStream<T> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.stream
-    }
-}
-
 struct RealAddrFuture<'a, T> {
     proxy_stream: &'a mut ProxyStream<T>,
 }
 
-impl<'a, T> RealAddrFuture<'a, T> {
-    fn format_header(&self, res: Header) -> IoResult<Option<SocketAddr>> {
-        let addr = match res.addresses {
-            Addresses::IPv4 {
-                source_address,
-                source_port,
-                ..
-            } => {
-                let port = source_port.unwrap_or_default();
-                SocketAddrV4::new(source_address.into(), port).into()
-            }
-            Addresses::IPv6 {
-                source_address,
-                source_port,
-                ..
-            } => {
-                let port = source_port.unwrap_or_default();
-                SocketAddrV6::new(source_address.into(), port, 0, 0).into()
-            }
-            address => {
-                return Err(IoError::new(
-                    ErrorKind::Other,
-                    format!("Cannot convert {:?} to a SocketAddr", address),
-                ))
-            }
-        };
+fn format_header(res: Header) -> IoResult<SocketAddr> {
+    let addr = match res.addresses {
+        Addresses::IPv4 {
+            source_address,
+            source_port,
+            ..
+        } => {
+            let port = source_port.unwrap_or_default();
+            SocketAddrV4::new(source_address.into(), port).into()
+        }
+        Addresses::IPv6 {
+            source_address,
+            source_port,
+            ..
+        } => {
+            let port = source_port.unwrap_or_default();
+            SocketAddrV6::new(source_address.into(), port, 0, 0).into()
+        }
+        address => {
+            return Err(IoError::new(
+                ErrorKind::Other,
+                format!("Cannot convert {:?} to a SocketAddr", address),
+            ))
+        }
+    };
 
-        Ok(Some(addr))
-    }
+    Ok(addr)
+}
 
+impl<'a, T> RealAddrFuture<'a, T> {
     fn get_header(&mut self) -> Poll<IoResult<Option<SocketAddr>>> {
         let data = match &mut self.proxy_stream.data {
             Some(data) => data,
@@ -212,7 +203,7 @@ impl<'a, T> RealAddrFuture<'a, T> {
             }
         };
 
-        Poll::Ready(self.format_header(res))
+        Poll::Ready(format_header(res).map(Some))
     }
 }
 
@@ -235,7 +226,7 @@ where
             Ok(0) => {
                 return Poll::Ready(Err(IoError::new(
                     ErrorKind::UnexpectedEof,
-                    "Streamed finished before end of proxy protocol header",
+                    "Stream finished before end of proxy protocol header",
                 )))
             }
             Ok(_) => {}
@@ -250,25 +241,79 @@ where
 mod tests {
     use ppp::model::{Addresses, Command, Header, Protocol, Version};
     use std::io::Cursor;
+    use std::net::SocketAddr;
+    use tokio::io::AsyncReadExt;
 
-    use super::ToProxyStream;
+    use super::{format_header, ToProxyStream};
     use crate::config::ProxyProtocol;
-    use std::net::SocketAddr;
 
     #[tokio::test]
-    async fn test_header_parsing() {
-        let header = Header::new(
+    async fn test_disabled() {
+        let mut proxy_stream = Cursor::new(vec![]).source(ProxyProtocol::Disabled);
+        assert!(proxy_stream.real_addr().await.unwrap().is_none());
+    }
+
+    fn generate_header(addresses: Addresses) -> Header {
+        Header::new(
             Version::Two,
             Command::Proxy,
             Protocol::Stream,
             vec![],
-            Addresses::from(([1, 1, 1, 1], [2, 2, 2, 2], 24034, 443)),
-        );
-        let header = ppp::to_bytes(header).unwrap();
+            addresses,
+        )
+    }
+
+    fn generate_ipv4() -> Header {
+        let adresses = Addresses::from(([1, 1, 1, 1], [2, 2, 2, 2], 24034, 443));
+        generate_header(adresses)
+    }
+
+    #[tokio::test]
+    async fn test_header_parsing() {
+        let mut header = ppp::to_bytes(generate_ipv4()).unwrap();
+        header.extend_from_slice("Test".as_ref());
         let mut header = Cursor::new(header).source(ProxyProtocol::Enabled);
 
         let actual = header.real_addr().await.unwrap().unwrap();
 
         assert_eq!(SocketAddr::from(([1, 1, 1, 1], 24034)), actual);
+
+        let mut actual = String::new();
+        let size = header.read_to_string(&mut actual).await.unwrap();
+        assert_eq!(4, size);
+        assert_eq!("Test", actual);
+    }
+
+    #[tokio::test]
+    async fn test_incomplete() {
+        let header = ppp::to_bytes(generate_ipv4()).unwrap();
+        let header = &mut &header[..10];
+        let mut header = header.source(ProxyProtocol::Enabled);
+
+        let actual = header.real_addr().await.unwrap_err();
+        assert_eq!(
+            format!("{}", actual),
+            "Stream finished before end of proxy protocol header"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_failure() {
+        let invalid = Vec::from("invalid header");
+        let invalid = &mut &invalid[..];
+
+        let mut invalid = invalid.source(ProxyProtocol::Enabled);
+
+        let actual = invalid.real_addr().await.unwrap_err();
+        assert_eq!(format!("{}", actual), "Proxy Parser Error");
+    }
+
+    #[test]
+    fn test_adresses() {
+        let address = [1, 1, 1, 1, 1, 1, 1, 1];
+        let addresses = Addresses::from((address, address, 24034, 443));
+
+        let actual = format_header(generate_header(addresses)).unwrap();
+        assert_eq!(SocketAddr::from((address, 24034)), actual);
     }
 }