chore(tcp): Add tests tcp proxy

src/socks5.rs

@@ -122,21 +122,21 @@ pub fn new_reply(error: &ReplyError, sock_addr: SocketAddr) -> Vec<u8> {
     reply
 }
 
-#[cfg(test)]
-mod test {
-    use super::*;
-    use futures_util::StreamExt;
-    use std::str::FromStr;
-
-    #[tokio::test]
-    async fn socks5_server() {
-        let mut x = run_server(SocketAddr::from_str("[::]:4343").unwrap())
-            .await
-            .unwrap();
-
-        loop {
-            let cnx = x.next().await.unwrap().unwrap();
-            eprintln!("{:?}", cnx);
-        }
-    }
-}
+//#[cfg(test)]
+//mod test {
+//    use super::*;
+//    use futures_util::StreamExt;
+//    use std::str::FromStr;
+//
+//    #[tokio::test]
+//    async fn socks5_server() {
+//        let mut x = run_server(SocketAddr::from_str("[::]:4343").unwrap())
+//            .await
+//            .unwrap();
+//
+//        loop {
+//            let cnx = x.next().await.unwrap().unwrap();
+//            eprintln!("{:?}", cnx);
+//        }
+//    }
+//}

src/tcp.rs

@@ -159,10 +159,11 @@ pub async fn connect_with_http_proxy(
         }
     }
 
-    static OK_RESPONSE: &[u8; 12] = b"HTTP/1.0 200";
+    static OK_RESPONSE_10: &[u8] = b"HTTP/1.0 200 ";
+    static OK_RESPONSE_11: &[u8] = b"HTTP/1.1 200 ";
     if !buf
-        .windows(OK_RESPONSE.len())
-        .any(|window| window == OK_RESPONSE)
+        .windows(OK_RESPONSE_10.len())
+        .any(|window| window == OK_RESPONSE_10 || window == OK_RESPONSE_11)
     {
         return Err(anyhow!(
             "Cannot connect to http proxy. Proxy returned an invalid response: {}",
@@ -182,3 +183,79 @@ pub async fn run_server(bind: SocketAddr) -> Result<TcpListenerStream, anyhow::E
         .with_context(|| format!("Cannot create TCP server {:?}", bind))?;
     Ok(TcpListenerStream::new(listener))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use futures_util::pin_mut;
+    use std::net::SocketAddr;
+    use testcontainers::core::WaitFor;
+    use testcontainers::{Image, ImageArgs, RunnableImage};
+
+    #[derive(Debug, Clone, Default)]
+    pub struct MitmProxy {}
+
+    impl ImageArgs for MitmProxy {
+        fn into_iterator(self) -> Box<dyn Iterator<Item = String>> {
+            Box::new(vec!["mitmdump".to_string()].into_iter())
+        }
+    }
+
+    impl Image for MitmProxy {
+        type Args = Self;
+
+        fn name(&self) -> String {
+            "mitmproxy/mitmproxy".to_string()
+        }
+
+        fn tag(&self) -> String {
+            "10.1.1".to_string()
+        }
+
+        fn ready_conditions(&self) -> Vec<WaitFor> {
+            vec![WaitFor::Duration {
+                length: Duration::from_secs(5),
+            }]
+        }
+    }
+
+    #[tokio::test]
+    async fn test_proxy_connection() {
+        let server_addr: SocketAddr = "[::1]:1236".parse().unwrap();
+        let server = TcpListener::bind(server_addr).await.unwrap();
+
+        let docker = testcontainers::clients::Cli::default();
+        let mitm_proxy: RunnableImage<MitmProxy> =
+            RunnableImage::from(MitmProxy {}).with_network("host".to_string());
+        let _node = docker.run(mitm_proxy);
+
+        let mut client = connect_with_http_proxy(
+            &"http://localhost:8080".parse().unwrap(),
+            &Host::Domain("[::1]".to_string()),
+            1236,
+            &None,
+            Duration::from_secs(1),
+        )
+        .await
+        .unwrap();
+
+        client
+            .write_all(b"GET / HTTP/1.1\r\n\r\n".as_slice())
+            .await
+            .unwrap();
+        let client_srv = server.accept().await.unwrap().0;
+        pin_mut!(client_srv);
+
+        let mut buf = [0u8; 25];
+        let ret = client_srv.read(&mut buf).await;
+        assert!(matches!(ret, Ok(18)));
+        client_srv
+            .write_all("HTTP/1.1 200 OK\r\n\r\n".as_bytes())
+            .await
+            .unwrap();
+
+        client_srv.get_mut().shutdown().await.unwrap();
+        let _ = client.read(&mut buf).await.unwrap();
+        assert!(buf.starts_with(b"HTTP/1.1 200 OK\r\n\r\n"));
+    }
+}

src/tunnel/io.rs

@@ -1,9 +1,7 @@
-
 use fastwebsockets::{Frame, OpCode, Payload, WebSocketError, WebSocketRead, WebSocketWrite};
 use futures_util::pin_mut;
 use hyper::upgrade::Upgraded;
 
-
 use std::time::Duration;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadHalf, WriteHalf};
 use tokio::select;

src/tunnel/server.rs

@@ -252,7 +252,7 @@ pub async fn run_server(server_config: Arc<WsServerConfig>) -> anyhow::Result<()
 
             let fut = async move {
                 if let Err(e) = conn_fut.await {
-                    error!("Error while upgrading cnx to weboscket: {:?}", e);
+                    error!("Error while upgrading cnx to websocket: {:?}", e);
                 }
             }
             .instrument(span);