wip

Author: andrewjstone

tls/examples/server.rs

@@ -88,7 +88,7 @@ async fn main() {
         resolve,
     };
 
-    let mut server = Server::new(server_config, listen_addr, log.clone())
+    let server = Server::new(server_config, listen_addr, log.clone())
         .await
         .unwrap();
 

tls/src/lib.rs

@@ -320,6 +320,8 @@ mod tests {
     use slog::Drain;
     use std::net::SocketAddrV6;
     use std::str::FromStr;
+    use std::sync::atomic::AtomicUsize;
+    use std::sync::Arc;
     use std::time::Duration;
     use tokio::io::{AsyncReadExt, AsyncWriteExt};
     use tokio::time::sleep;
@@ -400,7 +402,7 @@ mod tests {
 
         tokio::spawn(async move {
             let server_config = local_config(1);
-            let mut server = Server::new(server_config, addr, log2.clone())
+            let server = Server::new(server_config, addr, log2.clone())
                 .await
                 .unwrap();
 
@@ -460,7 +462,7 @@ mod tests {
         ];
 
         tokio::spawn(async move {
-            let mut server = Server::new(server_config, addr, log2.clone())
+            let server = Server::new(server_config, addr, log2.clone())
                 .await
                 .unwrap();
 
@@ -522,7 +524,7 @@ mod tests {
         ];
 
         let handle = tokio::spawn(async move {
-            let mut server = Server::new(server_config, addr, log2.clone())
+            let server = Server::new(server_config, addr, log2.clone())
                 .await
                 .unwrap();
 
@@ -578,57 +580,83 @@ mod tests {
 
         let server_config = local_config(1);
 
+        let max_connections = 3;
+        let done_count = Arc::new(AtomicUsize::new(0));
+
         // Message to send over TLS
         const MSG: &str = "Hello Joe";
 
-        let (done_tx, done_rx) = tokio::sync::oneshot::channel();
         let log2 = log.clone();
         let corpus = vec![
             pki_keydir.join("corim-rot.cbor"),
             pki_keydir.join("corim-sp.cbor"),
         ];
 
+        // Accept connections from 3 clients in different tasks
+        //
+        // For this test, the clients all share a set of keys, because
+        // we only generate 2 sets of keys from the KDL.
+        let dc2 = done_count.clone();
         tokio::spawn(async move {
-            let mut server = Server::new(server_config, addr, log2.clone())
+            let server = Server::new(server_config, addr, log2.clone())
                 .await
                 .unwrap();
 
-            let (mut stream, _) =
-                server.accept(corpus.clone()).await.await.unwrap();
-            let mut buf = String::new();
-            stream.read_to_string(&mut buf).await.unwrap();
-
-            assert_eq!(buf.as_str(), MSG);
-
-            // Inform the main task that the test is complete.
-            let _ = done_tx.send(());
+            for _ in 0..max_connections {
+                let accept_fut = server.accept(corpus.clone()).await;
+                let done_count = dc2.clone();
+                tokio::spawn(async move {
+                    let (mut stream, _) = accept_fut.await.unwrap();
+                    let mut buf = String::new();
+                    stream.read_to_string(&mut buf).await.unwrap();
+
+                    assert_eq!(buf.as_str(), MSG);
+                    done_count
+                        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+                });
+            }
         });
 
         // Loop until we succesfully connect
-        let mut stream = loop {
-            let client_config = local_config(2);
-
-            let corpus = vec![
-                pki_keydir.join("corim-rot.cbor"),
-                pki_keydir.join("corim-sp.cbor"),
-            ];
-
-            if let Ok(stream) =
-                Client::connect(client_config, addr, corpus, log.clone()).await
-            {
-                break stream;
-            }
-            sleep(Duration::from_millis(1)).await;
-        };
-
-        stream.write_all(MSG.as_bytes()).await.unwrap();
-
-        // Trigger an EOF so that `read_to_string` in the acceptor task
-        // completes.
-        stream.shutdown().await.unwrap();
+        for _ in 0..max_connections {
+            let pki_keydir = pki_keydir.clone();
+            let log = log.clone();
+            tokio::spawn(async move {
+                let mut stream = loop {
+                    let client_config = local_config(2);
+
+                    let corpus = vec![
+                        pki_keydir.join("corim-rot.cbor"),
+                        pki_keydir.join("corim-sp.cbor"),
+                    ];
+
+                    if let Ok(stream) = Client::connect(
+                        client_config,
+                        addr,
+                        corpus,
+                        log.clone(),
+                    )
+                    .await
+                    {
+                        break stream;
+                    }
+                    sleep(Duration::from_millis(1)).await;
+                };
+
+                stream.write_all(MSG.as_bytes()).await.unwrap();
+
+                // Trigger an EOF so that `read_to_string` in the acceptor task
+                // completes.
+                stream.shutdown().await.unwrap();
+            });
+        }
 
         // Wait for the other side of the connection to receive and assert the
         // message
-        let _ = done_rx.await;
+        while done_count.load(std::sync::atomic::Ordering::Relaxed)
+            != max_connections
+        {
+            sleep(Duration::from_millis(1)).await;
+        }
     }
 }

tls/src/server.rs

@@ -250,17 +250,25 @@ impl Server {
     }
 
     pub async fn accept(
-        &mut self,
+        &self,
         corpus: Vec<Utf8PathBuf>,
     ) -> impl Future<
         Output = Result<(Stream<TcpStream>, core::net::SocketAddr), Error>,
-    > + use<'_> {
+    > {
         // This is where the actual socket accept occurs.
         //
         // Everything else takes a while and therefore we return a future that
         // can be awaited inside a spawned taks.
         let accept_res = self.tcp_listener.accept().await;
-        async {
+
+        // Clone all our fields so that we don't hold a reference to `self`
+        // This allows the future to be `Send` so it can be spawned into a task.
+        let tls_acceptor = self.tls_acceptor.clone();
+        let log = self.log.clone();
+        let attest_config = self.attest_config.clone();
+        let roots = self.roots.clone();
+
+        async move {
             let (stream, addr) = accept_res?;
             // load corims into a set of ReferenceMeasurements
             let mut corims = Vec::new();
@@ -269,7 +277,7 @@ impl Server {
             }
             let corpus = ReferenceMeasurements::try_from(corims.as_slice())?;
 
-            let mut stream = self.tls_acceptor.clone().accept(stream).await?;
+            let mut stream = tls_acceptor.clone().accept(stream).await?;
 
             // get PlatformId from server TLS / Trust Quorum cert chain
             let (_, conn) = stream.get_ref();
@@ -333,7 +341,7 @@ impl Server {
             }
 
             // Right now all protocols are the same
-            info!(self.log, "Running with protocol version {version}");
+            info!(log, "Running with protocol version {version}");
 
             // get Nonce from client
             let client_nonce = recv_msg(&mut stream).await?;
@@ -347,8 +355,7 @@ impl Server {
             // The attesation protocol has an inherent race condition between
             // getting the log and the attestation. We verify our own attestation
             // before sending it to the challenger to fail as early as possible.
-            let attest_data =
-                get_attest_data(&self.attest_config, &client_nonce)?;
+            let attest_data = get_attest_data(&attest_config, &client_nonce)?;
             dice_verifier::verify_attestation(
                 &attest_data.certs[0],
                 &attest_data.attestation,
@@ -361,12 +368,12 @@ impl Server {
             let client_cert_chain = certs_from_der(&client_cert_chain)?;
             let root = dice_verifier::verify_cert_chain(
                 &client_cert_chain,
-                Some(&self.roots),
+                Some(&roots),
             )?;
             let client_platform_id =
                 dice_mfg_msgs::PlatformId::try_from(&client_cert_chain)?;
             info!(
-                self.log,
+                log,
                 "Cert chain from peer \"{}\" verified against root \"{}\"",
                 client_platform_id.as_str()?,
                 root.tbs_certificate.subject,
@@ -375,10 +382,7 @@ impl Server {
             if tq_platform_id != client_platform_id {
                 return Err(Error::PlatformIdMismatch);
             }
-            info!(
-                self.log,
-                "TQ & attestation cert chains agree on platform id"
-            );
+            info!(log, "TQ & attestation cert chains agree on platform id");
 
             // send server attestation cert chain to client
             let cert_chain_der = certs_to_der(&attest_data.certs)?;
@@ -405,7 +409,7 @@ impl Server {
                 &client_log,
                 &nonce,
             )?;
-            info!(self.log, "Peer attestation verified");
+            info!(log, "Peer attestation verified");
 
             // appraise measurements from client attestation against reference
             // measurements
@@ -418,14 +422,11 @@ impl Server {
                 &corpus,
             ) {
                 Ok(()) => {
-                    info!(self.log, "Peer measurements appraised successfully");
+                    info!(log, "Peer measurements appraised successfully");
                     true
                 }
                 Err(e) => {
-                    info!(
-                        self.log,
-                        "Peer measurements appraisal failed: {}", e
-                    );
+                    info!(log, "Peer measurements appraisal failed: {}", e);
                     false
                 }
             };