pierre/tcp-over-http
Archived
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases 2 Wiki Activity

cargo fmt

Browse Source
This commit is contained in:
Pierre HUBERT
2022-08-31 14:36:07 +02:00
parent cd0f6fea94
commit 3cbbd72a14
10 changed files with 170 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
tcp_relay_client/src/client_config.rs
View File

@ -42,17 +42,20 @@ pub struct ClientConfig {
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._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_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!"));
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
@ -69,18 +72,19 @@ impl ClientConfig {
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 }
} 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
})
self.get_client_keypair().map(|(c, k)| {
let mut out = k.to_vec();
out.put_slice("\n".as_bytes());
out.put_slice(c);
out
})
}
}
@ -93,4 +97,4 @@ mod test {
use clap::CommandFactory;
ClientConfig::command().debug_assert()
}
}
}

2
tcp_relay_client/src/lib.rs
View File

@ -1,2 +1,2 @@
pub mod client_config;
pub mod relay_client;
pub mod relay_client;

35
tcp_relay_client/src/main.rs
View File

@ -24,20 +24,25 @@ async fn get_server_config(config: &ClientConfig) -> Result<RemoteConfig, Box<dy
// Specify client certificate, if any
if let Some(kp) = config.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 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 {}", config.get_auth_token()))
let req = client
.get(url)
.header(
"Authorization",
format!("Bearer {}", config.get_auth_token()),
)
.send()
.await?;
if req.status().as_u16() != 200 {
log::error!("Could not retrieve configuration! (got status {})", req.status());
log::error!(
"Could not retrieve configuration! (got status {})",
req.status()
);
std::process::exit(2);
}
@ -54,7 +59,9 @@ async fn main() -> Result<(), Box<dyn Error>> {
// 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!");
log::error!(
"If you specify one of TLS certificate / key, you must then specify the other!"
);
panic!();
}
@ -71,9 +78,13 @@ async fn main() -> Result<(), Box<dyn Error>> {
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"),
format!(
"{}/ws?id={}&token={}",
args.relay_url,
port.id,
urlencoding::encode(args.get_auth_token())
)
.replace("http", "ws"),
listen_address,
args.clone(),
));
@ -83,4 +94,4 @@ async fn main() -> Result<(), Box<dyn Error>> {
join_all(handles).await;
Ok(())
}
}

45
tcp_relay_client/src/relay_client.rs
View File

@ -4,7 +4,7 @@ use std::sync::Arc;
use futures::{SinkExt, StreamExt};
use hyper_rustls::ConfigBuilderExt;
use rustls::{Certificate, PrivateKey, RootCertStore};
use rustls_pemfile::{Item, read_one};
use rustls_pemfile::{read_one, Item};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio_tungstenite::tungstenite::Message;
@ -22,7 +22,9 @@ pub async fn relay_client(ws_url: String, listen_address: String, config: Arc<Cl
};
loop {
let (socket, _) = listener.accept().await
let (socket, _) = listener
.accept()
.await
.expect("Failed to accept new connection!");
tokio::spawn(relay_connection(ws_url.clone(), socket, config.clone()));
@ -37,8 +39,7 @@ async fn relay_connection(ws_url: String, socket: TcpStream, conf: Arc<ClientCon
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 = rustls::ClientConfig::builder().with_safe_defaults();
let config = match conf.get_root_certificate() {
None => config.with_native_roots(),
@ -65,7 +66,8 @@ async fn relay_connection(ws_url: String, socket: TcpStream, conf: Arc<ClientCon
.collect::<Vec<_>>();
let key = match read_one(&mut Cursor::new(key))
.expect("Failed to read client private key!") {
.expect("Failed to read client private key!")
{
None => {
log::error!("Failed to extract private key!");
panic!();
@ -78,30 +80,29 @@ async fn relay_connection(ws_url: String, socket: TcpStream, conf: Arc<ClientCon
}
};
config.with_single_cert(certs, PrivateKey(key))
config
.with_single_cert(certs, PrivateKey(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!");
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!");
.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();
let (mut ws_write, mut ws_read) = ws_stream.split();
// TCP read -> WS write
let future = async move {
@ -136,12 +137,18 @@ async fn relay_connection(ws_url: String, socket: TcpStream, conf: Arc<ClientCon
while let Some(m) = ws_read.next().await {
match m {
Err(e) => {
log::error!("Failed to read from WebSocket. Breaking read loop... {:?}", 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);
log::error!(
"Failed to forward message to websocket. Closing reading end... {:?}",
e
);
break;
};
}
@ -149,7 +156,7 @@ async fn relay_connection(ws_url: String, socket: TcpStream, conf: Arc<ClientCon
log::info!("Server asked to close this WebSocket connection");
break;
}
Ok(m) => log::info!("{:?}", m)
Ok(m) => log::info!("{:?}", m),
}
}
}
}