fix(dgw) : improvement in network-scanner-net

crates/network-scanner-net/Cargo.toml

@@ -5,14 +5,15 @@ authors = ["Devolutions Inc. <[email protected]>"]
 edition = "2021"
 publish = false
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 [dependencies]
 anyhow = "1.0.79"
 crossbeam = { version = "0.8.4", features = ["crossbeam-channel"] }
+futures = "0.3.30"
 parking_lot = "0.12.1"
-polling = "3.3.2"
+polling = { git = "https://github.com/smol-rs/polling", rev = "62430fd56e668559d08ca7071ab13a0e116ba515" }
 socket2 = { version = "0.5.5", features = ["all"] }
 thiserror = "1.0.56"
+tokio-stream = "0.1.14"
 tracing = "0.1.40"
 
 [dev-dependencies]

crates/network-scanner-net/examples/intense_tcp.rs

@@ -0,0 +1,137 @@
+use std::{
+    mem::MaybeUninit,
+    net::{SocketAddr, SocketAddrV4},
+    sync::atomic::AtomicU32,
+    time::Instant,
+};
+
+use network_scanner_net::socket::AsyncRawSocket;
+use socket2::SockAddr;
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
+    task::JoinHandle,
+};
+
+#[tokio::main(flavor = "multi_thread", worker_threads = 12)]
+pub async fn main() -> anyhow::Result<()> {
+    tracing_subscriber::fmt::SubscriberBuilder::default()
+        .with_max_level(tracing::Level::DEBUG)
+        .with_thread_names(true)
+        .init();
+    // trace info all args
+    let args: Vec<String> = std::env::args().collect();
+    tracing::info!("Args: {:?}", std::env::args().collect::<Vec<String>>());
+    if args.len() < 4 {
+        println!("Usage: {} [server|client] -p <port>", args[0]);
+        return Ok(());
+    }
+
+    let port = args[args.len() - 1].parse::<u16>()?;
+    let addr = format!("127.0.0.1:{}", port);
+    match args[1].as_str() {
+        "server" => {
+            tcp_server(&addr).await?;
+        }
+        "client" => {
+            tcp_client().await?;
+        }
+        _ => {
+            println!("Usage: {} [server|client] -p <port>", args[0]);
+        }
+    }
+
+    Ok(())
+}
+
+async fn tcp_client() -> anyhow::Result<()> {
+    let async_runtime = network_scanner_net::runtime::Socket2Runtime::new(None)?;
+    let mut socket_arr = vec![];
+    for _ in 0..100 {
+        tracing::info!("Creating socket");
+        let socket = async_runtime.new_socket(socket2::Domain::IPV4, socket2::Type::STREAM, None)?;
+        socket_arr.push(socket);
+    }
+
+    fn connect_and_write_something(
+        mut socket: AsyncRawSocket,
+        addr: std::net::SocketAddr,
+    ) -> JoinHandle<anyhow::Result<()>> {
+        tokio::task::spawn(async move {
+            let addr: SockAddr = addr.into();
+            socket.connect(&addr).await?;
+            let mut buffer = [MaybeUninit::uninit(); 1024];
+            for i in 0..1000 {
+                let data = format!("hello world {} times", i);
+                tracing::info!("Sending: {} from socket {:?}", &data, &socket);
+                let write_future = socket.send(data.as_bytes());
+
+                let size = tokio::time::timeout(std::time::Duration::from_secs(1), write_future).await??;
+
+                if size == 0 {
+                    return Ok(());
+                }
+
+                let recv_future = socket.recv(&mut buffer);
+                tokio::time::timeout(std::time::Duration::from_secs(1), recv_future).await??;
+                let received = buffer[..size]
+                    .iter()
+                    .map(|x| unsafe { x.assume_init() })
+                    .collect::<Vec<u8>>();
+                assert_eq!(received, data.as_bytes());
+                tracing::debug!("Received: {}", std::str::from_utf8(&received)?);
+            }
+            Ok(())
+        })
+    }
+
+    let mut handles = vec![];
+    for socket in socket_arr {
+        let addr: SocketAddr = SocketAddrV4::new(std::net::Ipv4Addr::new(127, 0, 0, 1), 8080).into();
+
+        let handle = connect_and_write_something(socket, addr);
+        handles.push(handle);
+    }
+
+    for future in handles {
+        future.await??;
+    }
+
+    Ok(())
+}
+
+async fn tcp_server(addr: &str) -> anyhow::Result<()> {
+    let listener = tokio::net::TcpListener::bind(addr).await?;
+    println!("Listening on: {}", addr);
+    let count = std::sync::Arc::new(AtomicU32::new(0));
+    loop {
+        let (mut socket, _) = listener.accept().await?;
+        let _now = Instant::now();
+        let count = count.clone();
+        tokio::spawn(async move {
+            let mut buf = vec![0; 1024];
+            loop {
+                let n = match socket.read(&mut buf).await {
+                    // socket closed
+                    Ok(n) => {
+                        if n == 0 {
+                            println!("Socket closed");
+                            return;
+                        }
+                        n
+                    }
+                    Err(e) => {
+                        eprintln!("Failed to read from socket; err = {:?}", e);
+                        return;
+                    }
+                };
+                println!("Received {} bytes", n);
+                // Write the data back
+                if let Err(e) = socket.write_all(&buf[0..n]).await {
+                    eprintln!("Failed to write to socket; err = {:?}", e);
+                    return;
+                }
+                count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+            }
+        });
+    }
+}

crates/network-scanner-net/examples/tcp_fail.rs

@@ -0,0 +1,83 @@
+use std::net::SocketAddr;
+
+#[tokio::main]
+pub async fn main() -> anyhow::Result<()> {
+    utils::start_server();
+    tracing_subscriber::fmt::SubscriberBuilder::default()
+        .with_max_level(tracing::Level::TRACE)
+        .with_thread_names(true)
+        .init();
+
+    let async_runtime = network_scanner_net::runtime::Socket2Runtime::new(None)?;
+    let good_socket = async_runtime.new_socket(socket2::Domain::IPV4, socket2::Type::STREAM, None)?;
+    let bad_socket = async_runtime.new_socket(socket2::Domain::IPV4, socket2::Type::STREAM, None)?;
+
+    tracing::info!("good_socket: {:?}", good_socket);
+    tracing::info!("bad_socket: {:?}", bad_socket);
+
+    let available_port = 8080;
+    let non_available_port = 12345;
+
+    let available_addr = SocketAddr::from(([127, 0, 0, 1], available_port));
+    let non_available_addr = SocketAddr::from(([127, 0, 0, 1], non_available_port));
+
+    let handle = tokio::task::spawn(async move { good_socket.connect(&socket2::SockAddr::from(available_addr)).await });
+
+    let handle2 =
+        tokio::task::spawn(async move { bad_socket.connect(&socket2::SockAddr::from(non_available_addr)).await });
+
+    let (a, b) = tokio::join!(handle, handle2);
+    // remove the outer error from tokio task
+    let a = a?;
+    let b = b?;
+    tracing::info!("should connect: {:?}", &a);
+    tracing::info!("should not connect: {:?}", &b);
+    assert!(a.is_ok());
+    assert!(b.is_err());
+    Ok(())
+}
+
+mod utils {
+    use std::io::Read;
+    use std::io::Write;
+    use std::net::{TcpListener, TcpStream};
+    use std::thread;
+
+    fn handle_client(mut stream: TcpStream) {
+        // read 20 bytes at a time from stream echoing back to stream
+        loop {
+            let mut read = [0; 1028];
+            match stream.read(&mut read) {
+                Ok(n) => {
+                    if n == 0 {
+                        // connection was closed
+                        break;
+                    }
+                    let _ = stream.write(&read[0..n]).unwrap();
+                }
+                Err(_) => {
+                    return;
+                }
+            }
+        }
+    }
+
+    pub(super) fn start_server() {
+        thread::spawn(|| {
+            let listener = TcpListener::bind("127.0.0.1:8080").unwrap();
+
+            for stream in listener.incoming() {
+                match stream {
+                    Ok(stream) => {
+                        thread::spawn(move || {
+                            handle_client(stream);
+                        });
+                    }
+                    Err(_) => {
+                        // println!("Error");
+                    }
+                }
+            }
+        });
+    }
+}