Return a future from accept after the TCP accept completes

Author: andrewjstone

tls/examples/server.rs

@@ -93,7 +93,8 @@ async fn main() {
         .unwrap();
 
     loop {
-        let (stream, _) = server.accept(args.corpus.as_slice()).await.unwrap();
+        let (stream, _) =
+            server.accept(args.corpus.clone()).await.await.unwrap();
         let platform_id = stream.peer_platform_id().as_str().unwrap();
         info!(log, "connected to attested peer: {platform_id}");
         let (mut reader, mut writer) = split(stream);

tls/src/lib.rs

@@ -388,7 +388,7 @@ mod tests {
                 .unwrap();
 
             let (mut stream, _) =
-                server.accept(corpus.as_slice()).await.unwrap();
+                server.accept(corpus.clone()).await.await.unwrap();
             let mut buf = String::new();
             stream.read_to_string(&mut buf).await.unwrap();
 
@@ -475,7 +475,7 @@ mod tests {
                 .unwrap();
 
             let (mut stream, _) =
-                server.accept(corpus.as_slice()).await.unwrap();
+                server.accept(corpus.clone()).await.await.unwrap();
             let mut buf = String::new();
             stream.read_to_string(&mut buf).await.unwrap();
 
@@ -563,7 +563,7 @@ mod tests {
                 .unwrap();
 
             // We never expect this to succeed
-            let _ = match server.accept(corpus.as_slice()).await {
+            let _ = match server.accept(corpus.clone()).await.await {
                 Ok(_) => panic!("This should not succed"),
                 Err(_) => done_tx.send(()),
             };

tls/src/server.rs

@@ -28,10 +28,11 @@ use rustls::{
     CipherSuite, ServerConfig, SignatureScheme,
 };
 use slog::{error, info};
+use std::future::Future;
 use std::net::SocketAddrV6;
 use std::sync::Arc;
 use tokio::net::{TcpListener, TcpStream};
-use tokio_rustls::TlsAcceptor;
+use tokio_rustls::{Accept as TlsAccept, TlsAcceptor};
 use x509_cert::{
     der::{Decode, DecodePem},
     Certificate,
@@ -250,157 +251,172 @@ impl Server {
 
     pub async fn accept(
         &mut self,
-        corpus: &[Utf8PathBuf],
-    ) -> Result<(Stream<TcpStream>, core::net::SocketAddr), Error> {
-        // load corims into a set of ReferenceMeasurements
-        let mut corims = Vec::new();
-        for c in corpus {
-            corims.push(Corim::from_file(c)?);
-        }
-        let corpus = ReferenceMeasurements::try_from(corims.as_slice())?;
-
-        let (stream, addr) = self.tcp_listener.accept().await?;
-        let mut stream = self.tls_acceptor.clone().accept(stream).await?;
-
-        // get PlatformId from server TLS / Trust Quorum cert chain
-        let (_, conn) = stream.get_ref();
-        let tq_platform_id = if let Some(tls_certs) = conn.peer_certificates() {
-            let mut pki_path = Vec::new();
-            for der in tls_certs.iter() {
-                pki_path.push(Certificate::from_der(der).map_err(|_| {
-                    rustls::Error::InvalidCertificate(
-                        rustls::CertificateError::BadEncoding,
-                    )
-                })?)
+        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 {
+            let (stream, addr) = accept_res?;
+            // load corims into a set of ReferenceMeasurements
+            let mut corims = Vec::new();
+            for c in corpus {
+                corims.push(Corim::from_file(c)?);
             }
-            dice_mfg_msgs::PlatformId::try_from(&pki_path)?
-        } else {
-            return Err(Error::NoTQCerts);
-        };
-
-        // get version from the client
-        let version_bytes = recv_msg(&mut stream).await?;
-        let version =
-            u32::from_le_bytes(version_bytes[..4].try_into().unwrap());
-
-        if version == CURRENT_PROTOCOL_VERSION {
-            // we're good to go
-            let mut buf = vec![0u8; ProtocolResult::MAX_SIZE];
-            let resp: ProtocolResult = Ok(version);
-            let resp_len = hubpack::serialize(&mut buf, &resp)?;
-            send_msg(&mut stream, &buf[..resp_len]).await?;
-        } else if version == PREVIOUS_PROTOCOL_VERSION {
-            // We eventually want to support older protocol
-            let mut buf = vec![0u8; ProtocolResult::MAX_SIZE];
-            let resp: ProtocolResult = Ok(version);
-            let resp_len = hubpack::serialize(&mut buf, &resp)?;
-            send_msg(&mut stream, &buf[..resp_len]).await?;
-        } else {
-            // We can't deal with this
-            // We eventually want to support older protocol
-            let mut buf = vec![0u8; ProtocolResult::MAX_SIZE];
-            let resp: ProtocolResult = Err(());
-            let resp_len = hubpack::serialize(&mut buf, &resp)?;
-            send_msg(&mut stream, &buf[..resp_len]).await?;
-            // Client has given us something bad, time to give up
-            return Err(Error::ProtocolVersion);
-        }
+            let corpus = ReferenceMeasurements::try_from(corims.as_slice())?;
+
+            let mut stream = self.tls_acceptor.clone().accept(stream).await?;
+
+            // get PlatformId from server TLS / Trust Quorum cert chain
+            let (_, conn) = stream.get_ref();
+            let tq_platform_id = if let Some(tls_certs) =
+                conn.peer_certificates()
+            {
+                let mut pki_path = Vec::new();
+                for der in tls_certs.iter() {
+                    pki_path.push(Certificate::from_der(der).map_err(|_| {
+                        rustls::Error::InvalidCertificate(
+                            rustls::CertificateError::BadEncoding,
+                        )
+                    })?)
+                }
+                dice_mfg_msgs::PlatformId::try_from(&pki_path)?
+            } else {
+                return Err(Error::NoTQCerts);
+            };
 
-        // Wait for the protocol ACK
-        let protocol_ack_bytes = recv_msg(&mut stream).await?;
-        let (protocol_ack, _): (ProtocolRequestAck, _) =
-            hubpack::deserialize(&protocol_ack_bytes)?;
+            // get version from the client
+            let version_bytes = recv_msg(&mut stream).await?;
+            let version =
+                u32::from_le_bytes(version_bytes[..4].try_into().unwrap());
+
+            if version == CURRENT_PROTOCOL_VERSION {
+                // we're good to go
+                let mut buf = vec![0u8; ProtocolResult::MAX_SIZE];
+                let resp: ProtocolResult = Ok(version);
+                let resp_len = hubpack::serialize(&mut buf, &resp)?;
+                send_msg(&mut stream, &buf[..resp_len]).await?;
+            } else if version == PREVIOUS_PROTOCOL_VERSION {
+                // We eventually want to support older protocol
+                let mut buf = vec![0u8; ProtocolResult::MAX_SIZE];
+                let resp: ProtocolResult = Ok(version);
+                let resp_len = hubpack::serialize(&mut buf, &resp)?;
+                send_msg(&mut stream, &buf[..resp_len]).await?;
+            } else {
+                // We can't deal with this
+                // We eventually want to support older protocol
+                let mut buf = vec![0u8; ProtocolResult::MAX_SIZE];
+                let resp: ProtocolResult = Err(());
+                let resp_len = hubpack::serialize(&mut buf, &resp)?;
+                send_msg(&mut stream, &buf[..resp_len]).await?;
+                // Client has given us something bad, time to give up
+                return Err(Error::ProtocolVersion);
+            }
 
-        match protocol_ack {
-            Ok(v) => {
-                if v != version {
-                    // this isn't right...
-                    return Err(Error::ClientMismatch);
+            // Wait for the protocol ACK
+            let protocol_ack_bytes = recv_msg(&mut stream).await?;
+            let (protocol_ack, _): (ProtocolRequestAck, _) =
+                hubpack::deserialize(&protocol_ack_bytes)?;
+
+            match protocol_ack {
+                Ok(v) => {
+                    if v != version {
+                        // this isn't right...
+                        return Err(Error::ClientMismatch);
+                    }
                 }
+                Err(_) => return Err(Error::ClientGaveUp),
             }
-            Err(_) => return Err(Error::ClientGaveUp),
-        }
 
-        // Right now all protocols are the same
-        info!(self.log, "Running with protocol version {version}");
-
-        // get Nonce from client
-        let client_nonce = recv_msg(&mut stream).await?;
-        let client_nonce = Nonce::try_from(client_nonce)?;
-
-        // generate & send Nonce to client
-        let nonce = Nonce::from_platform_rng()?;
-        send_msg(&mut stream, nonce.as_ref()).await?;
-
-        // get attestation & verify it before sending it
-        // 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)?;
-        dice_verifier::verify_attestation(
-            &attest_data.certs[0],
-            &attest_data.attestation,
-            &attest_data.log,
-            &client_nonce,
-        )?;
-
-        // get & verify client attestation cert chain
-        let client_cert_chain = recv_msg(&mut stream).await?;
-        let client_cert_chain = certs_from_der(&client_cert_chain)?;
-        let root = dice_verifier::verify_cert_chain(
-            &client_cert_chain,
-            Some(&self.roots),
-        )?;
-        let client_platform_id =
-            dice_mfg_msgs::PlatformId::try_from(&client_cert_chain)?;
-        info!(
-            self.log,
-            "Cert chain from peer \"{}\" verified against root \"{}\"",
-            client_platform_id.as_str()?,
-            root.tbs_certificate.subject,
-        );
-
-        if tq_platform_id != client_platform_id {
-            return Err(Error::PlatformIdMismatch);
-        }
-        info!(
-            self.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)?;
-        send_msg(&mut stream, &cert_chain_der).await?;
-
-        // get measurement log from client
-        let client_log = recv_msg(&mut stream).await?;
-        let (client_log, _): (Log, _) = hubpack::deserialize(&client_log)?;
-
-        // send server measurement log to client
-        let mut buf = vec![0u8; Log::MAX_SIZE];
-        let len = hubpack::serialize(&mut buf, &attest_data.log)?;
-        send_msg(&mut stream, &buf[..len]).await?;
-
-        // get attestation from client
-        let client_attestation = recv_msg(&mut stream).await?;
-        let (client_attestation, _): (Attestation, _) =
-            hubpack::deserialize(&client_attestation)?;
-
-        // verify client attestation
-        dice_verifier::verify_attestation(
-            &client_cert_chain[0],
-            &client_attestation,
-            &client_log,
-            &nonce,
-        )?;
-        info!(self.log, "Peer attestation verified");
-
-        // appraise measurements from client attestation against reference
-        // measurements
-        let measurements =
-            MeasurementSet::from_artifacts(&client_cert_chain, &client_log)?;
-        let result =
-            match dice_verifier::verify_measurements(&measurements, &corpus) {
+            // Right now all protocols are the same
+            info!(self.log, "Running with protocol version {version}");
+
+            // get Nonce from client
+            let client_nonce = recv_msg(&mut stream).await?;
+            let client_nonce = Nonce::try_from(client_nonce)?;
+
+            // generate & send Nonce to client
+            let nonce = Nonce::from_platform_rng()?;
+            send_msg(&mut stream, nonce.as_ref()).await?;
+
+            // get attestation & verify it before sending it
+            // 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)?;
+            dice_verifier::verify_attestation(
+                &attest_data.certs[0],
+                &attest_data.attestation,
+                &attest_data.log,
+                &client_nonce,
+            )?;
+
+            // get & verify client attestation cert chain
+            let client_cert_chain = recv_msg(&mut stream).await?;
+            let client_cert_chain = certs_from_der(&client_cert_chain)?;
+            let root = dice_verifier::verify_cert_chain(
+                &client_cert_chain,
+                Some(&self.roots),
+            )?;
+            let client_platform_id =
+                dice_mfg_msgs::PlatformId::try_from(&client_cert_chain)?;
+            info!(
+                self.log,
+                "Cert chain from peer \"{}\" verified against root \"{}\"",
+                client_platform_id.as_str()?,
+                root.tbs_certificate.subject,
+            );
+
+            if tq_platform_id != client_platform_id {
+                return Err(Error::PlatformIdMismatch);
+            }
+            info!(
+                self.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)?;
+            send_msg(&mut stream, &cert_chain_der).await?;
+
+            // get measurement log from client
+            let client_log = recv_msg(&mut stream).await?;
+            let (client_log, _): (Log, _) = hubpack::deserialize(&client_log)?;
+
+            // send server measurement log to client
+            let mut buf = vec![0u8; Log::MAX_SIZE];
+            let len = hubpack::serialize(&mut buf, &attest_data.log)?;
+            send_msg(&mut stream, &buf[..len]).await?;
+
+            // get attestation from client
+            let client_attestation = recv_msg(&mut stream).await?;
+            let (client_attestation, _): (Attestation, _) =
+                hubpack::deserialize(&client_attestation)?;
+
+            // verify client attestation
+            dice_verifier::verify_attestation(
+                &client_cert_chain[0],
+                &client_attestation,
+                &client_log,
+                &nonce,
+            )?;
+            info!(self.log, "Peer attestation verified");
+
+            // appraise measurements from client attestation against reference
+            // measurements
+            let measurements = MeasurementSet::from_artifacts(
+                &client_cert_chain,
+                &client_log,
+            )?;
+            let result = match dice_verifier::verify_measurements(
+                &measurements,
+                &corpus,
+            ) {
                 Ok(()) => {
                     info!(self.log, "Peer measurements appraised successfully");
                     true
@@ -414,11 +430,12 @@ impl Server {
                 }
             };
 
-        // hubpack the attestation and send to client
-        let mut buf = vec![0u8; Attestation::MAX_SIZE];
-        let len = hubpack::serialize(&mut buf, &attest_data.attestation)?;
-        send_msg(&mut stream, &buf[..len]).await?;
+            // hubpack the attestation and send to client
+            let mut buf = vec![0u8; Attestation::MAX_SIZE];
+            let len = hubpack::serialize(&mut buf, &attest_data.attestation)?;
+            send_msg(&mut stream, &buf[..len]).await?;
 
-        Ok((Stream::new(stream.into(), client_platform_id, result), addr))
+            Ok((Stream::new(stream.into(), client_platform_id, result), addr))
+        }
     }
 }