Handle client auth with CLI

Author: ameba23

src/lib.rs

@@ -1,6 +1,7 @@
 mod attestation;
 
 pub use attestation::{AttestationPlatform, MockAttestation, NoAttestation};
+use tokio_rustls::rustls::server::WebPkiClientVerifier;
 
 #[cfg(test)]
 mod test_helpers;
@@ -18,6 +19,11 @@ use tokio_rustls::{
 /// The label used when exporting key material from a TLS session
 const EXPORTER_LABEL: &[u8; 24] = b"EXPORTER-Channel-Binding";
 
+pub struct TlsCertAndKey {
+    pub cert_chain: Vec<CertificateDer<'static>>,
+    pub key: PrivateKeyDer<'static>,
+}
+
 struct Proxy<L, R>
 where
     L: AttestationPlatform,
@@ -48,20 +54,36 @@ where
 
 impl<L: AttestationPlatform, R: AttestationPlatform> ProxyServer<L, R> {
     pub async fn new(
-        cert_chain: Vec<CertificateDer<'static>>,
-        key: PrivateKeyDer<'static>,
+        cert_and_key: TlsCertAndKey,
         local: impl ToSocketAddrs,
         target: SocketAddr,
         local_attestation_platform: L,
         remote_attestation_platform: R,
+        client_auth: bool,
     ) -> Self {
-        let server_config = ServerConfig::builder()
-            .with_no_client_auth()
-            .with_single_cert(cert_chain.clone(), key)
-            .expect("Failed to create rustls server config");
+        if remote_attestation_platform.is_cvm() && !client_auth {
+            panic!("Client auth is required when the client is running in a CVM");
+        }
+
+        let server_config = if client_auth {
+            let root_store =
+                RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
+            let verifier = WebPkiClientVerifier::builder(Arc::new(root_store))
+                .build()
+                .expect("invalid client verifier");
+            ServerConfig::builder()
+                .with_client_cert_verifier(verifier)
+                .with_single_cert(cert_and_key.cert_chain.clone(), cert_and_key.key)
+                .expect("Failed to create rustls server config")
+        } else {
+            ServerConfig::builder()
+                .with_no_client_auth()
+                .with_single_cert(cert_and_key.cert_chain.clone(), cert_and_key.key)
+                .expect("Failed to create rustls server config")
+        };
 
         Self::new_with_tls_config(
-            cert_chain,
+            cert_and_key.cert_chain,
             server_config.into(),
             local,
             target,
@@ -187,17 +209,28 @@ where
 
 impl<L: AttestationPlatform, R: AttestationPlatform> ProxyClient<L, R> {
     pub async fn new(
+        cert_and_key: Option<TlsCertAndKey>,
         address: impl ToSocketAddrs,
         server_address: SocketAddr,
         server_name: ServerName<'static>,
         local_attestation_platform: L,
         remote_attestation_platform: R,
-        cert_chain: Option<Vec<CertificateDer<'static>>>,
     ) -> Self {
         let root_store = RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
-        let client_config = ClientConfig::builder()
-            .with_root_certificates(root_store)
-            .with_no_client_auth();
+
+        let client_config = if let Some(ref cert_and_key) = cert_and_key {
+            ClientConfig::builder()
+                .with_root_certificates(root_store)
+                .with_client_auth_cert(
+                    cert_and_key.cert_chain.clone(),
+                    cert_and_key.key.clone_key(),
+                )
+                .unwrap()
+        } else {
+            ClientConfig::builder()
+                .with_root_certificates(root_store)
+                .with_no_client_auth()
+        };
 
         Self::new_with_tls_config(
             client_config.into(),
@@ -206,7 +239,7 @@ impl<L: AttestationPlatform, R: AttestationPlatform> ProxyClient<L, R> {
             server_name,
             local_attestation_platform,
             remote_attestation_platform,
-            cert_chain,
+            cert_and_key.map(|c| c.cert_chain),
         )
         .await
     }

src/main.rs

@@ -1,8 +1,8 @@
 use clap::{Parser, Subcommand};
-use std::{fs::File, net::SocketAddr};
+use std::{fs::File, net::SocketAddr, path::PathBuf};
 use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer};
 
-use attested_tls_proxy::{MockAttestation, NoAttestation, ProxyClient, ProxyServer};
+use attested_tls_proxy::{MockAttestation, NoAttestation, ProxyClient, ProxyServer, TlsCertAndKey};
 
 #[derive(Parser, Debug, Clone)]
 #[clap(version, about, long_about = None)]
@@ -22,11 +22,26 @@ enum CliCommand {
         server_address: SocketAddr,
         #[arg(long)]
         server_name: String,
+        /// The path to a PEM encoded private key for client authentication
+        #[arg(long)]
+        private_key: Option<PathBuf>,
+        /// The path to a PEM encoded certificate chain for client authentication
+        #[arg(long)]
+        cert_chain: Option<PathBuf>,
     },
     /// Run a proxy server
     Server {
+        /// Socket address of the target service to forward traffic to
         #[arg(short, long)]
-        client_address: SocketAddr,
+        target_address: SocketAddr,
+        /// The path to a PEM encoded private key
+        #[arg(long)]
+        private_key: PathBuf,
+        /// The path to a PEM encoded certificate chain
+        #[arg(long)]
+        cert_chain: PathBuf,
+        #[arg(long)]
+        client_auth: bool,
     },
 }
 
@@ -38,34 +53,43 @@ async fn main() {
         CliCommand::Client {
             server_name,
             server_address,
+            private_key,
+            cert_chain,
         } => {
+            let tls_cert_and_chain = private_key
+                .map(|private_key| load_tls_cert_and_key(cert_chain.unwrap(), private_key));
+
             let client = ProxyClient::new(
+                tls_cert_and_chain,
                 cli.address,
                 server_address,
                 server_name.try_into().unwrap(),
                 NoAttestation,
                 MockAttestation,
-                None,
             )
             .await;
 
             loop {
                 client.accept().await.unwrap();
             }
         }
-        CliCommand::Server { client_address } => {
-            let cert_chain = load_certs_pem("certs.pem").unwrap();
-            let key = load_private_key_pem("key.pem");
+        CliCommand::Server {
+            target_address,
+            private_key,
+            cert_chain,
+            client_auth,
+        } => {
+            let tls_cert_and_chain = load_tls_cert_and_key(cert_chain, private_key);
             let local_attestation = MockAttestation;
             let remote_attestation = NoAttestation;
 
             let server = ProxyServer::new(
-                cert_chain,
-                key,
+                tls_cert_and_chain,
                 cli.address,
-                client_address,
+                target_address,
                 local_attestation,
                 remote_attestation,
+                client_auth,
             )
             .await;
 
@@ -76,15 +100,21 @@ async fn main() {
     }
 }
 
-pub fn load_certs_pem(path: &str) -> std::io::Result<Vec<CertificateDer<'static>>> {
+fn load_tls_cert_and_key(cert_chain: PathBuf, private_key: PathBuf) -> TlsCertAndKey {
+    let key = load_private_key_pem(private_key);
+    let cert_chain = load_certs_pem(cert_chain).unwrap();
+    TlsCertAndKey { key, cert_chain }
+}
+
+pub fn load_certs_pem(path: PathBuf) -> std::io::Result<Vec<CertificateDer<'static>>> {
     Ok(
         rustls_pemfile::certs(&mut std::io::BufReader::new(File::open(path)?))
             .map(|res| res.unwrap())
             .collect(),
     )
 }
 
-pub fn load_private_key_pem(path: &str) -> PrivateKeyDer<'static> {
+pub fn load_private_key_pem(path: PathBuf) -> PrivateKeyDer<'static> {
     let mut reader = std::io::BufReader::new(File::open(path).unwrap());
 
     // Tries to read the key as PKCS#8, PKCS#1, or SEC1