147 lines
4.8 KiB
Rust
147 lines
4.8 KiB
Rust
use std::sync::Arc;
|
|
|
|
use futures::{SinkExt, StreamExt};
|
|
use hyper_rustls::ConfigBuilderExt;
|
|
use rustls::RootCertStore;
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
use tokio::net::{TcpListener, TcpStream};
|
|
use tokio_tungstenite::tungstenite::Message;
|
|
|
|
use crate::base::cert_utils;
|
|
|
|
use crate::tcp_relay_client::client_config::ClientConfig;
|
|
|
|
pub async fn relay_client(ws_url: String, listen_address: String, config: Arc<ClientConfig>) {
|
|
log::info!("Start to listen on {}", listen_address);
|
|
let listener = match TcpListener::bind(&listen_address).await {
|
|
Ok(l) => l,
|
|
Err(e) => {
|
|
log::error!("Failed to start to listen on {}! {}", listen_address, e);
|
|
std::process::exit(3);
|
|
}
|
|
};
|
|
|
|
loop {
|
|
let (socket, _) = listener
|
|
.accept()
|
|
.await
|
|
.expect("Failed to accept new connection!");
|
|
|
|
tokio::spawn(relay_connection(ws_url.clone(), socket, config.clone()));
|
|
}
|
|
}
|
|
|
|
/// Relay connection
|
|
///
|
|
/// WS read => TCP write
|
|
/// TCP read => WS write
|
|
async fn relay_connection(ws_url: String, socket: TcpStream, conf: Arc<ClientConfig>) {
|
|
log::debug!("Connecting to {}...", ws_url);
|
|
|
|
let ws_stream = if ws_url.starts_with("wss") {
|
|
let config = rustls::ClientConfig::builder().with_safe_defaults();
|
|
|
|
let config = match conf.get_root_certificate() {
|
|
None => config.with_native_roots(),
|
|
Some(cert) => {
|
|
log::debug!("Using custom root certificates");
|
|
let mut store = RootCertStore::empty();
|
|
cert_utils::parse_pem_certificates(&cert)
|
|
.unwrap()
|
|
.iter()
|
|
.for_each(|c| store.add(c).expect("Failed to add certificate to chain!"));
|
|
|
|
config.with_root_certificates(store)
|
|
}
|
|
};
|
|
|
|
let config = match conf.get_client_keypair() {
|
|
None => config.with_no_client_auth(),
|
|
Some((certs, key)) => {
|
|
let certs = cert_utils::parse_pem_certificates(certs)
|
|
.expect("Failed to parse client certificate!");
|
|
|
|
let key = cert_utils::parse_pem_private_key(key)
|
|
.expect("Failed to parse client auth private key!");
|
|
|
|
config
|
|
.with_single_cert(certs, key)
|
|
.expect("Failed to set client certificate!")
|
|
}
|
|
};
|
|
let connector = tokio_tungstenite::Connector::Rustls(Arc::new(config));
|
|
|
|
let (ws_stream, _) =
|
|
tokio_tungstenite::connect_async_tls_with_config(ws_url, None, Some(connector))
|
|
.await
|
|
.expect("Failed to connect to server relay!");
|
|
|
|
ws_stream
|
|
} else {
|
|
let (ws_stream, _) = tokio_tungstenite::connect_async(ws_url)
|
|
.await
|
|
.expect("Failed to connect to server relay!");
|
|
|
|
ws_stream
|
|
};
|
|
|
|
let (mut tcp_read, mut tcp_write) = socket.into_split();
|
|
let (mut ws_write, mut ws_read) = ws_stream.split();
|
|
|
|
// TCP read -> WS write
|
|
let future = async move {
|
|
let mut buff: [u8; 5000] = [0; 5000];
|
|
loop {
|
|
match tcp_read.read(&mut buff).await {
|
|
Ok(s) => {
|
|
if let Err(e) = ws_write.send(Message::Binary(Vec::from(&buff[0..s]))).await {
|
|
log::error!(
|
|
"Failed to write to WS connection! {:?} Exiting TCP read -> WS write loop...",e);
|
|
break;
|
|
}
|
|
|
|
if s == 0 {
|
|
log::info!("Got empty read TCP buffer. Stopping...");
|
|
break;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
log::error!(
|
|
"Failed to read from TCP connection! {:?} Exitin TCP read -> WS write loop...",
|
|
e
|
|
);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
};
|
|
tokio::spawn(future);
|
|
|
|
// WS read -> TCP write
|
|
while let Some(m) = ws_read.next().await {
|
|
match m {
|
|
Err(e) => {
|
|
log::error!(
|
|
"Failed to read from WebSocket. Breaking read loop... {:?}",
|
|
e
|
|
);
|
|
break;
|
|
}
|
|
Ok(Message::Binary(b)) => {
|
|
if let Err(e) = tcp_write.write_all(&b).await {
|
|
log::error!(
|
|
"Failed to forward message to websocket. Closing reading end... {:?}",
|
|
e
|
|
);
|
|
break;
|
|
};
|
|
}
|
|
Ok(Message::Close(_r)) => {
|
|
log::info!("Server asked to close this WebSocket connection");
|
|
break;
|
|
}
|
|
Ok(m) => log::info!("{:?}", m),
|
|
}
|
|
}
|
|
}
|