259 lines
No EOL
12 KiB
Rust
259 lines
No EOL
12 KiB
Rust
use crate::{
|
|
FullDuplexTcp,
|
|
HandlerCallbacks,
|
|
DataStreamType,
|
|
TCPDataType,
|
|
Duration,
|
|
Arc,
|
|
Mutex,
|
|
DownStreamInner,
|
|
UpStreamInner,
|
|
InnerHandlers,
|
|
Shutdown,
|
|
Sender,
|
|
Receiver,
|
|
FullDuplexTcpState,
|
|
DataPipe,
|
|
mpsc,
|
|
thread,
|
|
CallbackRet,
|
|
TcpStream,
|
|
SslVerifyMode,
|
|
SslConnector,
|
|
SslMethod,
|
|
};
|
|
|
|
impl<H: HandlerCallbacks + std::marker::Sync + std::marker::Send + Clone + 'static> FullDuplexTcp<H> {
|
|
|
|
pub fn new(ds_tcp_stream: DataStreamType, us_tcp_stream_type: TCPDataType, remote_host: String, remote_port: String, handlers: InnerHandlers<H>) -> Result<Self, i8> {
|
|
|
|
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))); },
|
|
}
|
|
|
|
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);
|
|
}
|
|
};
|
|
|
|
Ok(
|
|
FullDuplexTcp {
|
|
remote_host,
|
|
remote_port,
|
|
ds_inner_m: Arc::new(Mutex::new(Some(DownStreamInner{ds_stream: ds_tcp_stream, internal_data_buffer: Vec::<u8>::new()}))),
|
|
us_inner_m: Arc::new(Mutex::new(Some(UpStreamInner{us_stream: us_tcp_stream, internal_data_buffer: Vec::<u8>::new()}))),
|
|
inner_handlers: handlers,
|
|
})
|
|
}
|
|
|
|
pub fn handle(&mut self) {
|
|
|
|
let (state_sender, state_receiver): (Sender<FullDuplexTcpState>, Receiver<FullDuplexTcpState>) = mpsc::channel();
|
|
let (ds_data_pipe_sender, ds_data_pipe_receiver): (Sender<DataPipe>, Receiver<DataPipe>) = mpsc::channel();
|
|
let (us_data_pipe_sender, us_data_pipe_receiver): (Sender<DataPipe>, Receiver<DataPipe>) = mpsc::channel();
|
|
|
|
let ds_method_pointer = self.ds_inner_m.clone();
|
|
let ds_state_bc = state_sender.clone();
|
|
|
|
let us_method_pointer = self.us_inner_m.clone();
|
|
let us_state_bc = state_sender.clone();
|
|
|
|
thread::spawn(move || {
|
|
ds_method_pointer.lock().unwrap().take().unwrap().ds_handler(ds_state_bc, ds_data_pipe_receiver);
|
|
});
|
|
|
|
thread::spawn(move || {
|
|
us_method_pointer.lock().unwrap().take().unwrap().us_handler(us_state_bc, us_data_pipe_receiver);
|
|
});
|
|
|
|
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)
|
|
*/
|
|
|
|
let inner_handlers_clone = self.inner_handlers.clone();
|
|
let in_data = data.clone();
|
|
|
|
thread::spawn(move || {
|
|
inner_handlers_clone.cb.us_nb_callback(in_data);
|
|
});
|
|
|
|
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) => {
|
|
|
|
/*
|
|
Callbacks that work with data from DownStream go here
|
|
*/
|
|
|
|
let inner_handlers_clone = self.inner_handlers.clone();
|
|
let in_data = data.clone();
|
|
|
|
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<DataStreamType, i8> {
|
|
|
|
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);
|
|
}
|
|
} |