Merged all workspace projects into a single binary project
This commit is contained in:
100
src/tcp_relay_client/client_config.rs
Normal file
100
src/tcp_relay_client/client_config.rs
Normal file
@ -0,0 +1,100 @@
|
||||
use bytes::BufMut;
|
||||
use clap::Parser;
|
||||
|
||||
/// TCP relay client
|
||||
#[derive(Parser, Debug, Clone)]
|
||||
#[clap(author, version, about, long_about = None)]
|
||||
pub struct ClientConfig {
|
||||
/// Access token
|
||||
#[clap(short, long)]
|
||||
pub token: Option<String>,
|
||||
|
||||
/// Relay server
|
||||
#[clap(short, long, default_value = "http://127.0.0.1:8000")]
|
||||
pub relay_url: String,
|
||||
|
||||
/// Listen address
|
||||
#[clap(short, long, default_value = "127.0.0.1")]
|
||||
pub listen_address: String,
|
||||
|
||||
/// Alternative root certificate to use for server authentication
|
||||
#[clap(short = 'c', long)]
|
||||
pub root_certificate: Option<String>,
|
||||
|
||||
#[clap(skip)]
|
||||
_root_certificate_cache: Option<Vec<u8>>,
|
||||
|
||||
/// TLS certificate for TLS authentication.
|
||||
#[clap(long)]
|
||||
pub tls_cert: Option<String>,
|
||||
|
||||
#[clap(skip)]
|
||||
_tls_cert_cache: Option<Vec<u8>>,
|
||||
|
||||
/// TLS key for TLS authentication.
|
||||
#[clap(long)]
|
||||
pub tls_key: Option<String>,
|
||||
|
||||
#[clap(skip)]
|
||||
_tls_key_cache: Option<Vec<u8>>,
|
||||
}
|
||||
|
||||
impl ClientConfig {
|
||||
/// Load certificates and put them in cache
|
||||
pub fn load_certificates(&mut self) {
|
||||
self._root_certificate_cache = self
|
||||
.root_certificate
|
||||
.as_ref()
|
||||
.map(|c| std::fs::read(c).expect("Failed to read root certificate!"));
|
||||
|
||||
self._tls_cert_cache = self
|
||||
.tls_cert
|
||||
.as_ref()
|
||||
.map(|c| std::fs::read(c).expect("Failed to read client certificate!"));
|
||||
|
||||
self._tls_key_cache = self
|
||||
.tls_key
|
||||
.as_ref()
|
||||
.map(|c| std::fs::read(c).expect("Failed to read client key!"));
|
||||
}
|
||||
|
||||
/// Get client token, returning a dummy token if none was specified
|
||||
pub fn get_auth_token(&self) -> &str {
|
||||
self.token.as_deref().unwrap_or("none")
|
||||
}
|
||||
|
||||
/// Get root certificate content
|
||||
pub fn get_root_certificate(&self) -> Option<Vec<u8>> {
|
||||
self._root_certificate_cache.clone()
|
||||
}
|
||||
|
||||
/// Get client certificate & key pair, if available
|
||||
pub fn get_client_keypair(&self) -> Option<(&Vec<u8>, &Vec<u8>)> {
|
||||
if let (Some(cert), Some(key)) = (&self._tls_cert_cache, &self._tls_key_cache) {
|
||||
Some((cert, key))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Get client certificate & key pair, in a single memory buffer
|
||||
pub fn get_merged_client_keypair(&self) -> Option<Vec<u8>> {
|
||||
self.get_client_keypair().map(|(c, k)| {
|
||||
let mut out = k.to_vec();
|
||||
out.put_slice("\n".as_bytes());
|
||||
out.put_slice(c);
|
||||
out
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use crate::client_config::ClientConfig;
|
||||
|
||||
#[test]
|
||||
fn verify_cli() {
|
||||
use clap::CommandFactory;
|
||||
ClientConfig::command().debug_assert()
|
||||
}
|
||||
}
|
||||
100
src/tcp_relay_client/mod.rs
Normal file
100
src/tcp_relay_client/mod.rs
Normal file
@ -0,0 +1,100 @@
|
||||
extern crate core;
|
||||
|
||||
use std::error::Error;
|
||||
use std::sync::Arc;
|
||||
|
||||
use futures::future::join_all;
|
||||
use reqwest::{Certificate, Identity};
|
||||
|
||||
use crate::base::RemoteConfig;
|
||||
use crate::tcp_relay_client::client_config::ClientConfig;
|
||||
use crate::tcp_relay_client::relay_client::relay_client;
|
||||
|
||||
pub mod client_config;
|
||||
mod relay_client;
|
||||
|
||||
/// Get remote server config i.e. get the list of forwarded ports
|
||||
async fn get_server_config(conf: &ClientConfig) -> Result<RemoteConfig, Box<dyn Error>> {
|
||||
let url = format!("{}/config", conf.relay_url);
|
||||
log::info!("Retrieving configuration on {}", url);
|
||||
|
||||
let mut client = reqwest::Client::builder();
|
||||
|
||||
// Specify root certificate, if any was specified in the command line
|
||||
if let Some(cert) = conf.get_root_certificate() {
|
||||
client = client.add_root_certificate(Certificate::from_pem(&cert)?);
|
||||
}
|
||||
|
||||
// Specify client certificate, if any
|
||||
if let Some(kp) = conf.get_merged_client_keypair() {
|
||||
let identity = Identity::from_pem(&kp).expect("Failed to load certificates for reqwest!");
|
||||
client = client.identity(identity).use_rustls_tls();
|
||||
}
|
||||
|
||||
let client = client.build().expect("Failed to build reqwest client");
|
||||
|
||||
let req = client
|
||||
.get(url)
|
||||
.header("Authorization", format!("Bearer {}", conf.get_auth_token()))
|
||||
.send()
|
||||
.await?;
|
||||
if req.status().as_u16() != 200 {
|
||||
log::error!(
|
||||
"Could not retrieve configuration! (got status {})",
|
||||
req.status()
|
||||
);
|
||||
std::process::exit(2);
|
||||
}
|
||||
|
||||
Ok(req.json::<RemoteConfig>().await?)
|
||||
}
|
||||
|
||||
/// Core logic of the application
|
||||
pub async fn run_app(mut args: ClientConfig) -> std::io::Result<()> {
|
||||
args.load_certificates();
|
||||
let args = Arc::new(args);
|
||||
|
||||
// Check arguments coherence
|
||||
if args.tls_cert.is_some() != args.tls_key.is_some() {
|
||||
log::error!(
|
||||
"If you specify one of TLS certificate / key, you must then specify the other!"
|
||||
);
|
||||
panic!();
|
||||
}
|
||||
|
||||
if args.get_client_keypair().is_some() {
|
||||
log::info!("Using client-side authentication");
|
||||
}
|
||||
|
||||
// Get server relay configuration (fetch the list of port to forward)
|
||||
let remote_conf = match get_server_config(&args).await {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
log::error!("Failed to fetch relay configuration from server! {}", e);
|
||||
panic!();
|
||||
}
|
||||
};
|
||||
|
||||
// Start to listen port
|
||||
let mut handles = vec![];
|
||||
for port in remote_conf {
|
||||
let listen_address = format!("{}:{}", args.listen_address, port.port);
|
||||
|
||||
let h = tokio::spawn(relay_client(
|
||||
format!(
|
||||
"{}/ws?id={}&token={}",
|
||||
args.relay_url,
|
||||
port.id,
|
||||
urlencoding::encode(args.get_auth_token())
|
||||
)
|
||||
.replace("http", "ws"),
|
||||
listen_address,
|
||||
args.clone(),
|
||||
));
|
||||
handles.push(h);
|
||||
}
|
||||
|
||||
join_all(handles).await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
146
src/tcp_relay_client/relay_client.rs
Normal file
146
src/tcp_relay_client/relay_client.rs
Normal file
@ -0,0 +1,146 @@
|
||||
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),
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user