tuxmain/sslrelay
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Organizing error handling

Browse source
This commit is contained in:
PinkP4nther 2021-09-16 18:23:25 -07:00
commit e6a3420814
2 changed files with 34 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

2
examples/modifydata/src/main.rs
View file

@ -1,4 +1,4 @@
use sslrelay::{self, ConfigType, HandlerCallbacks}; use sslrelay::{self, ConfigType, HandlerCallbacks, CallbackRet};
// Handler object // Handler object
#[derive(Clone)] // Must have Clone trait implemented. #[derive(Clone)] // Must have Clone trait implemented.

61
src/data.rs
View file

@ -28,6 +28,12 @@ struct DownStreamInner {
} }
impl DownStreamInner { impl DownStreamInner {
fn handle_error(&self, error_description: &str) {
println!("[SSLRelay DownStream Thread Error]: {}", error_description);
let _ = self.ds_stream.as_ref().unwrap().lock().unwrap().get_ref().shutdown(Shutdown::Both);
}
pub fn ds_handler(&mut self, data_out: Sender<FullDuplexTcpState>, data_in: Receiver<DataPipe>) { pub fn ds_handler(&mut self, data_out: Sender<FullDuplexTcpState>, data_in: Receiver<DataPipe>) {
loop { loop {
@ -43,12 +49,10 @@ impl DownStreamInner {
let mut stream_lock = match self.ds_stream.as_ref().unwrap().lock() { let mut stream_lock = match self.ds_stream.as_ref().unwrap().lock() {
Ok(sl) => sl, Ok(sl) => sl,
Err(_e) => { Err(_e) => {
println!("[!] Failed to get stream lock!"); self.handle_error("Failed to get stream lock!");
if let Err(e) = data_out.send(FullDuplexTcpState::DownStreamShutDown) { if let Err(e) = data_out.send(FullDuplexTcpState::DownStreamShutDown) {
println!("[SSLRelay DownStream Thread Error]: Failed to send shutdown signal to main thread from DownStream thread: {}", e); self.handle_error(format!("Failed to send shutdown signal to main thread from DownStream thread: {}", e).as_str());
return;
} }
let _ = self.ds_stream.as_ref().unwrap().lock().unwrap().get_ref().shutdown(Shutdown::Both);
return; return;
} }
}; };
@ -56,12 +60,11 @@ impl DownStreamInner {
match stream_lock.write_all(&data) { match stream_lock.write_all(&data) {
Ok(()) => {}, Ok(()) => {},
Err(_e) => { Err(_e) => {
println!("[!] Failed to write data to DownStream tcp stream!"); self.handle_error("Failed to write data to DownStream tcp stream!");
if let Err(e) = data_out.send(FullDuplexTcpState::DownStreamShutDown) { if let Err(e) = data_out.send(FullDuplexTcpState::DownStreamShutDown) {
println!("[SSLRelay DownStream Thread Error]: Failed to send shutdown signal to main thread from DownStream thread: {}", e); self.handle_error(format!("Failed to send shutdown signal to main thread from DownStream thread: {}", e).as_str());
return;
} }
let _ = self.ds_stream.as_ref().unwrap().lock().unwrap().get_ref().shutdown(Shutdown::Both); return;
} }
} }
let _ = stream_lock.flush(); let _ = stream_lock.flush();
@ -78,7 +81,7 @@ impl DownStreamInner {
match _e { match _e {
mpsc::RecvTimeoutError::Timeout => {}, mpsc::RecvTimeoutError::Timeout => {},
mpsc::RecvTimeoutError::Disconnected => { mpsc::RecvTimeoutError::Disconnected => {
println!("[!] DownStream data_in channel is disconnected!"); self.handle_error("DownStream data_in channel is disconnected!");
return; return;
} }
} }
@ -90,7 +93,7 @@ impl DownStreamInner {
if byte_count > 0 { if byte_count > 0 {
if let Err(e) = data_out.send(FullDuplexTcpState::UpStreamWrite(self.internal_data_buffer.clone())) { if let Err(e) = data_out.send(FullDuplexTcpState::UpStreamWrite(self.internal_data_buffer.clone())) {
println!("[SSLRelay DownStream Thread Error]: Failed to send UpStreamWrite to main thread: {}", e); self.handle_error(format!("Failed to send UpStreamWrite to main thread: {}", e).as_str());
return; return;
} }
@ -99,11 +102,10 @@ impl DownStreamInner {
} else if byte_count == 0 || byte_count == -2 { } else if byte_count == 0 || byte_count == -2 {
if let Err(e) = data_out.send(FullDuplexTcpState::DownStreamShutDown) { if let Err(e) = data_out.send(FullDuplexTcpState::DownStreamShutDown) {
println!("[SSLRelay DownStream Thread Error]: Failed to send shutdown signal to main thread from DownStream thread: {}", e); self.handle_error(format!("Failed to send shutdown signal to main thread from DownStream thread: {}", e).as_str());
return;
} }
let _ = self.ds_stream.as_ref().unwrap().lock().unwrap().get_ref().shutdown(Shutdown::Both);
return; return;
} else if byte_count == -1 { } else if byte_count == -1 {
continue; continue;
} }
@ -177,6 +179,12 @@ struct UpStreamInner{
} }
impl UpStreamInner { impl UpStreamInner {
fn handle_error(&self, error_description: &str) {
println!("[SSLRelay UpStream Thread Error]: {}", error_description);
let _ = self.us_stream.as_ref().unwrap().lock().unwrap().get_ref().shutdown(Shutdown::Both);
}
pub fn us_handler(&mut self, data_out: Sender<FullDuplexTcpState>, data_in: Receiver<DataPipe>) { pub fn us_handler(&mut self, data_out: Sender<FullDuplexTcpState>, data_in: Receiver<DataPipe>) {
loop { loop {
@ -191,12 +199,11 @@ impl UpStreamInner {
let mut stream_lock = match self.us_stream.as_ref().unwrap().lock() { let mut stream_lock = match self.us_stream.as_ref().unwrap().lock() {
Ok(sl) => sl, Ok(sl) => sl,
Err(_e) => { Err(_e) => {
println!("[!] Failed to get stream lock!"); self.handle_error("Failed to get stream lock!");
if let Err(e) = data_out.send(FullDuplexTcpState::UpStreamShutDown) { if let Err(_e) = data_out.send(FullDuplexTcpState::UpStreamShutDown) {
println!("[SSLRelay UpStream Thread Error]: Failed to send shutdown signal to main thread from UpStream thread: {}", e); //println!("[SSLRelay UpStream Thread Error]: Failed to send shutdown signal to main thread from UpStream thread: {}", e);
return; //return;
} }
let _ = self.us_stream.as_ref().unwrap().lock().unwrap().get_ref().shutdown(Shutdown::Both);
return; return;
} }
}; };
@ -204,12 +211,12 @@ impl UpStreamInner {
match stream_lock.write_all(&data) { match stream_lock.write_all(&data) {
Ok(()) => {}, Ok(()) => {},
Err(_e) => { Err(_e) => {
println!("[!] Failed to write data to UpStream tcp stream!"); println!("Failed to write data to UpStream tcp stream!");
if let Err(e) = data_out.send(FullDuplexTcpState::UpStreamShutDown) { if let Err(_e) = data_out.send(FullDuplexTcpState::UpStreamShutDown) {
println!("[SSLRelay UpStream Thread Error]: Failed to send shutdown signal to main thread from UpStream thread: {}", e); //println!("[SSLRelay UpStream Thread Error]: Failed to send shutdown signal to main thread from UpStream thread: {}", e);
return; //return;
} }
let _ = self.us_stream.as_ref().unwrap().lock().unwrap().get_ref().shutdown(Shutdown::Both); return;
} }
} }
let _ = stream_lock.flush(); let _ = stream_lock.flush();
@ -226,7 +233,7 @@ impl UpStreamInner {
match e { match e {
mpsc::RecvTimeoutError::Timeout => {}, mpsc::RecvTimeoutError::Timeout => {},
mpsc::RecvTimeoutError::Disconnected => { mpsc::RecvTimeoutError::Disconnected => {
println!("[!] UpStream data_in channel is disconnected!"); self.handle_error("UpStream data_in channel is disconnected!");
return; return;
} }
} }
@ -237,7 +244,7 @@ impl UpStreamInner {
if byte_count > 0 { if byte_count > 0 {
if let Err(e) = data_out.send(FullDuplexTcpState::DownStreamWrite(self.internal_data_buffer.clone())) { if let Err(e) = data_out.send(FullDuplexTcpState::DownStreamWrite(self.internal_data_buffer.clone())) {
println!("[SSLRelay UpStream Thread Error]: Failed to send DownStreamWrite to main thread: {}", e); self.handle_error(format!("Failed to send DownStreamWrite to main thread: {}", e).as_str());
return; return;
} }
@ -246,10 +253,8 @@ impl UpStreamInner {
} else if byte_count == 0 || byte_count == -2 { } else if byte_count == 0 || byte_count == -2 {
if let Err(e) = data_out.send(FullDuplexTcpState::UpStreamShutDown) { if let Err(e) = data_out.send(FullDuplexTcpState::UpStreamShutDown) {
println!("[SSLRelay UpStream Thread Error]: Failed to send shutdown signal to main thread from UpStream thread: {}", e); self.handle_error(format!("Failed to send shutdown signal to main thread from UpStream thread: {}", e).as_str());
return;
} }
let _ = self.us_stream.as_ref().unwrap().lock().unwrap().get_ref().shutdown(Shutdown::Both);
return; return;
} else if byte_count == -1 { } else if byte_count == -1 {
continue; continue;