Handled Cert/Key File errors. Changed data stream R/W to use match

This commit is contained in:
PinkP4nther 2021-08-23 00:56:32 -07:00
commit b25ae45948
2 changed files with 68 additions and 53 deletions

View file

@ -31,72 +31,78 @@ impl DataHandler {
pub fn get_data_stream(&mut self, data: &mut Vec<u8>) -> usize {
let mut data_length: usize = 0;
if self.stream_direction == StreamDirection::DownStream {
loop {
let mut r_buf = [0; 1024];
match self.stream_direction {
match self.tcp_stream.as_mut().unwrap().read(&mut r_buf) {
StreamDirection::DownStream => {
Ok(bytes_read) => {
loop {
if bytes_read == 0 {
break;
let mut r_buf = [0; 1024];
} else if bytes_read != 0 && bytes_read <= 1024 {
match self.tcp_stream.as_mut().unwrap().read(&mut r_buf) {
let mut tmp_buf = r_buf.to_vec();
tmp_buf.truncate(bytes_read);
Ok(bytes_read) => {
let _bw = data.write(&tmp_buf).unwrap();
data_length += bytes_read;
} else {
println!("[+] Else hit!!!!!!!!!!!!!!!!!!!!!!");
}
},
Err(e) => {
match e.kind() {
io::ErrorKind::WouldBlock => {
if bytes_read == 0 {
break;
},
_ => {println!("[!!!] Got error: {}",e);}
}
},
} else if bytes_read != 0 && bytes_read <= 1024 {
let mut tmp_buf = r_buf.to_vec();
tmp_buf.truncate(bytes_read);
let _bw = data.write(&tmp_buf).unwrap();
data_length += bytes_read;
} else {
println!("[+] Else hit!!!!!!!!!!!!!!!!!!!!!!");
}
},
Err(e) => {
match e.kind() {
io::ErrorKind::WouldBlock => {
break;
},
_ => {println!("[!!!] Got error: {}",e);}
}
},
}
}
}
} else if self.stream_direction == StreamDirection::Upstream {
loop {
},
StreamDirection::Upstream => {
loop {
let mut r_buf = [0; 1024];
let mut r_buf = [0; 1024];
match self.relay_stream.as_mut().unwrap().read(&mut r_buf) {
match self.relay_stream.as_mut().unwrap().read(&mut r_buf) {
Ok(bytes_read) => {
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 = data.write(&tmp_buf).unwrap();
data_length += bytes_read;
} else {
println!("[+] Else hit!!!!!!!!!!!!!!!!!!!!!!");
}
},
Err(e) => {
match e.kind() {
io::ErrorKind::WouldBlock => {
if bytes_read == 0 {
break;
},
_ => {println!("[!!!] Got error: {}",e);}
}
},
} else if bytes_read != 0 && bytes_read <= 1024 {
let mut tmp_buf = r_buf.to_vec();
tmp_buf.truncate(bytes_read);
let _bw = data.write(&tmp_buf).unwrap();
data_length += bytes_read;
} else {
println!("[+] Else hit!!!!!!!!!!!!!!!!!!!!!!");
}
},
Err(e) => {
match e.kind() {
io::ErrorKind::WouldBlock => {
break;
},
_ => {println!("[!!!] Got error: {}",e);}
}
},
}
}
}
}

View file

@ -2,9 +2,10 @@ use openssl::ssl::{SslAcceptor, SslFiletype, SslMethod, SslStream};
use std::io::Write;
use std::net::{TcpListener, TcpStream};
use std::sync::{Arc, Mutex};
use std::thread;
use std::{process, thread};
use std::env;
use std::fs;
use std::path::Path;
use toml::Value as TValue;
@ -143,6 +144,14 @@ impl<H: HandlerCallbacks + std::marker::Sync + std::marker::Send + 'static> SSLR
}
fn setup_ssl_config(&self, priv_key: String, cert: String) -> Arc<SslAcceptor> {
if !Path::new(priv_key.as_str()).exists() {
println!("[-] [{}] does not exist!", priv_key);
process::exit(-1);
} else if !Path::new(cert.as_str()).exists() {
println!("[-] [{}] does not exist!", cert);
process::exit(-1);
}
let mut acceptor = SslAcceptor::mozilla_intermediate(SslMethod::tls()).unwrap();
acceptor.set_private_key_file(priv_key, SslFiletype::PEM).unwrap();