Improved error handling

src/gateway.rs

@@ -5,17 +5,17 @@ use std::{
 
 use crate::route_tcp;
 
-pub fn router() {
+pub fn router() -> std::io::Result<()> {
-    let listener = TcpListener::bind("0.0.0.0:7879").unwrap();
+    let listener = TcpListener::bind("0.0.0.0:7879")?;
 
     for stream in listener.incoming() {
         match stream {
             Ok(mut client_recv) => {
-                std::thread::spawn(|| {
+                std::thread::spawn(|| -> std::io::Result<()> {
                     let mut address = [0; 6];
-                    let bytes_read = client_recv.read(&mut address).unwrap();
+                    let bytes_read = client_recv.read(&mut address)?;
                     if bytes_read == 0 {
-                        return;
+                        return Ok(());
                     }
 
                     let address = format!(
@@ -27,24 +27,29 @@ pub fn router() {
                         u16::from_be_bytes([address[4], address[5]])
                     );
 
-                    let mut client_send = client_recv.try_clone().unwrap();
+                    let mut client_send = client_recv.try_clone()?;
-                    let mut server_recv = TcpStream::connect(&address).unwrap();
+                    let mut server_recv = TcpStream::connect(&address)?;
-                    let mut server_send = server_recv.try_clone().unwrap();
+                    let mut server_send = server_recv.try_clone()?;
 
                     // Client -> Server
-                    std::thread::spawn(move || {
+                    std::thread::spawn(move || -> std::io::Result<()> {
-                        while !route_tcp(&mut client_recv, &mut server_send) {}
+                        while !route_tcp(&mut client_recv, &mut server_send)? {}
-                        server_send.shutdown(Shutdown::Both).unwrap();
+                        server_send.shutdown(Shutdown::Both)?;
+                        Ok(())
                     });
 
                     // Server -> Client
-                    while !route_tcp(&mut server_recv, &mut client_send) {}
+                    while !route_tcp(&mut server_recv, &mut client_send)? {}
-                    client_send.shutdown(Shutdown::Both).unwrap();
+                    client_send.shutdown(Shutdown::Both)?;
+
+                    Ok(())
                 });
             }
             Err(_) => {}
         }
     }
+
+    Ok(())
 }
 
 pub fn run() -> std::io::Result<()> {
@@ -71,10 +76,12 @@ pub fn run() -> std::io::Result<()> {
 
                 connections.push(address);
 
-                while stream.read(&mut []).unwrap() != 0 {}
+                while stream.read(&mut [])? != 0 {}
 
-                stream.shutdown(Shutdown::Both).unwrap();
+                stream.shutdown(Shutdown::Both)?;
-                connections.remove(connections.iter().position(|v| v == &address).unwrap());
+                if let Some(v) = connections.iter().position(|v| v == &address) {
+                    connections.remove(v);
+                }
             }
             Err(_) => {}
         }

src/lib.rs

@@ -8,34 +8,28 @@ pub mod gateway;
 pub struct GhostStream(TcpStream);
 
 impl GhostStream {
-    pub fn connect(ip: (u8, u8, u8, u8), port: u16) -> Self {
+    pub fn connect(ip: (u8, u8, u8, u8), port: u16) -> std::io::Result<Self> {
-        let addr = get_addr().unwrap();
+        let addr = get_addr()?;
-
+        let mut stream = TcpStream::connect(addr)?;
-        println!("Connecting to: {addr}");
-
-        let mut stream = TcpStream::connect(addr).unwrap();
-
         let bep = port.to_be_bytes();
-        stream
+        stream.write_all(&[ip.0, ip.1, ip.2, ip.3, bep[0], bep[1]])?;
-            .write_all(&[ip.0, ip.1, ip.2, ip.3, bep[0], bep[1]])
-            .unwrap();
 
-        Self(stream)
+        Ok(Self(stream))
     }
 
-    pub fn write(&mut self, data: &[u8]) {
+    pub fn write(&mut self, data: &[u8]) -> std::io::Result<()> {
-        self.0.write_all(data).unwrap()
+        self.0.write_all(data)
     }
 
-    pub fn read(&mut self, data: &mut [u8]) -> usize {
+    pub fn read(&mut self, data: &mut [u8]) -> std::io::Result<usize> {
-        self.0.read(data).unwrap()
+        self.0.read(data)
     }
 }
 
 pub fn get_addr() -> std::io::Result<String> {
     let mut stream = TcpStream::connect("0.0.0.0:7878")?;
 
-    stream.write_all(&[0, 0, 0, 0, 0x1E, 0xC7]).unwrap();
+    stream.write_all(&[0, 0, 0, 0, 0x1E, 0xC7])?;
 
     let mut address = [0; 6];
 
@@ -51,15 +45,15 @@ pub fn get_addr() -> std::io::Result<String> {
     ))
 }
 
-pub fn route_tcp(from: &mut TcpStream, to: &mut TcpStream) -> bool {
+pub fn route_tcp(from: &mut TcpStream, to: &mut TcpStream) -> std::io::Result<bool> {
     let mut data = [0; 2048];
 
-    let bytes_read = from.read(&mut data).unwrap();
+    let bytes_read = from.read(&mut data)?;
     if bytes_read == 0 {
-        return true;
+        return Ok(true);
     }
 
-    to.write_all(&data[..bytes_read]).unwrap();
+    to.write_all(&data[..bytes_read])?;
 
-    false
+    Ok(false)
 }

src/main.rs

@@ -1,15 +1,13 @@
 use ghostnet_rs::GhostStream;
 
 fn main() {
-    if let Err(_) = ghostnet_rs::gateway::run() {
+    let mut stream = GhostStream::connect((0, 0, 0, 0), 7525).unwrap();
-        let mut stream = GhostStream::connect((0, 0, 0, 0), 7525);
+    stream.write("Hello, World".as_bytes()).unwrap();
-        stream.write("Hello, World".as_bytes());
     let mut data = [0; 12];
-        let data_len = stream.read(&mut data);
+    let data_len = stream.read(&mut data).unwrap();
     if data_len == 0 {
         return;
     }
 
     println!("From server: {}", String::from_utf8(data.to_vec()).unwrap());
 }
-}