diff --git a/src/listener_server.rs b/src/listener_server.rs index 0fbd68a..a4c7c6d 100644 --- a/src/listener_server.rs +++ b/src/listener_server.rs @@ -54,5 +54,6 @@ pub fn listen_to_client(tx: mpsc::Sender, rx: mpsc::Receiver, po stream_handler(&tx, &rx, stream); println!("Closed connection with {ip}"); + let _ = tx.send("DUMMY CLOSE_CONNECTION".to_string()); } } diff --git a/src/main.rs b/src/main.rs index 7dd653c..d983e1d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,7 +4,6 @@ use eyre::Result; use pgp::{Deserializable, SignedPublicKey, SignedSecretKey}; use std::collections::HashMap; use std::fs; -use std::net::{Shutdown, TcpStream}; use std::sync::mpsc; use std::thread; use std::time::Duration; @@ -55,7 +54,11 @@ fn main() -> Result<()> { let mut ap = ArgumentParser::new(); ap.set_description("Encrypted IRC Bouncer"); ap.refer(&mut server) - .add_argument("server", Store, "The Address Of The Server The Bouncer Connects To") + .add_argument( + "server", + Store, + "The Address Of The Server The Bouncer Connects To", + ) .required(); ap.refer(&mut port) .add_option(&["-p", "--port"], Store, "The Port The Bouncer Binds To"); @@ -67,44 +70,25 @@ fn main() -> Result<()> { ap.parse_args_or_exit(); } - let server = &server; - - let stream = TcpStream::connect(format!("{server}:{server_port}"))?; - let public_key = fs::read(public_key_location)?; let secret_key = SignedSecretKey::from_bytes(fs::read(secret_key_location)?.as_slice())?; - let reader_stream = match stream.try_clone() { - Ok(stream) => stream, - Err(_error) => { - let _ = stream.shutdown(Shutdown::Both); - panic!("Failed to create the reader stream") - } - }; - - let writer_stream = match stream.try_clone() { - Ok(stream) => stream, - Err(_error) => { - let _ = stream.shutdown(Shutdown::Both); - let _ = reader_stream.shutdown(Shutdown::Both); - panic!("Failed to create the writer stream") - } - }; - let (listener_channel_send_tx, listener_channel_rx) = mpsc::channel(); let (listener_channel_tx, listener_channel_recv_rx) = mpsc::channel(); let (writer_channel_tx, writer_channel_send_rx) = mpsc::channel(); let (writer_channel_recv_tx, writer_channel_rx) = mpsc::channel(); + let tmp_port = port.clone(); thread::spawn(move || { - listener_server::listen_to_client(listener_channel_send_tx, listener_channel_recv_rx, port) + listener_server::listen_to_client(listener_channel_send_tx, listener_channel_recv_rx, tmp_port) }); + let tmp_port = server_port.clone(); let tmp_server = server.clone(); - thread::spawn(|| { + thread::spawn(move || { writer_client::write_to_server( - writer_stream, - tmp_server, + &tmp_server, + &tmp_port, writer_channel_send_rx, writer_channel_recv_tx, ) @@ -118,7 +102,7 @@ fn main() -> Result<()> { let _ = client_handler::handle_message_from_client( &message, &public_key, - server, + &server, &mut keys, &writer_channel_tx, &writer_channel_rx, @@ -138,7 +122,7 @@ fn main() -> Result<()> { &message, &public_key, &secret_key, - server, + &server, passwd, &mut keys, &writer_channel_tx, diff --git a/src/writer_client.rs b/src/writer_client.rs index c69e017..d3ae7cd 100644 --- a/src/writer_client.rs +++ b/src/writer_client.rs @@ -6,54 +6,68 @@ use std::thread; use std::time::Duration; pub fn write_to_server( - tcp_stream: TcpStream, - server: String, + server: &str, + port: &str, rx: mpsc::Receiver, tx: mpsc::Sender, ) { - let connector = SslConnector::builder(SslMethod::tls()).unwrap().build(); - let mut stream = connector - .connect(&server, tcp_stream) - .expect("Couldn't start TLS"); + 'big: loop { + println!("Connecting to {server}:{port}"); + let tcp_stream = + TcpStream::connect(format!("{server}:{port}")).expect("Couldn't connect to server"); - stream - .get_mut() - .set_nonblocking(true) - .expect("Failed to set nonblocking"); + let connector = SslConnector::builder(SslMethod::tls()).unwrap().build(); + let mut stream = connector + .connect(&server, &tcp_stream) + .expect("Couldn't start TLS"); - loop { - let mut buffer: Vec = Vec::new(); - let mut buf: [u8; 1] = [0]; - let newline: u8 = b'\n'; + stream + .get_mut() + .set_nonblocking(true) + .expect("Failed to set nonblocking"); - while buf[0] != newline { - match stream.ssl_read(&mut buf) { - Ok(_length) => { - if _length > 0 { - buffer.push(buf[0]); - } - } - Err(_error) => match _error.io_error() { - None => { - dbg!(_error.ssl_error()); - } - Some(error) => match error.kind() { - ErrorKind::WouldBlock => {} - _ => { - dbg!(error.kind()); - println!("Couldn't read the stream"); + loop { + let mut buffer: Vec = Vec::new(); + let mut buf: [u8; 1] = [0]; + let newline: u8 = b'\n'; + + while buf[0] != newline { + match stream.ssl_read(&mut buf) { + Ok(_length) => { + if _length > 0 { + buffer.push(buf[0]); } + } + Err(_error) => match _error.io_error() { + None => { + dbg!(_error.ssl_error()); + continue 'big; + } + Some(error) => match error.kind() { + ErrorKind::WouldBlock => {} + _ => { + dbg!(error.kind()); + println!("Couldn't read the stream"); + continue 'big; + } + }, }, - }, + } + let value = rx.try_recv().unwrap_or("".to_string()); + match value.as_str() { + "DUMMY CLOSE_CONNECTION" => { + continue 'big; + } + _ => {} + } + match stream.write_all(value.as_bytes()) { + Ok(_) => {} + Err(_e) => println!("Couldn't send {value}"), + }; + thread::sleep(Duration::from_micros(100)); } - let value = rx.try_recv().unwrap_or("".to_string()); - match stream.write_all(value.as_bytes()) { - Ok(_) => {} - Err(_e) => println!("Couldn't send {value}"), - }; - thread::sleep(Duration::from_micros(100)); - } - let _ = tx.send(dbg!(String::from_utf8_lossy(&buffer).to_string())); + let _ = tx.send(dbg!(String::from_utf8_lossy(&buffer).to_string())); + } } }