Merge pull request #77 from flashbots/peg/url-measurements-file

Allow giving measurement file as url, and improve measurement checking logic

Author: ameba23

src/attestation/measurements.rs

@@ -162,6 +162,8 @@ pub enum MeasurementFormatError {
     BadRegisterIndex,
     #[error("ParseInt: {0}")]
     ParseInt(#[from] std::num::ParseIntError),
+    #[error("Failed to read measurements from URL: {0}")]
+    Reqwest(#[from] reqwest::Error),
 }
 
 /// An accepted measurement value given in the measurements file
@@ -261,9 +263,11 @@ impl MeasurementPolicy {
             .any(|measurement_record| match measurements {
                 MultiMeasurements::Dcap(dcap_measurements) => {
                     if let MultiMeasurements::Dcap(d) = measurement_record.measurements.clone() {
-                        for (k, v) in dcap_measurements.iter() {
-                            if d.get(k).is_some_and(|x| x != v) {
-                                return false;
+                        // All measurements in our policy must be given and must match
+                        for (k, v) in d.iter() {
+                            match dcap_measurements.get(k) {
+                                Some(value) if value == v => {}
+                                _ => return false,
                             }
                         }
                         return true;
@@ -272,9 +276,10 @@ impl MeasurementPolicy {
                 }
                 MultiMeasurements::Azure(azure_measurements) => {
                     if let MultiMeasurements::Azure(a) = measurement_record.measurements.clone() {
-                        for (k, v) in azure_measurements.iter() {
-                            if a.get(k).is_some_and(|x| x != v) {
-                                return false;
+                        for (k, v) in a.iter() {
+                            match azure_measurements.get(k) {
+                                Some(value) if value == v => {}
+                                _ => return false,
                             }
                         }
                         return true;
@@ -303,6 +308,16 @@ impl MeasurementPolicy {
             .any(|a| a.measurements == MultiMeasurements::NoAttestation)
     }
 
+    /// Given either a URL or the path to a file, parse the measurement policy from JSON
+    pub async fn from_file_or_url(file_or_url: String) -> Result<Self, MeasurementFormatError> {
+        if file_or_url.starts_with("https://") || file_or_url.starts_with("http://") {
+            let measurements_json = reqwest::get(file_or_url).await?.bytes().await?;
+            Self::from_json_bytes(measurements_json.to_vec()).await
+        } else {
+            Self::from_file(file_or_url.into()).await
+        }
+    }
+
     /// Given the path to a JSON file containing measurements, return a [MeasurementPolicy]
     pub async fn from_file(measurement_file: PathBuf) -> Result<Self, MeasurementFormatError> {
         let measurements_json = tokio::fs::read(measurement_file).await?;
@@ -449,6 +464,14 @@ mod tests {
                 .unwrap_err(),
             AttestationError::MeasurementsNotAccepted
         ));
+
+        // A non-specific measurement fails
+        assert!(matches!(
+            specific_measurements
+                .check_measurement(&MultiMeasurements::Azure(HashMap::new()))
+                .unwrap_err(),
+            AttestationError::MeasurementsNotAccepted
+        ));
     }
 
     #[tokio::test]
@@ -471,4 +494,31 @@ mod tests {
             AttestationError::MeasurementsNotAccepted
         ));
     }
+
+    #[tokio::test]
+    async fn test_read_remote_buildernet_measurements() {
+        // Check that the buildernet measurements are available and parse correctly
+        let policy = MeasurementPolicy::from_file_or_url(
+            "https://measurements.builder.flashbots.net".to_string(),
+        )
+        .await
+        .unwrap();
+
+        assert!(!policy.accepted_measurements.is_empty());
+
+        assert!(matches!(
+            policy
+                .check_measurement(&MultiMeasurements::NoAttestation)
+                .unwrap_err(),
+            AttestationError::MeasurementsNotAccepted
+        ));
+
+        // A non-specific measurement fails
+        assert!(matches!(
+            policy
+                .check_measurement(&MultiMeasurements::Azure(HashMap::new()))
+                .unwrap_err(),
+            AttestationError::MeasurementsNotAccepted
+        ));
+    }
 }

src/main.rs

@@ -17,9 +17,9 @@ use attested_tls_proxy::{
 struct Cli {
     #[clap(subcommand)]
     command: CliCommand,
-    /// Optional path to file containing JSON measurements to be enforced on the remote party
+    /// Path to file, or URL, containing JSON measurements to be enforced on the remote party
     #[arg(long, global = true, env = "MEASUREMENTS_FILE")]
-    measurements_file: Option<PathBuf>,
+    measurements_file: Option<String>,
     /// If no measurements file is specified, a single attestion type to allow
     #[arg(long, global = true)]
     allowed_remote_attestation_type: Option<String>,
@@ -171,7 +171,9 @@ async fn main() -> anyhow::Result<()> {
     }
 
     let measurement_policy = match cli.measurements_file {
-        Some(server_measurements) => MeasurementPolicy::from_file(server_measurements).await?,
+        Some(server_measurements) => {
+            MeasurementPolicy::from_file_or_url(server_measurements).await?
+        }
         None => {
             let allowed_server_attestation_type: AttestationType = serde_json::from_value(
                 serde_json::Value::String(cli.allowed_remote_attestation_type.ok_or(anyhow!(