Merge branch 'dev' of github.com:conblem/acme-dns-rust into dev

conblem · conblem · commit 8a7bfca3bb8d · 2021-02-09T17:20:46.000+01:00
diff --git a/src/api/proxy.rs b/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};
@@ -17,15 +16,35 @@ use tracing::{error, Instrument, Span};
 
 use crate::config::ProxyProtocol;
 
+trait RemoteAddr {
+    fn remote_addr(&self) -> IoResult<SocketAddr>;
+}
+
+impl RemoteAddr for TcpStream {
+    fn remote_addr(&self) -> IoResult<SocketAddr> {
+        self.peer_addr()
+    }
+}
+
+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,
-) -> impl Stream<Item = IoResult<impl Future<Output = IoResult<ProxyStream>>>> + Send {
+) -> impl Stream<
+    Item = IoResult<
+        impl Future<Output = IoResult<impl AsyncRead + AsyncWrite + Send + Unpin + 'static>>,
+    >,
+> + Send {
     TcpListenerStream::new(listener)
         .map_ok(move |conn| conn.source(proxy))
         .map_ok(|mut conn| {
             let span = Span::current();
-            span.record("remote.addr", &debug(conn.peer_addr()));
+            span.record("remote.addr", &debug(conn.remote_addr()));
             let span_clone = span.clone();
 
             async move {
@@ -44,12 +63,15 @@ pub(super) fn wrap(
         })
 }
 
-trait ToProxyStream {
-    fn source(self, proxy: ProxyProtocol) -> ProxyStream;
+trait ToProxyStream: Sized {
+    fn source(self, proxy: ProxyProtocol) -> ProxyStream<Self>;
 }
 
-impl ToProxyStream for TcpStream {
-    fn source(self, proxy: ProxyProtocol) -> ProxyStream {
+impl<T> ToProxyStream for T
+where
+    T: AsyncRead + Unpin,
+{
+    fn source(self, proxy: ProxyProtocol) -> ProxyStream<T> {
         let data = match proxy {
             ProxyProtocol::Enabled => Some(Default::default()),
             ProxyProtocol::Disabled => None,
@@ -62,19 +84,25 @@ impl ToProxyStream for TcpStream {
     }
 }
 
-pub(super) struct ProxyStream {
-    stream: TcpStream,
+pub(super) struct ProxyStream<T> {
+    stream: T,
     data: Option<Vec<u8>>,
     start_of_data: usize,
 }
 
-impl ProxyStream {
-    fn real_addr(&mut self) -> RealAddrFuture<'_> {
+impl<T> ProxyStream<T>
+where
+    T: AsyncRead + Unpin,
+{
+    fn real_addr(&mut self) -> RealAddrFuture<'_, T> {
         RealAddrFuture { proxy_stream: self }
     }
 }
 
-impl AsyncRead for ProxyStream {
+impl<T> AsyncRead for ProxyStream<T>
+where
+    T: AsyncRead + Unpin,
+{
     fn poll_read(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
@@ -88,7 +116,10 @@ impl AsyncRead for ProxyStream {
     }
 }
 
-impl AsyncWrite for ProxyStream {
+impl<T> AsyncWrite for ProxyStream<T>
+where
+    T: AsyncWrite + Unpin,
+{
     fn poll_write(
         mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
@@ -118,55 +149,41 @@ impl AsyncWrite for ProxyStream {
     }
 }
 
-impl Deref for ProxyStream {
-    type Target = TcpStream;
-
-    fn deref(&self) -> &Self::Target {
-        &self.stream
-    }
+struct RealAddrFuture<'a, T> {
+    proxy_stream: &'a mut ProxyStream<T>,
 }
 
-impl DerefMut for ProxyStream {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.stream
-    }
-}
+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),
+            ))
+        }
+    };
 
-struct RealAddrFuture<'a> {
-    proxy_stream: &'a mut ProxyStream,
+    Ok(addr)
 }
 
-impl<'a> RealAddrFuture<'a> {
-    fn format_header(&self, res: Header) -> <Self as Future>::Output {
-        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))
-    }
-
-    fn get_header(&mut self) -> Poll<<Self as Future>::Output> {
+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,
             None => unreachable!("Future cannot be pulled anymore"),
@@ -186,11 +203,14 @@ impl<'a> RealAddrFuture<'a> {
             }
         };
 
-        Poll::Ready(self.format_header(res))
+        Poll::Ready(format_header(res).map(Some))
     }
 }
 
-impl Future for RealAddrFuture<'_> {
+impl<T> Future for RealAddrFuture<'_, T>
+where
+    T: AsyncRead + Unpin,
+{
     type Output = IoResult<Option<SocketAddr>>;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
@@ -206,7 +226,7 @@ impl Future for RealAddrFuture<'_> {
             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(_) => {}
@@ -216,3 +236,84 @@ impl Future for RealAddrFuture<'_> {
         this.get_header()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use ppp::model::{Addresses, Command, Header, Protocol, Version};
+    use std::io::Cursor;
+    use std::net::SocketAddr;
+    use tokio::io::AsyncReadExt;
+
+    use super::{format_header, ToProxyStream};
+    use crate::config::ProxyProtocol;
+
+    #[tokio::test]
+    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,
+        )
+    }
+
+    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);
+    }
+}
diff --git a/src/api/tls.rs b/src/api/tls.rs
@@ -11,11 +11,10 @@ use tokio::io::{AsyncRead, AsyncWrite, Result as IoResult};
 use tokio_rustls::TlsAcceptor;
 use tracing::{error, info};
 
-use super::proxy::ProxyStream;
 use crate::cert::{Cert, CertFacade};
 use crate::util::to_u64;
 
-pub(super) fn wrap<L, I>(
+pub(super) fn wrap<L, I, S>(
     listener: L,
     pool: PgPool,
 ) -> impl Stream<
@@ -25,7 +24,8 @@ pub(super) fn wrap<L, I>(
 > + Send
 where
     L: Stream<Item = IoResult<I>> + Send,
-    I: Future<Output = IoResult<ProxyStream>> + Send,
+    I: Future<Output = IoResult<S>> + Send,
+    S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
 {
     let acceptor = Acceptor::new(pool);
 
