Can specify custom server root certificate for client

tcp_relay_client/src/relay_client.rs

@@ -1,9 +1,14 @@
+use std::io::Cursor;
+use std::sync::Arc;
+
 use futures::{SinkExt, StreamExt};
+use hyper_rustls::ConfigBuilderExt;
+use rustls::{Certificate, RootCertStore};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::{TcpListener, TcpStream};
 use tokio_tungstenite::tungstenite::Message;
 
-pub async fn relay_client(ws_url: String, listen_address: String) {
+pub async fn relay_client(ws_url: String, listen_address: String, root_cert: Option<Vec<u8>>) {
     log::info!("Start to listen on {}", listen_address);
     let listener = match TcpListener::bind(&listen_address).await {
         Ok(l) => l,
@@ -17,7 +22,7 @@ pub async fn relay_client(ws_url: String, listen_address: String) {
         let (socket, _) = listener.accept().await
             .expect("Failed to accept new connection!");
 
-        tokio::spawn(relay_connection(ws_url.clone(), socket));
+        tokio::spawn(relay_connection(ws_url.clone(), socket, root_cert.clone()));
     }
 }
 
@@ -25,10 +30,44 @@ pub async fn relay_client(ws_url: String, listen_address: String) {
 ///
 /// WS read => TCP write
 /// TCP read => WS write
-async fn relay_connection(ws_url: String, socket: TcpStream) {
+async fn relay_connection(ws_url: String, socket: TcpStream, root_cert: Option<Vec<u8>>) {
     log::debug!("Connecting to {}...", ws_url);
-    let (ws_stream, _) = tokio_tungstenite::connect_async(ws_url)
-        .await.expect("Failed to connect to server relay!");
+
+    let ws_stream = if ws_url.starts_with("wss") {
+        let config = rustls::ClientConfig::builder()
+            .with_safe_defaults();
+
+        let config = match root_cert {
+            None => config.with_native_roots(),
+            Some(cert) => {
+                log::debug!("Using custom root certificates");
+                let mut store = RootCertStore::empty();
+                rustls_pemfile::certs(&mut Cursor::new(cert))
+                    .expect("Failed to parse root certificates!")
+                    .into_iter()
+                    .map(Certificate)
+                    .for_each(|c| store.add(&c).expect("Failed to add certificate to chain!"));
+
+                config.with_root_certificates(store)
+            }
+        };
+
+        let config = config.with_no_client_auth();
+        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();