Split off connection module

isototest/Cargo.toml

@@ -8,5 +8,6 @@ repository = "https://github.com/ByteOtter/isotest-ng/tree/main/isototest"
 license = "GPL-2.0"
 
 [dependencies]
+mockito = "1.4.0"
 tokio = "1.38.1"
 vnc-rs = "0.5.1"

isototest/src/connection.rs

@@ -0,0 +1,23 @@
+use tokio::{self, net::TcpStream};
+use vnc::{PixelFormat, VncClient, VncConnector, VncError};
+
+pub async fn create_vnc_client(
+    target_ip: String,
+    psw: Option<String>,
+) -> Result<VncClient, VncError> {
+    let tcp: TcpStream = TcpStream::connect(target_ip).await?;
+    let vnc: VncClient = VncConnector::new(tcp)
+        .set_auth_method(async move { Ok(psw.unwrap()) })
+        .add_encoding(vnc::VncEncoding::Tight)
+        .add_encoding(vnc::VncEncoding::Zrle)
+        .add_encoding(vnc::VncEncoding::CopyRect)
+        .add_encoding(vnc::VncEncoding::Raw)
+        .allow_shared(true)
+        .set_pixel_format(PixelFormat::bgra())
+        .build()?
+        .try_start()
+        .await?
+        .finish()?;
+
+    Ok(vnc)
+}

isototest/src/lib.rs

@@ -1,33 +1 @@
-use vnc::{PixelFormat, VncClient, VncConnector, VncError};
-use tokio::{self, net::TcpStream};
-
-pub async fn create_vnc_client(target_ip: String, psw: Option<String>) -> Result<VncClient, VncError> {
-    let tcp: TcpStream = TcpStream::connect(target_ip).await?;
-    let vnc: VncClient = VncConnector::new(tcp)
-        .set_auth_method(async move { Ok(psw.unwrap())})
-        .add_encoding(vnc::VncEncoding::Tight)
-        .add_encoding(vnc::VncEncoding::Zrle)
-        .add_encoding(vnc::VncEncoding::CopyRect)
-        .add_encoding(vnc::VncEncoding::Raw)
-        .allow_shared(true)
-        .set_pixel_format(PixelFormat::bgra())
-        .build()?
-        .try_start()
-        .await?
-        .finish()?;
-
-    Ok(vnc)
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use std::any::type_name_of_val;
-
-
-    #[tokio::test]
-    async fn test_build_client() {
-        let res = create_vnc_client("192.168.0.1".to_string(), Some("Hello".to_string()));
-        assert!(type_name_of_val(&res).contains("VncClient"));
-   }
-}
+pub mod connection;

isototest/tests/common.rs

@@ -0,0 +1,116 @@
+//! This module contains common setup code required for the testuite.
+use std::io;
+
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
+    net::{TcpListener, TcpStream},
+};
+use vnc::VncError;
+
+/// Mock VNC server.
+pub async fn start_mock_vnc_srv() -> Result<TcpListener, VncError> {
+    eprintln!("Starting mock VNC server...");
+    let srv = match TcpListener::bind("127.0.0.1:0").await {
+        Ok(srv) => {
+            eprintln!("Mock VNC server started on {:?}", srv.local_addr());
+            Ok(srv)
+        },
+        Err(e) => return Err(VncError::IoError(e)),
+    };
+    srv
+}
+
+// TODO: Implement mock vnc server handshake.
+pub async fn mock_vnc_handshake(mut socket: TcpStream) -> Result<(), io::Error> {
+    eprintln!("Starting VNC handshake...");
+
+    // Send the protocol version response
+    let response = b"RFB 003.003\n";
+    eprintln!("VNC Version response created '{:?}'.", response);
+    if let Err(e) = socket.write_all(response).await {
+        eprintln!("Failed to write protocol version response: {}", e);
+        return Err(e);
+    }
+
+    eprintln!("Sent protocol version response.");
+    // Read the client's protocol version message (12 bytes)
+    let mut buff = [0; 12];
+    eprintln!("First message buffer created '{:?}'.", &buff);
+    match socket.read_exact(&mut buff).await {
+        Ok(_) => eprintln!("Received protocol version: {:?}", &buff),
+        Err(e) => {
+            eprintln!("Failed to read protocol version: {}", e);
+            return Err(e);
+        }
+    }
+
+    // Verify the protocol version
+    let version_message = String::from_utf8_lossy(&buff);
+    if !version_message.starts_with("RFB 003.003") {
+        return Err(io::Error::new(
+            io::ErrorKind::InvalidData,
+            "Unsupported VNC protocol version",
+        ));
+    }
+
+    // Send the security type response (0 for successful authentication)
+    let auth_response = b"\0\0\0\0\0\0\0\0";
+    eprintln!("Created security response buffer '{:?}'.", &auth_response);
+    if let Err(e) = socket.write_all(auth_response).await {
+        eprintln!("Failed to write security type response: {}", e);
+        return Err(e);
+    }
+    eprintln!("Sent security type response.");
+
+
+    // Read the client's security type message (8 bytes)
+    let mut buf = [0; 8];
+    eprintln!("Created security receiving buffer '{:?}'.", &buf);
+    match socket.read_exact(&mut buf).await {
+        Ok(_) => eprintln!("Received security type: {:?}", &buf),
+        Err(e) => {
+            eprintln!("Failed to read security type: {}", e);
+            return Err(e);
+        }
+    }
+
+    // Check if the security type is supported (0 for no authentication)
+    if buf[0] != 0 {
+        return Err(io::Error::new(
+            io::ErrorKind::InvalidData,
+            "Unsupported security type",
+        ));
+    }
+
+    // Send the initialization response (example response)
+    let init_response = [0; 4];
+    if let Err(e) = socket.write_all(&init_response).await {
+        eprintln!("Failed to write initialization response: {}", e);
+        return Err(e);
+    }
+    eprintln!("Sent initialization response.");
+
+    // Read the client's initialization message (16 bytes as an example)
+    let mut buf = [0; 16];
+    match socket.read_exact(&mut buf).await {
+        Ok(_) => eprintln!("Received initialization message: {:?}", &buf),
+        Err(e) => {
+            eprintln!("Failed to read initialization message: {}", e);
+            return Err(e);
+        }
+    }
+
+    // Properly shutdown the connection
+    if let Err(e) = socket.shutdown().await {
+        eprintln!("Failed to shutdown the connection: {}", e);
+        return Err(e);
+    }
+    eprintln!("Connection shutdown successfully.");
+
+    Ok(())
+}
+
+/// Kill server connection.
+pub async fn kill_connection(mut socket: tokio::net::TcpStream) {
+    let _ = socket.shutdown();
+}

isototest/tests/test.rs

@@ -0,0 +1,82 @@
+use std::net::SocketAddr;
+use tokio::net::TcpListener;
+use vnc::client;
+use vnc::PixelFormat;
+use vnc::VncError;
+
+use isototest::connection::create_vnc_client;
+mod common;
+
+#[tokio::test]
+async fn test_create_success() {
+    let srv = common::start_mock_vnc_srv()
+        .await
+        .expect("Failed to start mock VNC server");
+    let addr = srv.local_addr().unwrap().to_string();
+    let psw = Some("password".to_string());
+
+    // Spawn a task to handle the VNC handshake
+    let handshake_task = tokio::spawn(async move {
+        if let Ok((socket, _)) = srv.accept().await {
+            common::mock_vnc_handshake(socket)
+                .await
+                .expect("Failed during mock VNC handshake");
+        }
+    });
+
+    // Create the VNC client
+    let result = create_vnc_client(addr, psw).await;
+    match result {
+        Ok(_) => assert!(true),
+        Err(e) => panic!("{}", e),
+    };
+
+    // Await the handshake task to ensure it completes
+    handshake_task
+        .await
+        .expect("Failed to await handshake task");
+}
+
+#[tokio::test]
+async fn test_create_invalid_ip() {
+    let target_ip = "256.256.256.256:5900".to_string();
+    let psw = Some("pass".to_string());
+    let result = create_vnc_client(target_ip, psw).await;
+    assert!(matches!(result, Err(VncError::IoError(_))));
+}
+
+#[tokio::test]
+async fn test_create_no_pass() {
+    let srv = common::start_mock_vnc_srv()
+        .await
+        .expect("Failed to start mock VNC server");
+    let addr = srv.local_addr().unwrap().to_string();
+    let psw = None;
+
+    // Spawn a task to handle the VNC handshake
+    let handshake_task = tokio::spawn(async move {
+        if let Ok((socket, _)) = srv.accept().await {
+            common::mock_vnc_handshake(socket)
+                .await
+                .expect("Failed during mock VNC handshake");
+        }
+    });
+
+    // Create the VNC client
+    let result = create_vnc_client(addr, psw).await;
+    assert!(result.is_ok());
+
+    // Await the handshake task to ensure it completes
+    handshake_task
+        .await
+        .expect("Failed to await handshake task");
+}
+
+#[tokio::test]
+async fn test_create_connect_fail() {
+    let target_ip = "127.0.0.1:9999".to_string();
+    let psw = Some("password".to_string());
+
+    let result = create_vnc_client(target_ip, psw).await;
+    assert!(matches!(result, Err(VncError::IoError(_))));
+}