From 953ff1fda585d8f35c935bba61060eb7dad48d41 Mon Sep 17 00:00:00 2001 From: PinkP4nther <0x0090@protonmail.com> Date: Thu, 16 Sep 2021 02:11:49 -0700 Subject: [PATCH] Added Callback return types. --- examples/basic/src/main.rs | 8 ++-- examples/modifydata/src/main.rs | 6 ++- src/data.rs | 76 +++++++++++++++++++++++++++------ src/lib.rs | 12 +++++- 4 files changed, 82 insertions(+), 20 deletions(-) diff --git a/examples/basic/src/main.rs b/examples/basic/src/main.rs index bf1440c..5c8d4d9 100644 --- a/examples/basic/src/main.rs +++ b/examples/basic/src/main.rs @@ -1,4 +1,4 @@ -use sslrelay::{self, ConfigType, RelayConfig, HandlerCallbacks}; +use sslrelay::{self, ConfigType, RelayConfig, HandlerCallbacks, CallbackRet}; // Handler object #[derive(Clone)] // Must have Clone trait implemented. @@ -16,8 +16,9 @@ impl HandlerCallbacks for Handler { } // DownStream blocking callback - fn ds_b_callback(&self, _in_data: &mut Vec) { + fn ds_b_callback(&self, _in_data: Vec) { println!("[CALLBACK] Down Stream Blocking CallBack!"); + CallbackRet::Relay(_in_data) } // UpStream non blocking callback @@ -26,8 +27,9 @@ impl HandlerCallbacks for Handler { } // UpStream blocking callback - fn us_b_callback(&self, _in_data: &mut Vec) { + fn us_b_callback(&self, _in_data: Vec) { println!("[CALLBACK] Up Stream Blocking CallBack!"); + CallbackRet::Relay(_in_data) } } diff --git a/examples/modifydata/src/main.rs b/examples/modifydata/src/main.rs index c05bc41..8b8fac7 100644 --- a/examples/modifydata/src/main.rs +++ b/examples/modifydata/src/main.rs @@ -16,9 +16,10 @@ impl HandlerCallbacks for Handler { } // DownStream blocking callback - fn ds_b_callback(&self, _in_data: &mut Vec) { + fn ds_b_callback(&self, _in_data: Vec) { _in_data.reverse(); println!("[+] Data rewritten to:\n{:#04X?}", _in_data); + CallbackRet::Relay(_in_data) } // UpStream non blocking callback @@ -27,9 +28,10 @@ impl HandlerCallbacks for Handler { } // UpStream blocking callback - fn us_b_callback(&self, _in_data: &mut Vec) { + fn us_b_callback(&self, _in_data: Vec) { _in_data.reverse(); println!("[+] Data rewritten to:\n{:#04X?}", _in_data); + CallbackRet::Relay(_in_data) } } diff --git a/src/data.rs b/src/data.rs index 6151d29..85e1981 100644 --- a/src/data.rs +++ b/src/data.rs @@ -6,7 +6,7 @@ use std::sync::mpsc::{self, Receiver, Sender}; use std::thread; use std::sync::{Arc, Mutex}; -use crate::{HandlerCallbacks, InnerHandlers}; +use crate::{HandlerCallbacks, CallbackRet, InnerHandlers}; #[derive(Debug)] enum FullDuplexTcpState { @@ -378,10 +378,15 @@ impl { + 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) */ let inner_handlers_clone = self.inner_handlers.clone(); @@ -391,17 +396,40 @@ impl {}, - Err(e) => { - println!("[SSLRelay Error]: Failed to send data write to DownStream thread: {}", e); + match self.inner_handlers.cb.us_b_callback(data) { + CallbackRet::Relay(retdata) => { + match ds_data_pipe_sender.send(DataPipe::DataWrite(retdata)) { + Ok(()) => {}, + Err(e) => { + println!("[SSLRelay Error]: Failed to send data write to DownStream thread: {}", e); + return; + } + } + }, + CallbackRet::Spoof(retdata) => { + match us_data_pipe_sender.send(DataPipe::DataWrite(retdata)) { + Ok(()) => {}, + Err(e) => { + println!("[SSLRelay Error]: Failed to send data write to DownStream thread: {}", e); + return; + } + } + }, + CallbackRet::Freeze => {}, + CallbackRet::Shutdown => { + if let Err(e) = us_data_pipe_sender.send(DataPipe::Shutdown) { + println!("[SSLRelay Error]: Failed to send Shutdown signal to UpStream thread: {}", e); + } + if let Err(e) = ds_data_pipe_sender.send(DataPipe::Shutdown) { + println!("[SSLRelay Error]: Failed to send Shutdown signal to DownStream thread: {}", e); + } return; } } + }, // UpStream Write Request - FullDuplexTcpState::UpStreamWrite(mut data) => { + FullDuplexTcpState::UpStreamWrite(data) => { /* Callbacks that work with data from DownStream go here @@ -414,12 +442,34 @@ impl { + match us_data_pipe_sender.send(DataPipe::DataWrite(retdata)) { + Ok(()) => {}, + Err(e) => { + println!("[SSLRelay Error]: Failed to send data write to UpStream thread: {}", e); + return; + } + } + }, + CallbackRet::Spoof(retdata) => { + match ds_data_pipe_sender.send(DataPipe::DataWrite(retdata)) { + Ok(()) => {}, + Err(e) => { + println!("[SSLRelay Error]: Failed to send data write to DownStream thread: {}", e); + return; + } + } + }, + CallbackRet::Freeze => {}, + CallbackRet::Shutdown => { + if let Err(e) = ds_data_pipe_sender.send(DataPipe::Shutdown) { + println!("[SSLRelay Error]: Failed to send Shutdown signal to DownStream thread: {}", e); - match us_data_pipe_sender.send(DataPipe::DataWrite(data)) { - Ok(()) => {}, - Err(e) => { - println!("[SSLRelay Error]: Failed to send data write to UpStream thread: {}", e); + } + if let Err(e) = us_data_pipe_sender.send(DataPipe::Shutdown) { + println!("[SSLRelay Error]: Failed to send Shutdown signal to UpStream thread: {}", e); + } return; } } diff --git a/src/lib.rs b/src/lib.rs index 62e71bb..1d559ef 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,9 +22,9 @@ pub struct RelayConfig { } pub trait HandlerCallbacks { - fn ds_b_callback(&self, _in_data: &mut Vec){} + fn ds_b_callback(&self, _in_data: Vec) -> CallbackRet {CallbackRet::Relay(_in_data)} fn ds_nb_callback(&self, _in_data: Vec){} - fn us_b_callback(&self, _in_data: &mut Vec){} + fn us_b_callback(&self, _in_data: Vec) -> CallbackRet {CallbackRet::Relay(_in_data)} fn us_nb_callback(&self, _in_data: Vec){} } @@ -35,6 +35,14 @@ pub enum ConfigType { Default, } +#[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) +} + #[derive(Clone)] pub struct SSLRelay where