Merge pull request #32 from flashbots/peg/attestation-encoding

Use SCALE rather than JSON for encoding attestation payloads

Author: ameba23

Cargo.lock

@@ -27,6 +27,7 @@ bytes = "1.10.1"
 http = "1.3.1"
 serde_json = "1.0.145"
 serde = "1.0.228"
+parity-scale-codec = "3.7.5"
 
 [dev-dependencies]
 rcgen = "0.14.5"

Cargo.toml

@@ -1,6 +1,7 @@
 pub mod measurements;
 
 use measurements::{CvmImageMeasurements, MeasurementRecord, Measurements, PlatformMeasurements};
+use parity_scale_codec::{Decode, Encode};
 use serde::{Deserialize, Serialize};
 use std::{
     fmt::{self, Display, Formatter},
@@ -22,13 +23,19 @@ use x509_parser::prelude::*;
 /// For fetching collateral directly from intel, if no PCCS is specified
 const PCS_URL: &str = "https://api.trustedservices.intel.com";
 
-#[derive(Debug, Serialize, Deserialize)]
+/// This is the type sent over the channel to provide an attestation
+#[derive(Debug, Serialize, Deserialize, Encode, Decode)]
 pub struct AttesationPayload {
+    /// What CVM platform is used (including none)
     pub attestation_type: AttestationType,
+    /// The attestation evidence as bytes - in the case of DCAP this is a quote
     pub attestation: Vec<u8>,
 }
 
 impl AttesationPayload {
+    /// Given an attestation generator (quote generation function for a specific platform)
+    /// return an attestation
+    /// This also takes the certificate chain and exporter as they are given as input to the attestation
     pub fn from_attestation_generator(
         cert_chain: &[CertificateDer<'_>],
         exporter: [u8; 32],
@@ -39,6 +46,15 @@ impl AttesationPayload {
             attestation: attesation_generator.create_attestation(cert_chain, exporter)?,
         })
     }
+
+    /// Create an empty attestation payload for the case that we are running in a non-confidential
+    /// environment
+    pub fn without_attestation() -> Self {
+        Self {
+            attestation_type: AttestationType::None,
+            attestation: Vec::new(),
+        }
+    }
 }
 
 /// Type of attestaion used
@@ -73,6 +89,7 @@ impl AttestationType {
         }
     }
 
+    /// Get a quote generator for this type of platform
     pub fn get_quote_generator(&self) -> Result<Arc<dyn QuoteGenerator>, AttestationError> {
         match self {
             AttestationType::None => Ok(Arc::new(NoQuoteGenerator)),
@@ -85,6 +102,23 @@ impl AttestationType {
     }
 }
 
+/// SCALE encode (used over the wire)
+impl Encode for AttestationType {
+    fn encode(&self) -> Vec<u8> {
+        self.as_str().encode()
+    }
+}
+
+/// SCALE decode
+impl Decode for AttestationType {
+    fn decode<I: parity_scale_codec::Input>(
+        input: &mut I,
+    ) -> Result<Self, parity_scale_codec::Error> {
+        let s: String = String::decode(input)?;
+        serde_json::from_str(&format!("\"{s}\"")).map_err(|_| "Failed to decode enum".into())
+    }
+}
+
 impl Display for AttestationType {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         f.write_str(self.as_str())
@@ -104,9 +138,15 @@ pub trait QuoteGenerator: Send + Sync + 'static {
     ) -> Result<Vec<u8>, AttestationError>;
 }
 
+/// Allows remote attestations to be verified
 #[derive(Clone, Debug)]
 pub struct AttestationVerifier {
+    /// The measurement values we accept
+    ///
+    /// If this is empty, anything will be accepted - but measurements are always injected into HTTP
+    /// headers, so that they can be verified upstream
     accepted_measurements: Vec<MeasurementRecord>,
+    /// A PCCS service to use - defaults to Intel PCS
     pccs_url: Option<String>,
 }
 

src/attestation/mod.rs

@@ -4,11 +4,10 @@ use attestation::{measurements::Measurements, AttestationError, AttestationType}
 pub use attestation::{DcapTdxQuoteGenerator, NoQuoteGenerator, QuoteGenerator};
 use bytes::Bytes;
 use http::HeaderValue;
-use http_body_util::combinators::BoxBody;
-use http_body_util::BodyExt;
-use hyper::service::service_fn;
-use hyper::Response;
+use http_body_util::{combinators::BoxBody, BodyExt};
+use hyper::{service::service_fn, Response};
 use hyper_util::rt::TokioIo;
+use parity_scale_codec::{Decode, Encode};
 use thiserror::Error;
 use tokio::sync::{mpsc, oneshot};
 use tokio_rustls::rustls::server::{VerifierBuilderError, WebPkiClientVerifier};
@@ -191,15 +190,12 @@ impl ProxyServer {
         let remote_cert_chain = connection.peer_certificates().map(|c| c.to_owned());
 
         // If we are in a CVM, generate an attestation
-        let attestation = if local_quote_generator.attestation_type() != AttestationType::None {
-            serde_json::to_vec(&AttesationPayload::from_attestation_generator(
-                &cert_chain,
-                exporter,
-                local_quote_generator,
-            )?)?
-        } else {
-            Vec::new()
-        };
+        let attestation = AttesationPayload::from_attestation_generator(
+            &cert_chain,
+            exporter,
+            local_quote_generator,
+        )?
+        .encode();
 
         // Write our attestation to the channel, with length prefix
         let attestation_length_prefix = length_prefix(&attestation);
@@ -215,24 +211,20 @@ impl ProxyServer {
         let mut buf = vec![0; length];
         tls_stream.read_exact(&mut buf).await?;
 
+        let remote_attestation_payload = AttesationPayload::decode(&mut &buf[..])?;
+        let remote_attestation_type = remote_attestation_payload.attestation_type;
+
         // If we expect an attestaion from the client, verify it and get measurements
-        let (measurements, remote_attestation_type) = if attestation_verifier.has_remote_attestion()
-        {
-            let remote_attestation_payload: AttesationPayload = serde_json::from_slice(&buf)?;
-
-            let remote_attestation_type = remote_attestation_payload.attestation_type;
-            (
-                attestation_verifier
-                    .verify_attestation(
-                        remote_attestation_payload,
-                        &remote_cert_chain.ok_or(ProxyError::NoClientAuth)?,
-                        exporter,
-                    )
-                    .await?,
-                remote_attestation_type,
-            )
+        let measurements = if attestation_verifier.has_remote_attestion() {
+            attestation_verifier
+                .verify_attestation(
+                    remote_attestation_payload,
+                    &remote_cert_chain.ok_or(ProxyError::NoClientAuth)?,
+                    exporter,
+                )
+                .await?
         } else {
-            (None, AttestationType::None)
+            None
         };
 
         // Setup an HTTP server
@@ -607,7 +599,7 @@ impl ProxyClient {
         let mut buf = vec![0; length];
         tls_stream.read_exact(&mut buf).await?;
 
-        let remote_attestation_payload: AttesationPayload = serde_json::from_slice(&buf)?;
+        let remote_attestation_payload = AttesationPayload::decode(&mut &buf[..])?;
         let remote_attestation_type = remote_attestation_payload.attestation_type;
 
         // Verify the remote attestation against our accepted measurements
@@ -617,13 +609,14 @@ impl ProxyClient {
 
         // If we are in a CVM, provide an attestation
         let attestation = if local_quote_generator.attestation_type() != AttestationType::None {
-            serde_json::to_vec(&AttesationPayload::from_attestation_generator(
+            AttesationPayload::from_attestation_generator(
                 &cert_chain.ok_or(ProxyError::NoClientAuth)?,
                 exporter,
                 local_quote_generator,
-            )?)?
+            )?
+            .encode()
         } else {
-            Vec::new()
+            AttesationPayload::without_attestation().encode()
         };
 
         // Send our attestation (or zero bytes) prefixed with length
@@ -705,7 +698,7 @@ async fn get_tls_cert_with_config(
     let mut buf = vec![0; length];
     tls_stream.read_exact(&mut buf).await?;
 
-    let remote_attestation_payload: AttesationPayload = serde_json::from_slice(&buf)?;
+    let remote_attestation_payload = AttesationPayload::decode(&mut &buf[..])?;
 
     let _measurements = attestation_verifier
         .verify_attestation(remote_attestation_payload, &remote_cert_chain, exporter)
@@ -741,6 +734,8 @@ pub enum ProxyError {
     OneShotRecv(#[from] oneshot::error::RecvError),
     #[error("Failed to send request, connection to proxy-server dropped")]
     MpscSend,
+    #[error("Serialization: {0}")]
+    Serialization(#[from] parity_scale_codec::Error),
 }
 
 impl From<mpsc::error::SendError<RequestWithResponseSender>> for ProxyError {

src/lib.rs

@@ -243,11 +243,13 @@ fn load_tls_cert_and_key(
     Ok(TlsCertAndKey { key, cert_chain })
 }
 
+/// load certificates from a PEM-encoded file
 fn load_certs_pem(path: PathBuf) -> std::io::Result<Vec<CertificateDer<'static>>> {
     rustls_pemfile::certs(&mut std::io::BufReader::new(File::open(path)?))
         .collect::<Result<Vec<_>, _>>()
 }
 
+/// load TLS private key from a PEM-encoded file
 fn load_private_key_pem(path: PathBuf) -> anyhow::Result<PrivateKeyDer<'static>> {
     let mut reader = std::io::BufReader::new(File::open(path)?);