From 11a4777cfebfc41695583c23c545478dc50bdbe0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pascal=20Eng=C3=A9libert?= Date: Thu, 30 Oct 2025 11:54:08 +0100 Subject: [PATCH] Fmt, add experimental features --- .vscode/settings.json | 3 + examples/basic.rs | 66 ++++ examples/basic/src/main.rs | 57 --- rustfmt.toml | 8 + src/data.rs | 741 +++++++++++++++++-------------------- src/lib.rs | 274 +++++++------- src/relay.rs | 324 ++++++++++------ src/tcp.rs | 581 +++++++++++++++++------------ 8 files changed, 1091 insertions(+), 963 deletions(-) create mode 100644 .vscode/settings.json create mode 100644 examples/basic.rs delete mode 100644 examples/basic/src/main.rs create mode 100644 rustfmt.toml diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..e931c34 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "rust-analyzer.showUnlinkedFileNotification": false +} \ No newline at end of file diff --git a/examples/basic.rs b/examples/basic.rs new file mode 100644 index 0000000..9a2da22 --- /dev/null +++ b/examples/basic.rs @@ -0,0 +1,66 @@ +use sslrelay::{self, CallbackRet, HandlerCallbacks, RelayConfig, TCPDataType, TLSConfig}; + +// Handler object +#[derive(Clone)] // Must have Clone trait implemented. +struct Handler; + +/* + Callback traits that can be used to read or inject data + into data upstream or downstream. +*/ +impl HandlerCallbacks for Handler { + // DownStream non blocking callback + fn ds_nb_callback(&self, in_data: Vec, _conn_id: u64) { + if let Ok(in_data) = str::from_utf8(&in_data) { + println!("[downstream] {in_data}"); + } else { + //println!("[downstream] {in_data:?}"); + } + } + + // DownStream blocking callback + fn ds_b_callback(&mut self, _in_data: Vec, _conn_id: u64) -> CallbackRet { + //println!("[CALLBACK] Down Stream Blocking CallBack!"); + CallbackRet::Relay(_in_data) + } + + // UpStream non blocking callback + fn us_nb_callback(&self, in_data: Vec, _conn_id: u64) { + if let Ok(in_data) = str::from_utf8(&in_data) { + println!("[upstream] {in_data}"); + } else { + //println!("[upstream] {in_data:?}"); + } + } + + // UpStream blocking callback + fn us_b_callback(&mut self, _in_data: Vec, _conn_id: u64) -> CallbackRet { + //println!("[CALLBACK] Up Stream Blocking CallBack!"); + CallbackRet::Relay(_in_data) + } +} + +fn main() { + // Create new SSLRelay object + let mut relay = sslrelay::SSLRelay::new( + Handler, + RelayConfig { + downstream_data_type: TCPDataType::TLS, + upstream_data_type: TCPDataType::TLS, + bind_host: "127.0.0.1".to_string(), + bind_port: "443".to_string(), + remote_host: |server_name| { + server_name.map_or_else(|| panic!("NO HOST"), str::to_string) + }, + remote_port: "443".to_string(), + tls_config: + TLSConfig::FILE { + certificate_path: "/dev/shm/exp/certs/prime256v1/apple.com.crt".to_string(), + private_key_path: "/dev/shm/exp/certs/prime256v1/apple.com.key".to_string(), + }, + }, + ); + + // Start listening + relay.start(); +} diff --git a/examples/basic/src/main.rs b/examples/basic/src/main.rs deleted file mode 100644 index cf8ff14..0000000 --- a/examples/basic/src/main.rs +++ /dev/null @@ -1,57 +0,0 @@ -use sslrelay::{self, RelayConfig, HandlerCallbacks, CallbackRet, TCPDataType, TLSConfig}; - -// Handler object -#[derive(Clone)] // Must have Clone trait implemented. -struct Handler; - -/* - Callback traits that can be used to read or inject data - into data upstream or downstream. -*/ -impl HandlerCallbacks for Handler { - - // DownStream non blocking callback - fn ds_nb_callback(&self, _in_data: Vec) { - println!("[CALLBACK] Down Stream Non Blocking CallBack!"); - } - - // DownStream blocking callback - fn ds_b_callback(&mut self, _in_data: Vec) -> CallbackRet { - println!("[CALLBACK] Down Stream Blocking CallBack!"); - CallbackRet::Relay(_in_data) - } - - // UpStream non blocking callback - fn us_nb_callback(&self, _in_data: Vec) { - println!("[CALLBACK] Up Stream Non Blocking CallBack!"); - } - - // UpStream blocking callback - fn us_b_callback(&mut self, _in_data: Vec) -> CallbackRet { - println!("[CALLBACK] Up Stream Blocking CallBack!"); - CallbackRet::Relay(_in_data) - } -} - -fn main() { - - // Create new SSLRelay object - let mut relay = sslrelay::SSLRelay::new( - Handler, - RelayConfig { - downstream_data_type: TCPDataType::TLS, - upstream_data_type: TCPDataType::TLS, - bind_host: "0.0.0.0".to_string(), - bind_port: "443".to_string(), - remote_host: "remote.com".to_string(), - remote_port: "443".to_string(), - tls_config: TLSConfig::FILE{ - certificate_path: "./tls.crt".to_string(), - private_key_path: "./tls.key".to_string(), - }, - } - ); - - // Start listening - relay.start(); -} \ No newline at end of file diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 0000000..ade1f81 --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1,8 @@ +hard_tabs = true +newline_style = "unix" + +unstable_features = true +format_code_in_doc_comments = true +format_macro_bodies = true +format_macro_matchers = true +format_strings = true diff --git a/src/data.rs b/src/data.rs index 3a20305..f28fc03 100644 --- a/src/data.rs +++ b/src/data.rs @@ -1,441 +1,370 @@ use crate::{ - DownStreamInner, - UpStreamInner, - FullDuplexTcpState, - DataPipe, - DataStreamType, - Sender, - Receiver, - Shutdown, - mpsc, - Duration, - Read, - Write, - io, + io, mpsc, DataPipe, DataStreamType, DownStreamInner, Duration, FullDuplexTcpState, Read, + Receiver, Sender, Shutdown, UpStreamInner, Write, }; impl DownStreamInner { + fn handle_error(error_description: &str) { + println!("[SSLRelay DownStream Thread Error]: {}", error_description); + } - fn handle_error(error_description: &str) { - println!("[SSLRelay DownStream Thread Error]: {}", error_description); - } + pub fn ds_handler(self, data_out: Sender, data_in: Receiver) { + match &self.ds_stream { + DataStreamType::RAW(_) => self.handle_raw(data_out, data_in), + DataStreamType::TLS(_) => self.handle_tls(data_out, data_in), + } + } - pub fn ds_handler(self, data_out: Sender, data_in: Receiver) { + fn handle_raw(mut self, data_out: Sender, data_in: Receiver) { + let mut raw_stream = match &self.ds_stream { + DataStreamType::RAW(ref s) => s, + _ => return, + }; - match &self.ds_stream { - DataStreamType::RAW(_) => self.handle_raw(data_out, data_in), - DataStreamType::TLS(_) => self.handle_tls(data_out, data_in), - } - } + loop { + match data_in.recv_timeout(Duration::from_millis(50)) { + // DataPipe Received + Ok(data_received) => match data_received { + DataPipe::DataWrite(data) => { + match raw_stream.write_all(&data) { + Ok(()) => {} + Err(_e) => { + Self::handle_error( + "Failed to write data to DownStream tcp stream!", + ); + let _ = data_out.send(FullDuplexTcpState::DownStreamShutDown); + let _ = raw_stream.shutdown(Shutdown::Both); + return; + } + } + let _ = raw_stream.flush(); + } + DataPipe::Shutdown => { + let _ = raw_stream.shutdown(Shutdown::Both); + return; + } + }, + Err(_e) => match _e { + mpsc::RecvTimeoutError::Timeout => {} + mpsc::RecvTimeoutError::Disconnected => { + Self::handle_error("DownStream data_in channel is disconnected!"); + return; + } + }, + } // End of data_in receive - fn handle_raw(mut self, data_out: Sender, data_in: Receiver) { + // If received data + if let Some(byte_count) = + Self::get_data_stream(&mut raw_stream, &mut self.internal_data_buffer) + { + if byte_count > 0 { + if let Err(_e) = data_out.send(FullDuplexTcpState::UpStreamWrite( + self.internal_data_buffer.clone(), + )) { + //Self::handle_error(format!("Failed to send UpStreamWrite to main thread: {}", e).as_str()); + let _ = raw_stream.shutdown(Shutdown::Both); + return; + } - let mut raw_stream = match &self.ds_stream { - DataStreamType::RAW(ref s) => s, - _ => return, - }; + self.internal_data_buffer.clear(); + } else if byte_count == 0 || byte_count == -2 { + let _ = data_out.send(FullDuplexTcpState::DownStreamShutDown); + let _ = raw_stream.shutdown(Shutdown::Both); + return; + } else if byte_count == -1 { + continue; + } + } else { + } + } + } - loop { + fn handle_tls(mut self, data_out: Sender, data_in: Receiver) { + let mut tls_stream = match self.ds_stream { + DataStreamType::TLS(ref mut s) => s, + _ => return, + }; - match data_in.recv_timeout(Duration::from_millis(50)) { + loop { + match data_in.recv_timeout(Duration::from_millis(50)) { + // DataPipe Received + Ok(data_received) => match data_received { + DataPipe::DataWrite(data) => { + match tls_stream.write_all(&data) { + Ok(()) => {} + Err(_e) => { + Self::handle_error( + "Failed to write data to DownStream tcp stream!", + ); + let _ = data_out.send(FullDuplexTcpState::DownStreamShutDown); + let _ = tls_stream.shutdown(); + return; + } + } + let _ = tls_stream.flush(); + } + DataPipe::Shutdown => { + let _ = tls_stream.shutdown(); + return; + } + }, + Err(_e) => match _e { + mpsc::RecvTimeoutError::Timeout => {} + mpsc::RecvTimeoutError::Disconnected => { + Self::handle_error("DownStream data_in channel is disconnected!"); + return; + } + }, + } // End of data_in receive - // DataPipe Received - Ok(data_received) => { + // If received data + if let Some(byte_count) = + Self::get_data_stream(&mut tls_stream, &mut self.internal_data_buffer) + { + if byte_count > 0 { + if let Err(_e) = data_out.send(FullDuplexTcpState::UpStreamWrite( + self.internal_data_buffer.clone(), + )) { + //Self::handle_error(format!("Failed to send UpStreamWrite to main thread: {}", e).as_str()); + let _ = tls_stream.shutdown(); + return; + } - match data_received { - DataPipe::DataWrite(data) => { + self.internal_data_buffer.clear(); + } else if byte_count == 0 || byte_count == -2 { + let _ = data_out.send(FullDuplexTcpState::DownStreamShutDown); + let _ = tls_stream.shutdown(); + return; + } else if byte_count == -1 { + continue; + } + } else { + } + } + } - match raw_stream.write_all(&data) { - Ok(()) => {}, - Err(_e) => { - Self::handle_error("Failed to write data to DownStream tcp stream!"); - let _ = data_out.send(FullDuplexTcpState::DownStreamShutDown); - let _ = raw_stream.shutdown(Shutdown::Both); - return; - } - } - let _ = raw_stream.flush(); + fn get_data_stream(stream: &mut S, internal_data_buffer: &mut Vec) -> Option { + let mut data_length: i64 = 0; - }, - DataPipe::Shutdown => { - let _ = raw_stream.shutdown(Shutdown::Both); - return; - }, - } - }, - Err(_e) => { - match _e { - mpsc::RecvTimeoutError::Timeout => {}, - mpsc::RecvTimeoutError::Disconnected => { - Self::handle_error("DownStream data_in channel is disconnected!"); - return; - } - } - } - }// End of data_in receive + loop { + let mut r_buf = [0; 1024]; - // If received data - if let Some(byte_count) = Self::get_data_stream(&mut raw_stream, &mut self.internal_data_buffer) { - if byte_count > 0 { + match stream.read(&mut r_buf) { + Ok(bytes_read) => { + if bytes_read == 0 { + break; + } else if bytes_read != 0 && bytes_read <= 1024 { + /* + let mut tmp_buf = r_buf.to_vec(); + tmp_buf.truncate(bytes_read); + */ - if let Err(_e) = data_out.send(FullDuplexTcpState::UpStreamWrite(self.internal_data_buffer.clone())) { - - //Self::handle_error(format!("Failed to send UpStreamWrite to main thread: {}", e).as_str()); - let _ = raw_stream.shutdown(Shutdown::Both); - return; - } + //let _bw = self.internal_data_buffer.write(&tmp_buf).unwrap(); - self.internal_data_buffer.clear(); - - } else if byte_count == 0 || byte_count == -2 { - - let _ = data_out.send(FullDuplexTcpState::DownStreamShutDown); - let _ = raw_stream.shutdown(Shutdown::Both); - return; - - } else if byte_count == -1 { - continue; - } - } else { - - } - } - } - - fn handle_tls(mut self, data_out: Sender, data_in: Receiver) { - - let mut tls_stream = match self.ds_stream { - DataStreamType::TLS(ref mut s) => s, - _ => return, - }; - - loop { - - match data_in.recv_timeout(Duration::from_millis(50)) { - - // DataPipe Received - Ok(data_received) => { - - match data_received { - DataPipe::DataWrite(data) => { - - match tls_stream.write_all(&data) { - Ok(()) => {}, - Err(_e) => { - Self::handle_error("Failed to write data to DownStream tcp stream!"); - let _ = data_out.send(FullDuplexTcpState::DownStreamShutDown); - let _ = tls_stream.shutdown(); - return; - } - } - let _ = tls_stream.flush(); - }, - DataPipe::Shutdown => { - let _ = tls_stream.shutdown(); - return; - }, - } - }, - Err(_e) => { - match _e { - mpsc::RecvTimeoutError::Timeout => {}, - mpsc::RecvTimeoutError::Disconnected => { - Self::handle_error("DownStream data_in channel is disconnected!"); - return; - } - } - } - }// End of data_in receive - - // If received data - if let Some(byte_count) = Self::get_data_stream(&mut tls_stream, &mut self.internal_data_buffer) { - if byte_count > 0 { - - if let Err(_e) = data_out.send(FullDuplexTcpState::UpStreamWrite(self.internal_data_buffer.clone())) { - - //Self::handle_error(format!("Failed to send UpStreamWrite to main thread: {}", e).as_str()); - let _ = tls_stream.shutdown(); - return; - } - - self.internal_data_buffer.clear(); - - } else if byte_count == 0 || byte_count == -2 { - - let _ = data_out.send(FullDuplexTcpState::DownStreamShutDown); - let _ = tls_stream.shutdown(); - return; - - } else if byte_count == -1 { - continue; - } - } else { - - } - } - } - - fn get_data_stream(stream: &mut S, internal_data_buffer: &mut Vec) -> Option { - - let mut data_length: i64 = 0; - - loop { - - let mut r_buf = [0; 1024]; - - match stream.read(&mut r_buf) { - - Ok(bytes_read) => { - - if bytes_read == 0 { - break; - - } else if bytes_read != 0 && bytes_read <= 1024 { - - /* - let mut tmp_buf = r_buf.to_vec(); - tmp_buf.truncate(bytes_read); - */ - - - //let _bw = self.internal_data_buffer.write(&tmp_buf).unwrap(); - - let _bw = internal_data_buffer.write(r_buf.split_at(bytes_read).0).unwrap(); - data_length += bytes_read as i64; - - } else { - println!("[+] Else hit!!!!!!!!!!!!!!!!!!!!!!"); - } - }, - Err(e) => { - match e.kind() { - - io::ErrorKind::WouldBlock => { - if data_length == 0 { - - data_length = -1; - } - break; - - }, - io::ErrorKind::ConnectionReset => { - data_length = -2; - break; - }, - _ => {println!("[!!!] Got error: {}",e);} - } - }, - } - } - return Some(data_length); - } + let _bw = internal_data_buffer + .write(r_buf.split_at(bytes_read).0) + .unwrap(); + data_length += bytes_read as i64; + } else { + println!("[+] Else hit!!!!!!!!!!!!!!!!!!!!!!"); + } + } + Err(e) => match e.kind() { + io::ErrorKind::WouldBlock => { + if data_length == 0 { + data_length = -1; + } + break; + } + io::ErrorKind::ConnectionReset => { + data_length = -2; + break; + } + _ => { + println!("[!!!downstream] Got error: {}", e); + } + }, + } + } + return Some(data_length); + } } - impl UpStreamInner { + fn handle_error(error_description: &str) { + println!("[SSLRelay UpStream Thread Error]: {}", error_description); + } - fn handle_error(error_description: &str) { - println!("[SSLRelay UpStream Thread Error]: {}", error_description); - } + pub fn us_handler(self, data_out: Sender, data_in: Receiver) { + match &self.us_stream { + DataStreamType::RAW(_) => self.handle_raw(data_out, data_in), + DataStreamType::TLS(_) => self.handle_tls(data_out, data_in), + } + } - pub fn us_handler(self, data_out: Sender, data_in: Receiver) { + fn handle_raw(mut self, data_out: Sender, data_in: Receiver) { + let mut raw_stream = match self.us_stream { + DataStreamType::RAW(ref s) => s, + _ => return, + }; - match &self.us_stream { - DataStreamType::RAW(_) => self.handle_raw(data_out, data_in), - DataStreamType::TLS(_) => self.handle_tls(data_out, data_in), - } + loop { + match data_in.recv_timeout(Duration::from_millis(50)) { + Ok(data_received) => match data_received { + DataPipe::DataWrite(data) => { + match raw_stream.write_all(&data) { + Ok(()) => {} + Err(_e) => { + Self::handle_error("Failed to write data to UpStream tcp stream!"); + let _ = data_out.send(FullDuplexTcpState::UpStreamShutDown); + let _ = raw_stream.shutdown(Shutdown::Both); + return; + } + } + let _ = raw_stream.flush(); + } + DataPipe::Shutdown => { + let _ = raw_stream.shutdown(Shutdown::Both); + return; + } + }, + Err(e) => match e { + mpsc::RecvTimeoutError::Timeout => {} + mpsc::RecvTimeoutError::Disconnected => { + Self::handle_error("UpStream data_in channel is disconnected!"); + return; + } + }, + } // End of data_in receive - } + if let Some(byte_count) = + Self::get_data_stream(&mut raw_stream, &mut self.internal_data_buffer) + { + if byte_count > 0 { + if let Err(_e) = data_out.send(FullDuplexTcpState::DownStreamWrite( + self.internal_data_buffer.clone(), + )) { + //Self::handle_error(format!("Failed to send DownStreamWrite to main thread: {}", e).as_str()); + let _ = raw_stream.shutdown(Shutdown::Both); + return; + } - fn handle_raw(mut self, data_out: Sender, data_in: Receiver) { + self.internal_data_buffer.clear(); + } else if byte_count == 0 || byte_count == -2 { + let _ = data_out.send(FullDuplexTcpState::UpStreamShutDown); + let _ = raw_stream.shutdown(Shutdown::Both); + return; + } else if byte_count == -1 { + continue; + } + } else { + } + } + } - let mut raw_stream = match self.us_stream { - DataStreamType::RAW(ref s) => s, - _ => return, - }; + fn handle_tls(mut self, data_out: Sender, data_in: Receiver) { + let mut tls_stream = match self.us_stream { + DataStreamType::TLS(ref mut s) => s, + _ => return, + }; - loop { + loop { + match data_in.recv_timeout(Duration::from_millis(50)) { + Ok(data_received) => match data_received { + DataPipe::DataWrite(data) => { + match tls_stream.write_all(&data) { + Ok(()) => {} + Err(_e) => { + Self::handle_error("Failed to write data to UpStream tcp stream!"); + let _ = data_out.send(FullDuplexTcpState::UpStreamShutDown); + let _ = tls_stream.shutdown(); + return; + } + } + let _ = tls_stream.flush(); + } + DataPipe::Shutdown => { + let _ = tls_stream.shutdown(); + return; + } + }, + Err(e) => match e { + mpsc::RecvTimeoutError::Timeout => {} + mpsc::RecvTimeoutError::Disconnected => { + Self::handle_error("UpStream data_in channel is disconnected!"); + return; + } + }, + } // End of data_in receive - match data_in.recv_timeout(Duration::from_millis(50)) { + if let Some(byte_count) = + Self::get_data_stream(&mut tls_stream, &mut self.internal_data_buffer) + { + if byte_count > 0 { + if let Err(_e) = data_out.send(FullDuplexTcpState::DownStreamWrite( + self.internal_data_buffer.clone(), + )) { + //Self::handle_error(format!("Failed to send DownStreamWrite to main thread: {}", e).as_str()); + let _ = tls_stream.shutdown(); + return; + } - Ok(data_received) => { + self.internal_data_buffer.clear(); + } else if byte_count == 0 || byte_count == -2 { + let _ = data_out.send(FullDuplexTcpState::UpStreamShutDown); + let _ = tls_stream.shutdown(); + return; + } else if byte_count == -1 { + continue; + } + } else { + } + } + } - match data_received { - DataPipe::DataWrite(data) => { + fn get_data_stream(stream: &mut S, internal_data_buffer: &mut Vec) -> Option { + let mut data_length: i64 = 0; - match raw_stream.write_all(&data) { - Ok(()) => {}, - Err(_e) => { - Self::handle_error("Failed to write data to UpStream tcp stream!"); - let _ = data_out.send(FullDuplexTcpState::UpStreamShutDown); - let _ = raw_stream.shutdown(Shutdown::Both); - return; - } - } - let _ = raw_stream.flush(); - }, - DataPipe::Shutdown => { - let _ = raw_stream.shutdown(Shutdown::Both); - return; - } - } - }, - Err(e) => { - match e { - mpsc::RecvTimeoutError::Timeout => {}, - mpsc::RecvTimeoutError::Disconnected => { - Self::handle_error("UpStream data_in channel is disconnected!"); - return; - } - } - } - }// End of data_in receive + loop { + let mut r_buf = [0; 1024]; - if let Some(byte_count) = Self::get_data_stream(&mut raw_stream, &mut self.internal_data_buffer) { - if byte_count > 0 { + match stream.read(&mut r_buf) { + Ok(bytes_read) => { + if bytes_read == 0 { + break; + } else if bytes_read != 0 && bytes_read <= 1024 { + /* + let mut tmp_buf = r_buf.to_vec(); + tmp_buf.truncate(bytes_read); + */ - if let Err(_e) = data_out.send(FullDuplexTcpState::DownStreamWrite(self.internal_data_buffer.clone())) { - - //Self::handle_error(format!("Failed to send DownStreamWrite to main thread: {}", e).as_str()); - let _ = raw_stream.shutdown(Shutdown::Both); - return; - } + //let _bw = self.internal_data_buffer.write(&tmp_buf).unwrap(); - self.internal_data_buffer.clear(); - - } else if byte_count == 0 || byte_count == -2 { - - let _ = data_out.send(FullDuplexTcpState::UpStreamShutDown); - let _ = raw_stream.shutdown(Shutdown::Both); - return; - } else if byte_count == -1 { - continue; - } - } else { - } - } - } - - fn handle_tls(mut self, data_out: Sender, data_in: Receiver) { - - let mut tls_stream = match self.us_stream { - DataStreamType::TLS(ref mut s) => s, - _ => return, - }; - - loop { - - match data_in.recv_timeout(Duration::from_millis(50)) { - - Ok(data_received) => { - - match data_received { - DataPipe::DataWrite(data) => { - - match tls_stream.write_all(&data) { - Ok(()) => {}, - Err(_e) => { - Self::handle_error("Failed to write data to UpStream tcp stream!"); - let _ = data_out.send(FullDuplexTcpState::UpStreamShutDown); - let _ = tls_stream.shutdown(); - return; - } - } - let _ = tls_stream.flush(); - }, - DataPipe::Shutdown => { - let _ = tls_stream.shutdown(); - return; - } - } - }, - Err(e) => { - match e { - mpsc::RecvTimeoutError::Timeout => {}, - mpsc::RecvTimeoutError::Disconnected => { - Self::handle_error("UpStream data_in channel is disconnected!"); - return; - } - } - } - }// End of data_in receive - - if let Some(byte_count) = Self::get_data_stream(&mut tls_stream, &mut self.internal_data_buffer) { - if byte_count > 0 { - - if let Err(_e) = data_out.send(FullDuplexTcpState::DownStreamWrite(self.internal_data_buffer.clone())) { - - //Self::handle_error(format!("Failed to send DownStreamWrite to main thread: {}", e).as_str()); - let _ = tls_stream.shutdown(); - return; - } - - self.internal_data_buffer.clear(); - - } else if byte_count == 0 || byte_count == -2 { - - let _ = data_out.send(FullDuplexTcpState::UpStreamShutDown); - let _ = tls_stream.shutdown(); - return; - } else if byte_count == -1 { - continue; - } - } else { - } - } - } - - fn get_data_stream(stream: &mut S, internal_data_buffer: &mut Vec) -> Option { - - let mut data_length: i64 = 0; - - loop { - - let mut r_buf = [0; 1024]; - - match stream.read(&mut r_buf) { - - Ok(bytes_read) => { - - if bytes_read == 0 { - - break; - - } else if bytes_read != 0 && bytes_read <= 1024 { - - /* - let mut tmp_buf = r_buf.to_vec(); - tmp_buf.truncate(bytes_read); - */ - - - //let _bw = self.internal_data_buffer.write(&tmp_buf).unwrap(); - - let _bw = internal_data_buffer.write(r_buf.split_at(bytes_read).0).unwrap(); - data_length += bytes_read as i64; - - } else { - println!("[+] Else hit!!!!!!!!!!!!!!!!!!!!!!"); - } - }, - Err(e) => { - match e.kind() { - - io::ErrorKind::WouldBlock => { - if data_length == 0 { - data_length = -1; - } - break; - }, - io::ErrorKind::ConnectionReset => { - data_length = -2; - break; - }, - _ => {println!("[!!!] Got error: {}",e);} - } - }, - } - } - return Some(data_length); - } -} \ No newline at end of file + let _bw = internal_data_buffer + .write(r_buf.split_at(bytes_read).0) + .unwrap(); + data_length += bytes_read as i64; + } else { + println!("[+] Else hit!!!!!!!!!!!!!!!!!!!!!!"); + } + } + Err(e) => match e.kind() { + io::ErrorKind::WouldBlock => { + if data_length == 0 { + data_length = -1; + } + break; + } + io::ErrorKind::ConnectionReset => { + data_length = -2; + break; + } + _ => { + println!("[!!!upstream] Got error: {}", e); + } + }, + } + } + return Some(data_length); + } +} diff --git a/src/lib.rs b/src/lib.rs index e6cfa2f..b30a64e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,157 +4,132 @@ //! This library allows you to implement callback functions for upstream and downstream traffic. //! These callbacks can R/W the data from a stream(Blocking) or only R the data(Non-Blocking). //!``` -//!pub trait HandlerCallbacks { -//! fn ds_b_callback(&mut self, _in_data: Vec) -> CallbackRet {CallbackRet::Relay(_in_data)} -//! fn ds_nb_callback(&self, _in_data: Vec){} -//! fn us_b_callback(&mut self, _in_data: Vec) -> CallbackRet {CallbackRet::Relay(_in_data)} -//! fn us_nb_callback(&self, _in_data: Vec){} -//!} -//!``` +//! pub trait HandlerCallbacks { +//! fn ds_b_callback(&mut self, _in_data: Vec) -> CallbackRet { +//! CallbackRet::Relay(_in_data) +//! } +//! fn ds_nb_callback(&self, _in_data: Vec) {} +//! fn us_b_callback(&mut self, _in_data: Vec) -> CallbackRet { +//! CallbackRet::Relay(_in_data) +//! } +//! fn us_nb_callback(&self, _in_data: Vec) {} +//! } +//! ``` //! The blocking callbacks return an enum called CallbackRet with four different variants. //! The variants control the flow of the tcp stream. //!``` //! pub enum CallbackRet { -//! Relay(Vec),// Relay data -//! Spoof(Vec),// Skip relaying and send data back -//! Shutdown,// Shutdown TCP connection -//! Freeze,// Dont send data (pretend as if stream never was recieved) +//! Relay(Vec), // Relay data +//! Spoof(Vec), // Skip relaying and send data back +//! Shutdown, // Shutdown TCP connection +//! Freeze, // Dont send data (pretend as if stream never was recieved) //! } //! ``` //! ## Example (basic.rs) //! ``` -//! use sslrelay::{self, RelayConfig, HandlerCallbacks, CallbackRet, TCPDataType, TLSConfig}; -//! +//! use sslrelay::{self, CallbackRet, HandlerCallbacks, RelayConfig, TCPDataType, TLSConfig}; +//! //! // Handler object //! #[derive(Clone)] // Must have Clone trait implemented. //! struct Handler; -//! +//! //! /* -//! Callback traits that can be used to read or inject data -//! into data upstream or downstream. +//! Callback traits that can be used to read or inject data +//! into data upstream or downstream. //! */ //! impl HandlerCallbacks for Handler { -//! -//! // DownStream non blocking callback (Read Only) -//! fn ds_nb_callback(&self, _in_data: Vec) { -//! println!("[CALLBACK] Down Stream Non Blocking CallBack!"); -//! } -//! -//! // DownStream blocking callback (Read & Write) -//! fn ds_b_callback(&mut self, _in_data: Vec) -> CallbackRet { -//! println!("[CALLBACK] Down Stream Blocking CallBack!"); -//! CallbackRet::Relay(_in_data) -//! } -//! -//! // UpStream non blocking callback (Read Only) -//! fn us_nb_callback(&self, _in_data: Vec) { -//! println!("[CALLBACK] Up Stream Non Blocking CallBack!"); -//! } -//! -//! // UpStream blocking callback (Read & Write) -//! fn us_b_callback(&mut self, _in_data: Vec) -> CallbackRet { -//! println!("[CALLBACK] Up Stream Blocking CallBack!"); -//! CallbackRet::Relay(_in_data) -//! } +//! // DownStream non blocking callback (Read Only) +//! fn ds_nb_callback(&self, _in_data: Vec) { +//! println!("[CALLBACK] Down Stream Non Blocking CallBack!"); +//! } +//! +//! // DownStream blocking callback (Read & Write) +//! fn ds_b_callback(&mut self, _in_data: Vec) -> CallbackRet { +//! println!("[CALLBACK] Down Stream Blocking CallBack!"); +//! CallbackRet::Relay(_in_data) +//! } +//! +//! // UpStream non blocking callback (Read Only) +//! fn us_nb_callback(&self, _in_data: Vec) { +//! println!("[CALLBACK] Up Stream Non Blocking CallBack!"); +//! } +//! +//! // UpStream blocking callback (Read & Write) +//! fn us_b_callback(&mut self, _in_data: Vec) -> CallbackRet { +//! println!("[CALLBACK] Up Stream Blocking CallBack!"); +//! CallbackRet::Relay(_in_data) +//! } //! } -//! +//! //! fn main() { -//! -//! // Create new SSLRelay object -//! let mut relay = sslrelay::SSLRelay::new( -//! Handler, -//! RelayConfig { -//! downstream_data_type: TCPDataType::TLS, -//! upstream_data_type: TCPDataType::TLS, -//! bind_host: "0.0.0.0".to_string(), -//! bind_port: "443".to_string(), -//! remote_host: "remote.com".to_string(), -//! remote_port: "443".to_string(), -//! tls_config: TLSConfig::FILE{ -//! certificate_path: "./tls.crt".to_string(), -//! private_key_path: "./tls.key".to_string(), -//! }, -//! } -//! ); -//! -//! // Start listening -//! relay.start(); +//! // Create new SSLRelay object +//! let mut relay = sslrelay::SSLRelay::new( +//! Handler, +//! RelayConfig { +//! downstream_data_type: TCPDataType::TLS, +//! upstream_data_type: TCPDataType::TLS, +//! bind_host: "0.0.0.0".to_string(), +//! bind_port: "443".to_string(), +//! remote_host: |_| "remote.com", +//! remote_port: "443".to_string(), +//! tls_config: TLSConfig::FILE { +//! certificate_path: "./tls.crt".to_string(), +//! private_key_path: "./tls.key".to_string(), +//! }, +//! }, +//! ); +//! +//! // Start listening +//! relay.start(); //! } //! ``` use openssl::{ - x509::X509, - pkey::PKey, - ssl::{ - SslVerifyMode, - SslConnector, - SslAcceptor, - SslStream, - SslFiletype, - SslMethod, - } + pkey::PKey, + ssl::{SslAcceptor, SslConnector, SslFiletype, SslMethod, SslStream, SslVerifyMode}, + x509::X509, }; -use std::net::{ - TcpListener, - TcpStream, - Shutdown -}; +use std::net::{Shutdown, TcpListener, TcpStream}; -use std::sync::{ - Arc, - Mutex -}; +use std::sync::{Arc, Mutex}; -use std::{ - thread -}; +use std::thread; -use std::{ - path::Path, - time::Duration, -}; +use std::{path::Path, time::Duration}; -use std::io::{ - self, - Read, - Write -}; +use std::io::{self, Read, Write}; -use std::sync::mpsc::{ - self, - Receiver, - Sender -}; +use std::sync::mpsc::{self, Receiver, Sender}; mod data; -mod tcp; mod relay; +mod tcp; #[derive(Debug)] enum FullDuplexTcpState { - DownStreamWrite(Vec), - UpStreamWrite(Vec), - DownStreamShutDown, - UpStreamShutDown, + DownStreamWrite(Vec), + UpStreamWrite(Vec), + DownStreamShutDown, + UpStreamShutDown, } #[derive(Debug)] enum DataPipe { - DataWrite(Vec), - Shutdown, + DataWrite(Vec), + Shutdown, } enum DataStreamType { - RAW(TcpStream), - TLS(SslStream), + RAW(TcpStream), + TLS(SslStream), } /// Specifies the upstream or downstream data type (TLS or RAW). #[derive(Copy, Clone)] pub enum TCPDataType { - TLS, - RAW, + TLS, + RAW, } /// TLSConfig is used to specify TLS options. @@ -163,78 +138,89 @@ pub enum TCPDataType { /// NONE is for when you are not using TLS on the listening/downstream side of the relay. #[derive(Clone)] pub enum TLSConfig { - FILE {certificate_path: String, private_key_path: String}, - DATA {certificate: Vec, private_key: Vec}, - NONE, + FILE { + certificate_path: String, + private_key_path: String, + }, + DATA { + certificate: Vec, + private_key: Vec, + }, + NONE, } /// Relay Config structure for passing into the SSLRelay::new() config parameter. #[derive(Clone)] pub struct RelayConfig { - pub downstream_data_type: TCPDataType, - pub upstream_data_type: TCPDataType, - pub bind_host: String, - pub bind_port: String, - pub remote_host: String, - pub remote_port: String, - pub tls_config: TLSConfig, + pub downstream_data_type: TCPDataType, + pub upstream_data_type: TCPDataType, + pub bind_host: String, + pub bind_port: String, + pub remote_host: fn(Option<&str>) -> String, + //pub remote_host: String, + pub remote_port: String, + pub tls_config: TLSConfig, } /// CallbackRet for blocking callback functions #[derive(Debug)] pub enum CallbackRet { - Relay(Vec),// Relay data - Spoof(Vec),// Skip relaying and send data back - Shutdown,// Shutdown TCP connection - Freeze,// Dont send data (pretend as if stream never was recieved) + Relay(Vec), // Relay data + Spoof(Vec), // Skip relaying and send data back + Shutdown, // Shutdown TCP connection + Freeze, // Dont send data (pretend as if stream never was recieved) } /// Callback functions a user may or may not implement. pub trait HandlerCallbacks { - fn ds_b_callback(&mut self, _in_data: Vec) -> CallbackRet {CallbackRet::Relay(_in_data)} - fn ds_nb_callback(&self, _in_data: Vec){} - fn us_b_callback(&mut self, _in_data: Vec) -> CallbackRet {CallbackRet::Relay(_in_data)} - fn us_nb_callback(&self, _in_data: Vec){} + fn ds_b_callback(&mut self, _in_data: Vec, _conn_id: u64) -> CallbackRet { + CallbackRet::Relay(_in_data) + } + fn ds_nb_callback(&self, _in_data: Vec, _conn_id: u64) {} + fn us_b_callback(&mut self, _in_data: Vec, _conn_id: u64) -> CallbackRet { + CallbackRet::Relay(_in_data) + } + fn us_nb_callback(&self, _in_data: Vec, _conn_id: u64) {} + fn set_server_name(&mut self, _server_name: Option<&str>) {} } /// The main SSLRelay object. #[derive(Clone)] pub struct SSLRelay where - H: HandlerCallbacks + std::marker::Sync + std::marker::Send + Clone + 'static, + H: HandlerCallbacks + std::marker::Sync + std::marker::Send + Clone + 'static, { - config: RelayConfig, - handlers: Option>, + config: RelayConfig, + handlers: Option>, } #[allow(dead_code)] struct FullDuplexTcp where - H: HandlerCallbacks + std::marker::Sync + std::marker::Send + Clone + 'static, + H: HandlerCallbacks + std::marker::Sync + std::marker::Send + Clone + 'static, { - remote_host: String, - remote_port: String, - ds_inner_m: Arc>>, - us_inner_m: Arc>>, - inner_handlers: InnerHandlers, + remote_host: String, + remote_port: String, + ds_inner_m: Arc>>, + us_inner_m: Arc>>, + inner_handlers: InnerHandlers, + conn_id: u64, } #[derive(Clone)] struct InnerHandlers where - H: HandlerCallbacks + std::marker::Sync + std::marker::Send + Clone + 'static, + H: HandlerCallbacks + std::marker::Sync + std::marker::Send + Clone + 'static, { - cb: H + cb: H, } -struct DownStreamInner -{ - ds_stream: DataStreamType, - internal_data_buffer: Vec, +struct DownStreamInner { + ds_stream: DataStreamType, + internal_data_buffer: Vec, } -struct UpStreamInner -{ - us_stream: DataStreamType, - internal_data_buffer: Vec -} \ No newline at end of file +struct UpStreamInner { + us_stream: DataStreamType, + internal_data_buffer: Vec, +} diff --git a/src/relay.rs b/src/relay.rs index 36342bb..c5be34c 100644 --- a/src/relay.rs +++ b/src/relay.rs @@ -1,137 +1,221 @@ //! SSLRelay use crate::{ - SSLRelay, - HandlerCallbacks, - InnerHandlers, - TCPDataType, - TcpListener, - thread, - FullDuplexTcp, - DataStreamType, - RelayConfig, - Arc, - SslAcceptor, - Path, - SslMethod, - SslFiletype, - TLSConfig, - PKey, - X509, + thread, Arc, DataStreamType, FullDuplexTcp, HandlerCallbacks, InnerHandlers, PKey, Path, + RelayConfig, SSLRelay, SslAcceptor, SslFiletype, SslMethod, TCPDataType, TLSConfig, + TcpListener, X509, }; impl SSLRelay { - /// Creates new SSLRelay instance. - pub fn new(handlers: H, config: RelayConfig) -> Self { + /// Creates new SSLRelay instance. + pub fn new(handlers: H, config: RelayConfig) -> Self { + SSLRelay { + config, + handlers: Some(InnerHandlers { cb: handlers }), + } + } + /// Starts the SSLRelay connection handling. + pub fn start(&mut self) { + let rhost = self.config.remote_host.clone(); + let rport = self.config.remote_port.clone(); + let upstream_data_stream_type = self.config.upstream_data_type; - SSLRelay { - config, - handlers: Some(InnerHandlers{cb: handlers}), - } - } - /// Starts the SSLRelay connection handling. - pub fn start(&mut self) { + let listener = TcpListener::bind(format!( + "{}:{}", + self.config.bind_host.clone(), + self.config.bind_port.clone() + )) + .unwrap(); - let rhost = self.config.remote_host.clone(); - let rport = self.config.remote_port.clone(); - let upstream_data_stream_type = self.config.upstream_data_type; + let mut conn_id = 0; + match self.config.downstream_data_type { + TCPDataType::TLS => { + let acceptor = self.setup_ssl_config(&self.config.tls_config); - let listener = TcpListener::bind(format!("{}:{}", self.config.bind_host.clone(), self.config.bind_port.clone())).unwrap(); + for stream in listener.incoming() { + match stream { + Ok(stream) => { + let acceptor = acceptor.clone(); + let mut handler_clone = self.handlers.as_ref().unwrap().clone(); - match self.config.downstream_data_type { + let r_host = rhost.clone(); + let r_port = rport.clone(); - TCPDataType::TLS => { - - let acceptor = self.setup_ssl_config(self.config.tls_config.clone()); + let this_conn_id = conn_id; + thread::spawn(move || { + match acceptor.accept(stream) { + Ok(stream) => { + let server_name = stream + .ssl() + .servername(openssl::ssl::NameType::HOST_NAME); + handler_clone.cb.set_server_name(server_name); + let remote_host = (r_host)(server_name); + + // FULL DUPLEX OBJECT CREATION HERE + match FullDuplexTcp::new( + DataStreamType::TLS(stream), + upstream_data_stream_type, + remote_host, + r_port, + handler_clone, + this_conn_id, + ) { + Ok(mut fdtcp) => fdtcp.handle(), + Err(_ec) => { + println!( + "[SSLRelay Error] Failed to handle TCP(TLS) \ + connection: {}", + _ec + ) + } + } + } + Err(e) => { + println!("[Error] {}", e); + } + } + }); + conn_id += 1; + } + Err(e) => { + println!("[Error] Tcp Connection Failed: {}", e) + } + } + } + } - for stream in listener.incoming() { - - match stream { - Ok(stream) => { - - let acceptor = acceptor.clone(); - let handler_clone = self.handlers.as_ref().unwrap().clone(); - - let r_host = rhost.clone(); - let r_port = rport.clone(); - - thread::spawn(move || { - - match acceptor.accept(stream) { - Ok(stream) => { - // FULL DUPLEX OBJECT CREATION HERE - match FullDuplexTcp::new(DataStreamType::TLS(stream), upstream_data_stream_type, r_host, r_port, handler_clone) { - Ok(mut fdtcp) => fdtcp.handle(), - Err(_ec) => {println!("[SSLRelay Error] Failed to handle TCP(TLS) connection: {}", _ec)} - } - }, - Err(e) => { - - println!("[Error] {}", e); - } - } - }); - }, - Err(e) => {println!("[Error] Tcp Connection Failed: {}", e)} - } - } - }, + TCPDataType::RAW => { + for stream in listener.incoming() { + match stream { + Ok(stream) => { + let handler_clone = self.handlers.as_ref().unwrap().clone(); - TCPDataType::RAW => { + let r_host = rhost.clone(); + let r_port = rport.clone(); - for stream in listener.incoming() { - - match stream { - Ok(stream) => { - - let handler_clone = self.handlers.as_ref().unwrap().clone(); - - let r_host = rhost.clone(); - let r_port = rport.clone(); - - thread::spawn(move || { - - // FULL DUPLEX OBJECT CREATION HERE - match FullDuplexTcp::new(DataStreamType::RAW(stream), upstream_data_stream_type, r_host, r_port, handler_clone) { - Ok(mut fdtcp) => fdtcp.handle(), - Err(_ec) => println!("[SSLRelay Error] Failed to handle TCP connection: {}", _ec), - } - }); - }, - Err(e) => {println!("[Error] Tcp Connection Failed: {}", e)} - } - } - } - } - } + let this_conn_id = conn_id; + thread::spawn(move || { + // FULL DUPLEX OBJECT CREATION HERE + match FullDuplexTcp::new( + DataStreamType::RAW(stream), + upstream_data_stream_type, + (r_host)(None), + r_port, + handler_clone, + this_conn_id, + ) { + Ok(mut fdtcp) => fdtcp.handle(), + Err(_ec) => println!( + "[SSLRelay Error] Failed to handle TCP connection: {}", + _ec + ), + } + }); + conn_id += 1; + } + Err(e) => { + println!("[Error] Tcp Connection Failed: {}", e) + } + } + } + } + } + } - fn setup_ssl_config(&self, tls_config: TLSConfig) -> Arc { + fn setup_ssl_config(&self, tls_config: &TLSConfig) -> Arc { + let mut acceptor = SslAcceptor::mozilla_intermediate(SslMethod::tls()).unwrap(); + match tls_config { + TLSConfig::FILE { + certificate_path, + private_key_path, + } => { + if !Path::new(&private_key_path).exists() { + panic!("[-] [{}] does not exist!", private_key_path); + } + if !Path::new(&certificate_path).exists() { + panic!("[-] [{}] does not exist!", certificate_path); + } + acceptor + .set_private_key_file(private_key_path, SslFiletype::PEM) + .unwrap(); + acceptor + .set_certificate_chain_file(certificate_path) + .unwrap(); + acceptor.check_private_key().unwrap(); + } + TLSConfig::DATA { + certificate, + private_key, + } => { + let x_509_certificate = X509::from_pem(certificate.as_slice()).unwrap(); + let private_key = PKey::private_key_from_pem(private_key.as_slice()).unwrap(); + acceptor + .set_certificate(x_509_certificate.as_ref()) + .unwrap(); + acceptor.set_private_key(private_key.as_ref()).unwrap(); + } + TLSConfig::NONE => { + panic!( + "[SSLRelay Error] Specified NONE for TLSConfig and downstream_data_type as \ + TLS." + ); + } + } + Arc::new(acceptor.build()) + } - let mut acceptor = SslAcceptor::mozilla_intermediate(SslMethod::tls()).unwrap(); + // fn setup_ssl_config(&self, tls_configs: &[(&[&str], TLSConfig)]) -> AcceptorDb { + // let mut acceptor_db = AcceptorDb { + // acceptors: HashMap::new(), + // }; + // for (names, tls_config) in tls_configs { + // let mut acceptor = SslAcceptor::mozilla_intermediate(SslMethod::tls()).unwrap(); + // match tls_config { + // TLSConfig::FILE { + // certificate_path, + // private_key_path, + // } => { + // if !Path::new(&private_key_path).exists() { + // panic!("[-] [{}] does not exist!", private_key_path); + // } + // if !Path::new(&certificate_path).exists() { + // panic!("[-] [{}] does not exist!", certificate_path); + // } + // acceptor + // .set_private_key_file(private_key_path, SslFiletype::PEM) + // .unwrap(); + // acceptor + // .set_certificate_chain_file(certificate_path) + // .unwrap(); + // acceptor.check_private_key().unwrap(); + // } + // TLSConfig::DATA { + // certificate, + // private_key, + // } => { + // let x_509_certificate = X509::from_pem(certificate.as_slice()).unwrap(); + // let private_key = PKey::private_key_from_pem(private_key.as_slice()).unwrap(); + // acceptor + // .set_certificate(x_509_certificate.as_ref()) + // .unwrap(); + // acceptor.set_private_key(private_key.as_ref()).unwrap(); + // } + // TLSConfig::NONE => { + // panic!( + // "[SSLRelay Error] Specified NONE for TLSConfig and downstream_data_type as \ + // TLS." + // ); + // } + // } + // let acceptor = Arc::new(acceptor.build()); + // for name in *names { + // acceptor_db.acceptors.insert(name.to_string(), acceptor.clone()); + // } + // } + // acceptor_db + // } +} // SSLRelay - match tls_config { - TLSConfig::FILE{certificate_path, private_key_path} => { - - if !Path::new(&private_key_path).exists() { - panic!("[-] [{}] does not exist!", private_key_path); - } - if !Path::new(&certificate_path).exists() { - panic!("[-] [{}] does not exist!", certificate_path); - } - acceptor.set_private_key_file(private_key_path, SslFiletype::PEM).unwrap(); - acceptor.set_certificate_chain_file(certificate_path).unwrap(); - acceptor.check_private_key().unwrap(); - }, - TLSConfig::DATA{certificate, private_key} => { - let x_509_certificate = X509::from_pem(certificate.as_slice()).unwrap(); - let private_key = PKey::private_key_from_pem(private_key.as_slice()).unwrap(); - acceptor.set_certificate(x_509_certificate.as_ref()).unwrap(); - acceptor.set_private_key(private_key.as_ref()).unwrap(); - }, - TLSConfig::NONE => { - panic!("[SSLRelay Error] Specified NONE for TLSConfig and downstream_data_type as TLS."); - } - } - Arc::new(acceptor.build()) - } -}// SSLRelay \ No newline at end of file +// struct AcceptorDb { +// acceptors: HashMap>, +// } diff --git a/src/tcp.rs b/src/tcp.rs index 99d0592..115ae6d 100644 --- a/src/tcp.rs +++ b/src/tcp.rs @@ -1,259 +1,368 @@ use crate::{ - FullDuplexTcp, - HandlerCallbacks, - DataStreamType, - TCPDataType, - Duration, - Arc, - Mutex, - DownStreamInner, - UpStreamInner, - InnerHandlers, - Shutdown, - Sender, - Receiver, - FullDuplexTcpState, - DataPipe, - mpsc, - thread, - CallbackRet, - TcpStream, - SslVerifyMode, - SslConnector, - SslMethod, + mpsc, thread, Arc, CallbackRet, DataPipe, DataStreamType, DownStreamInner, Duration, + FullDuplexTcp, FullDuplexTcpState, HandlerCallbacks, InnerHandlers, Mutex, Receiver, Sender, + Shutdown, SslConnector, SslMethod, SslVerifyMode, TCPDataType, TcpStream, UpStreamInner, }; -impl FullDuplexTcp { +impl + FullDuplexTcp +{ + pub fn new( + ds_tcp_stream: DataStreamType, + us_tcp_stream_type: TCPDataType, + remote_host: String, + remote_port: String, + handlers: InnerHandlers, + conn_id: u64, + ) -> Result { + match ds_tcp_stream { + DataStreamType::RAW(ref s) => { + let _ = s.set_read_timeout(Some(Duration::from_millis(50))); + } + DataStreamType::TLS(ref s) => { + let _ = s + .get_ref() + .set_read_timeout(Some(Duration::from_millis(50))); + } + } - pub fn new(ds_tcp_stream: DataStreamType, us_tcp_stream_type: TCPDataType, remote_host: String, remote_port: String, handlers: InnerHandlers) -> Result { + let us_tcp_stream = match Self::connect_endpoint( + us_tcp_stream_type, + remote_host.clone(), + remote_port.clone(), + ) { + Ok(s) => s, + Err(ec) => { + match ds_tcp_stream { + DataStreamType::RAW(s) => { + let _ = s.shutdown(Shutdown::Both); + } + DataStreamType::TLS(mut s) => { + let _ = s.shutdown(); + } + } + return Err(ec); + } + }; - match ds_tcp_stream { - DataStreamType::RAW(ref s) => { let _ = s.set_read_timeout(Some(Duration::from_millis(50))); }, - DataStreamType::TLS(ref s) => { let _ = s.get_ref().set_read_timeout(Some(Duration::from_millis(50))); }, - } + Ok(FullDuplexTcp { + remote_host, + remote_port, + ds_inner_m: Arc::new(Mutex::new(Some(DownStreamInner { + ds_stream: ds_tcp_stream, + internal_data_buffer: Vec::::new(), + }))), + us_inner_m: Arc::new(Mutex::new(Some(UpStreamInner { + us_stream: us_tcp_stream, + internal_data_buffer: Vec::::new(), + }))), + inner_handlers: handlers, + conn_id, + }) + } - let us_tcp_stream = match Self::connect_endpoint(us_tcp_stream_type, remote_host.clone(), remote_port.clone()) { - Ok(s) => s, - Err(ec) => { - match ds_tcp_stream { - DataStreamType::RAW(s) => { let _ = s.shutdown(Shutdown::Both); }, - DataStreamType::TLS(mut s) => { let _ = s.shutdown(); }, - } - return Err(ec); - } - }; + pub fn handle(&mut self) { + let conn_id = self.conn_id; + let (state_sender, state_receiver): ( + Sender, + Receiver, + ) = mpsc::channel(); + let (ds_data_pipe_sender, ds_data_pipe_receiver): (Sender, Receiver) = + mpsc::channel(); + let (us_data_pipe_sender, us_data_pipe_receiver): (Sender, Receiver) = + mpsc::channel(); - Ok( - FullDuplexTcp { - remote_host, - remote_port, - ds_inner_m: Arc::new(Mutex::new(Some(DownStreamInner{ds_stream: ds_tcp_stream, internal_data_buffer: Vec::::new()}))), - us_inner_m: Arc::new(Mutex::new(Some(UpStreamInner{us_stream: us_tcp_stream, internal_data_buffer: Vec::::new()}))), - inner_handlers: handlers, - }) - } + let ds_method_pointer = self.ds_inner_m.clone(); + let ds_state_bc = state_sender.clone(); - pub fn handle(&mut self) { + let us_method_pointer = self.us_inner_m.clone(); + let us_state_bc = state_sender.clone(); - let (state_sender, state_receiver): (Sender, Receiver) = mpsc::channel(); - let (ds_data_pipe_sender, ds_data_pipe_receiver): (Sender, Receiver) = mpsc::channel(); - let (us_data_pipe_sender, us_data_pipe_receiver): (Sender, Receiver) = mpsc::channel(); + thread::spawn(move || { + ds_method_pointer + .lock() + .unwrap() + .take() + .unwrap() + .ds_handler(ds_state_bc, ds_data_pipe_receiver); + }); - let ds_method_pointer = self.ds_inner_m.clone(); - let ds_state_bc = state_sender.clone(); + thread::spawn(move || { + us_method_pointer + .lock() + .unwrap() + .take() + .unwrap() + .us_handler(us_state_bc, us_data_pipe_receiver); + }); - let us_method_pointer = self.us_inner_m.clone(); - let us_state_bc = state_sender.clone(); + loop { + match state_receiver.recv() { + Ok(state_request) => { + match state_request { + // DownStream Write Request + FullDuplexTcpState::DownStreamWrite(data) => { + /* + Callbacks that work with data from UpStream go here + Add callback return types for blocking callback subroutines + Shutdown - Shutdown TCP connection + Relay - Relay TCP stream + Spoof - Spoof back to received stream direction + Freeze - Freeze data (dont relay and destroy data) + */ - thread::spawn(move || { - ds_method_pointer.lock().unwrap().take().unwrap().ds_handler(ds_state_bc, ds_data_pipe_receiver); - }); + let inner_handlers_clone = self.inner_handlers.clone(); + let in_data = data.clone(); - thread::spawn(move || { - us_method_pointer.lock().unwrap().take().unwrap().us_handler(us_state_bc, us_data_pipe_receiver); - }); + thread::spawn(move || { + inner_handlers_clone.cb.us_nb_callback(in_data, conn_id); + }); - loop { + match self.inner_handlers.cb.us_b_callback(data, conn_id) { + CallbackRet::Relay(retdata) => { + match ds_data_pipe_sender.send(DataPipe::DataWrite(retdata)) { + Ok(()) => {} + Err(e) => { + Self::handle_error( + format!( + "Failed to send data write to DownStream \ + thread: {}", + e + ) + .as_str(), + ); + return; + } + } + } + CallbackRet::Spoof(retdata) => { + match us_data_pipe_sender.send(DataPipe::DataWrite(retdata)) { + Ok(()) => {} + Err(e) => { + Self::handle_error( + format!( + "Failed to send data write to DownStream \ + thread: {}", + e + ) + .as_str(), + ); + return; + } + } + } + CallbackRet::Freeze => {} + CallbackRet::Shutdown => { + if let Err(e) = us_data_pipe_sender.send(DataPipe::Shutdown) { + Self::handle_error( + format!( + "Failed to send Shutdown signal to UpStream \ + thread: {}", + e + ) + .as_str(), + ); + } + if let Err(e) = ds_data_pipe_sender.send(DataPipe::Shutdown) { + Self::handle_error( + format!( + "Failed to send Shutdown signal to DownStream \ + thread: {}", + e + ) + .as_str(), + ); + } + return; + } + } + } + // UpStream Write Request + FullDuplexTcpState::UpStreamWrite(data) => { + /* + Callbacks that work with data from DownStream go here + */ - match state_receiver.recv() { - Ok(state_request) => { - match state_request { + let inner_handlers_clone = self.inner_handlers.clone(); + let in_data = data.clone(); - // DownStream Write Request - FullDuplexTcpState::DownStreamWrite(data) => { + thread::spawn(move || { + inner_handlers_clone.cb.ds_nb_callback(in_data, conn_id); + }); - /* - Callbacks that work with data from UpStream go here - Add callback return types for blocking callback subroutines - Shutdown - Shutdown TCP connection - Relay - Relay TCP stream - Spoof - Spoof back to received stream direction - Freeze - Freeze data (dont relay and destroy data) - */ + match self.inner_handlers.cb.ds_b_callback(data, conn_id) { + CallbackRet::Relay(retdata) => { + match us_data_pipe_sender.send(DataPipe::DataWrite(retdata)) { + Ok(()) => {} + Err(e) => { + Self::handle_error( + format!( + "Failed to send data write to UpStream \ + thread: {}", + e + ) + .as_str(), + ); + return; + } + } + } + CallbackRet::Spoof(retdata) => { + match ds_data_pipe_sender.send(DataPipe::DataWrite(retdata)) { + Ok(()) => {} + Err(e) => { + Self::handle_error( + format!( + "Failed to send data write to DownStream \ + thread: {}", + e + ) + .as_str(), + ); + return; + } + } + } + CallbackRet::Freeze => {} + CallbackRet::Shutdown => { + if let Err(e) = ds_data_pipe_sender.send(DataPipe::Shutdown) { + Self::handle_error( + format!( + "Failed to send Shutdown signal to DownStream \ + thread: {}", + e + ) + .as_str(), + ); + } + if let Err(e) = us_data_pipe_sender.send(DataPipe::Shutdown) { + Self::handle_error( + format!( + "Failed to send Shutdown signal to UpStream \ + thread: {}", + e + ) + .as_str(), + ); + } + return; + } + } + } + // DownStreamShutDown Request + FullDuplexTcpState::DownStreamShutDown => { + if let Err(e) = us_data_pipe_sender.send(DataPipe::Shutdown) { + Self::handle_error( + format!( + "Failed to send Shutdown signal to UpStream thread: {}", + e + ) + .as_str(), + ); + return; + } + return; + } + // UpStreamShutDown Request + FullDuplexTcpState::UpStreamShutDown => { + if let Err(e) = ds_data_pipe_sender.send(DataPipe::Shutdown) { + Self::handle_error( + format!( + "Failed to send Shutdown signal to DownStream thread: {}", + e + ) + .as_str(), + ); + return; + } + return; + } + } + } + Err(_e) => { + Self::handle_error("State receiver communication channel has closed!"); + if let Err(e) = ds_data_pipe_sender.send(DataPipe::Shutdown) { + Self::handle_error( + format!("Failed to send Shutdown signal to DownStream thread: {}", e) + .as_str(), + ); + } + if let Err(e) = us_data_pipe_sender.send(DataPipe::Shutdown) { + Self::handle_error( + format!("Failed to send Shutdown signal to UpStream thread: {}", e) + .as_str(), + ); + } + return; + } + } // State Receiver + } + } - let inner_handlers_clone = self.inner_handlers.clone(); - let in_data = data.clone(); + fn connect_endpoint( + stream_data_type: TCPDataType, + remote_host: String, + remote_port: String, + ) -> Result { + match stream_data_type { + TCPDataType::RAW => { + let s = match TcpStream::connect(format!("{}:{}", remote_host, remote_port)) { + Ok(s) => s, + Err(e) => { + Self::handle_error( + format!( + "Can't connect to remote host: {}\nErr: {}", + format!("{}:{}", remote_host, remote_port), + e + ) + .as_str(), + ); + return Result::Err(-1); + } + }; + let _ = s.set_read_timeout(Some(Duration::from_millis(50))); + return Ok(DataStreamType::RAW(s)); + } + TCPDataType::TLS => { + let mut sslbuilder = SslConnector::builder(SslMethod::tls()).unwrap(); + sslbuilder.set_verify(SslVerifyMode::NONE); - thread::spawn(move || { - inner_handlers_clone.cb.us_nb_callback(in_data); - }); + let connector = sslbuilder.build(); - match self.inner_handlers.cb.us_b_callback(data) { - CallbackRet::Relay(retdata) => { - match ds_data_pipe_sender.send(DataPipe::DataWrite(retdata)) { - Ok(()) => {}, - Err(e) => { - Self::handle_error(format!("Failed to send data write to DownStream thread: {}", e).as_str()); - return; - } - } - }, - CallbackRet::Spoof(retdata) => { - match us_data_pipe_sender.send(DataPipe::DataWrite(retdata)) { - Ok(()) => {}, - Err(e) => { - Self::handle_error(format!("Failed to send data write to DownStream thread: {}", e).as_str()); - return; - } - } - }, - CallbackRet::Freeze => {}, - CallbackRet::Shutdown => { - if let Err(e) = us_data_pipe_sender.send(DataPipe::Shutdown) { - Self::handle_error(format!("Failed to send Shutdown signal to UpStream thread: {}", e).as_str()); - } - if let Err(e) = ds_data_pipe_sender.send(DataPipe::Shutdown) { - Self::handle_error(format!("Failed to send Shutdown signal to DownStream thread: {}", e).as_str()); - } - return; - } - } - }, - // UpStream Write Request - FullDuplexTcpState::UpStreamWrite(data) => { + let s = match TcpStream::connect(format!("{}:{}", remote_host, remote_port)) { + Ok(s) => s, + Err(e) => { + Self::handle_error( + format!( + "Can't connect to remote host: {}\nErr: {}", + format!("{}:{}", remote_host, remote_port), + e + ) + .as_str(), + ); + return Result::Err(-1); + } + }; - /* - Callbacks that work with data from DownStream go here - */ + let s = match connector.connect(remote_host.as_str(), s) { + Ok(s) => s, + Err(e) => { + Self::handle_error( + format!("Failed to accept TLS/SSL handshake: {}", e).as_str(), + ); + return Result::Err(-2); + } + }; - let inner_handlers_clone = self.inner_handlers.clone(); - let in_data = data.clone(); + let _ = s + .get_ref() + .set_read_timeout(Some(Duration::from_millis(50))); + return Ok(DataStreamType::TLS(s)); + } + } + } - thread::spawn(move || { - inner_handlers_clone.cb.ds_nb_callback(in_data); - }); - - match self.inner_handlers.cb.ds_b_callback(data) { - CallbackRet::Relay(retdata) => { - match us_data_pipe_sender.send(DataPipe::DataWrite(retdata)) { - Ok(()) => {}, - Err(e) => { - Self::handle_error(format!("Failed to send data write to UpStream thread: {}", e).as_str()); - return; - } - } - }, - CallbackRet::Spoof(retdata) => { - match ds_data_pipe_sender.send(DataPipe::DataWrite(retdata)) { - Ok(()) => {}, - Err(e) => { - Self::handle_error(format!("Failed to send data write to DownStream thread: {}", e).as_str()); - return; - } - } - }, - CallbackRet::Freeze => {}, - CallbackRet::Shutdown => { - if let Err(e) = ds_data_pipe_sender.send(DataPipe::Shutdown) { - Self::handle_error(format!("Failed to send Shutdown signal to DownStream thread: {}", e).as_str()); - } - if let Err(e) = us_data_pipe_sender.send(DataPipe::Shutdown) { - Self::handle_error(format!("Failed to send Shutdown signal to UpStream thread: {}", e).as_str()); - } - return; - } - } - }, - // DownStreamShutDown Request - FullDuplexTcpState::DownStreamShutDown => { - - if let Err(e) = us_data_pipe_sender.send(DataPipe::Shutdown) { - Self::handle_error(format!("Failed to send Shutdown signal to UpStream thread: {}", e).as_str()); - return; - } - return; - }, - // UpStreamShutDown Request - FullDuplexTcpState::UpStreamShutDown => { - - if let Err(e) = ds_data_pipe_sender.send(DataPipe::Shutdown) { - Self::handle_error(format!("Failed to send Shutdown signal to DownStream thread: {}", e).as_str()); - return; - } - return; - }, - } - }, - Err(_e) => { - Self::handle_error("State receiver communication channel has closed!"); - if let Err(e) = ds_data_pipe_sender.send(DataPipe::Shutdown) { - Self::handle_error(format!("Failed to send Shutdown signal to DownStream thread: {}", e).as_str()); - } - if let Err(e) = us_data_pipe_sender.send(DataPipe::Shutdown) { - Self::handle_error(format!("Failed to send Shutdown signal to UpStream thread: {}", e).as_str()); - } - return; - } - }// State Receiver - } - } - - fn connect_endpoint(stream_data_type: TCPDataType, remote_host: String, remote_port: String) -> Result { - - match stream_data_type { - - TCPDataType::RAW => { - let s = match TcpStream::connect(format!("{}:{}", remote_host, remote_port)) { - Ok(s) => s, - Err(e) => { - Self::handle_error(format!("Can't connect to remote host: {}\nErr: {}", format!("{}:{}", remote_host, remote_port), e).as_str()); - return Result::Err(-1); - } - }; - let _ = s.set_read_timeout(Some(Duration::from_millis(50))); - return Ok(DataStreamType::RAW(s)); - - }, - TCPDataType::TLS => { - - let mut sslbuilder = SslConnector::builder(SslMethod::tls()).unwrap(); - sslbuilder.set_verify(SslVerifyMode::NONE); - - let connector = sslbuilder.build(); - - let s = match TcpStream::connect(format!("{}:{}", remote_host, remote_port)) { - Ok(s) => s, - Err(e) => { - Self::handle_error(format!("Can't connect to remote host: {}\nErr: {}", format!("{}:{}", remote_host, remote_port), e).as_str()); - return Result::Err(-1); - } - }; - - let s = match connector.connect(remote_host.as_str(), s) { - Ok(s) => s, - Err(e) => { - Self::handle_error(format!("Failed to accept TLS/SSL handshake: {}", e).as_str()); - return Result::Err(-2); - } - }; - - let _ = s.get_ref().set_read_timeout(Some(Duration::from_millis(50))); - return Ok(DataStreamType::TLS(s)); - } - } - } - - fn handle_error(error_description: &str) { - println!("[SSLRelay Master Thread Error]: {}", error_description); - } -} \ No newline at end of file + fn handle_error(error_description: &str) { + println!("[SSLRelay Master Thread Error]: {}", error_description); + } +}