diff --git a/src/gateway.rs b/src/gateway.rs index ffdc889..6e54b8b 100644 --- a/src/gateway.rs +++ b/src/gateway.rs @@ -5,17 +5,17 @@ use std::{ use crate::route_tcp; -pub fn router() { - let listener = TcpListener::bind("0.0.0.0:7879").unwrap(); +pub fn router() -> std::io::Result<()> { + 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 server_recv = TcpStream::connect(&address).unwrap(); - let mut server_send = server_recv.try_clone().unwrap(); + let mut client_send = client_recv.try_clone()?; + let mut server_recv = TcpStream::connect(&address)?; + let mut server_send = server_recv.try_clone()?; // Client -> Server - std::thread::spawn(move || { - while !route_tcp(&mut client_recv, &mut server_send) {} - server_send.shutdown(Shutdown::Both).unwrap(); + std::thread::spawn(move || -> std::io::Result<()> { + while !route_tcp(&mut client_recv, &mut server_send)? {} + server_send.shutdown(Shutdown::Both)?; + Ok(()) }); // Server -> Client - while !route_tcp(&mut server_recv, &mut client_send) {} - client_send.shutdown(Shutdown::Both).unwrap(); + while !route_tcp(&mut server_recv, &mut client_send)? {} + 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(); - connections.remove(connections.iter().position(|v| v == &address).unwrap()); + stream.shutdown(Shutdown::Both)?; + if let Some(v) = connections.iter().position(|v| v == &address) { + connections.remove(v); + } } Err(_) => {} } diff --git a/src/lib.rs b/src/lib.rs index e820f3a..c5326ad 100644 --- a/src/lib.rs +++ b/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 { - let addr = get_addr().unwrap(); - - println!("Connecting to: {addr}"); - - let mut stream = TcpStream::connect(addr).unwrap(); - + pub fn connect(ip: (u8, u8, u8, u8), port: u16) -> std::io::Result { + let addr = get_addr()?; + let mut stream = TcpStream::connect(addr)?; let bep = port.to_be_bytes(); - stream - .write_all(&[ip.0, ip.1, ip.2, ip.3, bep[0], bep[1]]) - .unwrap(); + stream.write_all(&[ip.0, ip.1, ip.2, ip.3, bep[0], bep[1]])?; - Self(stream) + Ok(Self(stream)) } - pub fn write(&mut self, data: &[u8]) { - self.0.write_all(data).unwrap() + pub fn write(&mut self, data: &[u8]) -> std::io::Result<()> { + self.0.write_all(data) } - pub fn read(&mut self, data: &mut [u8]) -> usize { - self.0.read(data).unwrap() + pub fn read(&mut self, data: &mut [u8]) -> std::io::Result { + self.0.read(data) } } pub fn get_addr() -> std::io::Result { 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 { )) } -pub fn route_tcp(from: &mut TcpStream, to: &mut TcpStream) -> bool { +pub fn route_tcp(from: &mut TcpStream, to: &mut TcpStream) -> std::io::Result { 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) } diff --git a/src/main.rs b/src/main.rs index f78091c..97f7219 100644 --- a/src/main.rs +++ b/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); - stream.write("Hello, World".as_bytes()); - let mut data = [0; 12]; - let data_len = stream.read(&mut data); - if data_len == 0 { - return; - } - - println!("From server: {}", String::from_utf8(data.to_vec()).unwrap()); + let mut stream = GhostStream::connect((0, 0, 0, 0), 7525).unwrap(); + stream.write("Hello, World".as_bytes()).unwrap(); + let mut data = [0; 12]; + let data_len = stream.read(&mut data).unwrap(); + if data_len == 0 { + return; } + + println!("From server: {}", String::from_utf8(data.to_vec()).unwrap()); }