Error handling

Author: ameba23

src/attestation.rs

@@ -1,30 +1,40 @@
-use crate::AttestationError;
 use sha2::{Digest, Sha256};
+use thiserror::Error;
 use tokio_rustls::rustls::pki_types::CertificateDer;
 use x509_parser::prelude::*;
 
+/// Represents a CVM technology with quote generation and verification
 pub trait AttestationPlatform: Clone + Send + 'static {
-    fn create_attestation(&self, cert_chain: &[CertificateDer<'_>], exporter: [u8; 32]) -> Vec<u8>;
+    fn create_attestation(
+        &self,
+        cert_chain: &[CertificateDer<'_>],
+        exporter: [u8; 32],
+    ) -> Result<Vec<u8>, AttestationError>;
 
     fn verify_attestation(
         &self,
         input: Vec<u8>,
         cert_chain: &[CertificateDer<'_>],
         exporter: [u8; 32],
-    ) -> bool;
+    ) -> Result<(), AttestationError>;
 }
 
+/// For testing
 #[derive(Clone)]
 pub struct MockAttestation;
 
 impl AttestationPlatform for MockAttestation {
     /// Mocks creating an attestation
-    fn create_attestation(&self, cert_chain: &[CertificateDer<'_>], exporter: [u8; 32]) -> Vec<u8> {
+    fn create_attestation(
+        &self,
+        cert_chain: &[CertificateDer<'_>],
+        exporter: [u8; 32],
+    ) -> Result<Vec<u8>, AttestationError> {
         let mut quote_input = [0u8; 64];
-        let pki_hash = get_pki_hash_from_certificate_chain(cert_chain).unwrap();
+        let pki_hash = get_pki_hash_from_certificate_chain(cert_chain)?;
         quote_input[..32].copy_from_slice(&pki_hash);
         quote_input[32..].copy_from_slice(&exporter);
-        quote_input.to_vec()
+        Ok(quote_input.to_vec())
     }
 
     /// Mocks verifying an attestation
@@ -33,16 +43,20 @@ impl AttestationPlatform for MockAttestation {
         input: Vec<u8>,
         cert_chain: &[CertificateDer<'_>],
         exporter: [u8; 32],
-    ) -> bool {
+    ) -> Result<(), AttestationError> {
         let mut quote_input = [0u8; 64];
-        let pki_hash = get_pki_hash_from_certificate_chain(cert_chain).unwrap();
+        let pki_hash = get_pki_hash_from_certificate_chain(cert_chain)?;
         quote_input[..32].copy_from_slice(&pki_hash);
         quote_input[32..].copy_from_slice(&exporter);
 
-        input == quote_input
+        if input != quote_input {
+            return Err(AttestationError::InputMismatch);
+        }
+        Ok(())
     }
 }
 
+/// For no CVM platform (eg: for one-sided remote-attested TLS)
 #[derive(Clone)]
 pub struct NoAttestation;
 
@@ -52,18 +66,22 @@ impl AttestationPlatform for NoAttestation {
         &self,
         _cert_chain: &[CertificateDer<'_>],
         _exporter: [u8; 32],
-    ) -> Vec<u8> {
-        Vec::new()
+    ) -> Result<Vec<u8>, AttestationError> {
+        Ok(Vec::new())
     }
 
     /// Mocks verifying an attestation
     fn verify_attestation(
         &self,
-        _input: Vec<u8>,
+        input: Vec<u8>,
         _cert_chain: &[CertificateDer<'_>],
         _exporter: [u8; 32],
-    ) -> bool {
-        true
+    ) -> Result<(), AttestationError> {
+        if input.is_empty() {
+            Ok(())
+        } else {
+            Err(AttestationError::AttestationGivenWhenNoneExpected)
+        }
     }
 }
 
@@ -80,3 +98,18 @@ fn get_pki_hash_from_certificate_chain(
     hasher.update(key_bytes);
     Ok(hasher.finalize().into())
 }
+
+/// An error when generating or verifying an attestation
+#[derive(Error, Debug)]
+pub enum AttestationError {
+    #[error("Certificate chain is empty")]
+    NoCertificate,
+    #[error("X509 parse: {0}")]
+    X509Parse(#[from] x509_parser::asn1_rs::Err<x509_parser::error::X509Error>),
+    #[error("X509: {0}")]
+    X509(#[from] x509_parser::error::X509Error),
+    #[error("Quote input is not as expected")]
+    InputMismatch,
+    #[error("Configuration mismatch - expected no remote attestation")]
+    AttestationGivenWhenNoneExpected,
+}

src/lib.rs

@@ -3,7 +3,6 @@ mod attestation;
 pub use attestation::{AttestationPlatform, MockAttestation, NoAttestation};
 
 use std::{net::SocketAddr, sync::Arc};
-use thiserror::Error;
 use tokio::io::{self, AsyncReadExt, AsyncWriteExt};
 use tokio::net::{TcpListener, TcpStream, ToSocketAddrs};
 use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer, ServerName};
@@ -118,7 +117,9 @@ impl<L: AttestationPlatform, R: AttestationPlatform> ProxyServer<L, R> {
                 )
                 .unwrap();
 
-            let attestation = local_attestation_platform.create_attestation(&cert_chain, exporter);
+            let attestation = local_attestation_platform
+                .create_attestation(&cert_chain, exporter)
+                .unwrap();
             let attestation_length_prefix = length_prefix(&attestation);
 
             tls_stream
@@ -135,9 +136,9 @@ impl<L: AttestationPlatform, R: AttestationPlatform> ProxyServer<L, R> {
             let mut buf = vec![0; length];
             tls_stream.read_exact(&mut buf).await.unwrap();
 
-            if !remote_attestation_platform.verify_attestation(buf, &cert_chain, exporter) {
-                panic!("Cannot verify attestation");
-            };
+            remote_attestation_platform
+                .verify_attestation(buf, &cert_chain, exporter)
+                .unwrap();
 
             let outbound = TcpStream::connect(target).await.unwrap();
 
@@ -252,11 +253,14 @@ impl<L: AttestationPlatform, R: AttestationPlatform> ProxyClient<L, R> {
             let mut buf = vec![0; length];
             tls_stream.read_exact(&mut buf).await.unwrap();
 
-            if !remote_attestation_platform.verify_attestation(buf, &cert_chain, exporter) {
-                panic!("Cannot verify attestation");
-            };
+            remote_attestation_platform
+                .verify_attestation(buf, &cert_chain, exporter)
+                .unwrap();
+
+            let attestation = local_attestation_platform
+                .create_attestation(&cert_chain, exporter)
+                .unwrap();
 
-            let attestation = local_attestation_platform.create_attestation(&cert_chain, exporter);
             let attestation_length_prefix = length_prefix(&attestation);
 
             tls_stream
@@ -287,17 +291,6 @@ fn length_prefix(input: &[u8]) -> [u8; 4] {
     len.to_be_bytes()
 }
 
-/// An error when generating an attestation
-#[derive(Error, Debug)]
-pub enum AttestationError {
-    #[error("Certificate chain is empty")]
-    NoCertificate,
-    #[error("X509 parse: {0}")]
-    X509Parse(#[from] x509_parser::asn1_rs::Err<x509_parser::error::X509Error>),
-    #[error("X509: {0}")]
-    X509(#[from] x509_parser::error::X509Error),
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;