Address code review comments

matter-labs · Jan 17, 2025 · afa524c · afa524c
1 parent 2d04ba0
commit afa524c
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 33 deletions.
diff --git a/bin/tee-key-preexec/src/main.rs b/bin/tee-key-preexec/src/main.rs
@@ -8,7 +8,6 @@
 
 use anyhow::{Context, Result};
 use clap::Parser;
-use core::convert::AsRef;
 use secp256k1::{rand, Secp256k1};
 use std::{ffi::OsString, os::unix::process::CommandExt, process::Command};
 use teepot::{
@@ -46,7 +45,7 @@ fn main_with_error() -> Result<()> {
     let ethereum_address = public_key_to_ethereum_address(&verifying_key);
     let report_data = ReportDataV1 { ethereum_address };
     let report_data_bytes: [u8; 64] = report_data.into();
-    let tee_type = match get_quote(report_data_bytes.as_ref()) {
+    let tee_type = match get_quote(&report_data_bytes) {
         Ok((tee_type, quote)) => {
             // save quote to file
             std::fs::write(TEE_QUOTE_FILE, quote)?;
@@ -87,24 +86,3 @@ fn main() -> Result<()> {
     }
     ret
 }
-
-#[cfg(test)]
-mod tests {
-    use secp256k1::{PublicKey, Secp256k1, SecretKey};
-
-    use super::*;
-
-    #[test]
-    fn test_public_key_to_address() {
-        let secp = Secp256k1::new();
-        let secret_key_bytes =
-            hex::decode("c87509a1c067bbde78beb793e6fa76530b6382a4c0241e5e4a9ec0a0f44dc0d3")
-                .unwrap();
-        let secret_key = SecretKey::from_slice(&secret_key_bytes).unwrap();
-        let public_key = PublicKey::from_secret_key(&secp, &secret_key);
-        let expected_address = hex::decode("627306090abaB3A6e1400e9345bC60c78a8BEf57").unwrap();
-        let address = public_key_to_ethereum_address(&public_key);
-
-        assert_eq!(address, expected_address.as_slice());
-    }
-}
diff --git a/bin/verify-attestation/src/main.rs b/bin/verify-attestation/src/main.rs
@@ -3,15 +3,16 @@
 
 //! Tool for SGX attestation and batch signature verification
 
-use anyhow::{Context, Result};
+use anyhow::{anyhow, Context, Result};
 use clap::{Args, Parser, Subcommand};
 use core::convert::TryInto;
 use hex::encode;
-use secp256k1::{Message, PublicKey};
+use secp256k1::Message;
 use std::{fs, io::Read, path::PathBuf, str::FromStr, time::UNIX_EPOCH};
 use teepot::{
     client::TcbLevel,
     ethereum::recover_signer,
+    prover::reportdata::ReportData,
     quote::{error, tee_qv_get_collateral, verify_quote_with_collateral, QuoteVerificationResult},
 };
 use zksync_basic_types::H256;
@@ -87,15 +88,14 @@ fn verify_signature(
     quote_verification_result: &QuoteVerificationResult,
     signature_args: &SignatureArgs,
 ) -> Result<()> {
-    let reportdata = &quote_verification_result.quote.get_report_data();
-    let public_key = PublicKey::from_slice(reportdata)?;
-    println!("Public key from attestation quote: {}", public_key);
+    let report_data = ReportData::try_from(quote_verification_result.quote.get_report_data())?;
+    let ethereum_address_from_quote = match report_data {
+        ReportData::V1(report_data_v1) => report_data_v1.ethereum_address,
+        _ => return Err(anyhow!("Unsupported report data version")),
+    };
     let signature_bytes: &[u8] = &fs::read(&signature_args.signature_file)?;
-    let ethereum_address_from_quote = &quote_verification_result.quote.get_report_data()[..20];
-    let root_hash_bytes = signature_args.root_hash.as_bytes();
-    let root_hash_msg = Message::from_digest_slice(root_hash_bytes)?;
-    let ethereum_address_from_signature =
-        recover_signer(&signature_bytes.try_into()?, &root_hash_msg)?;
+    let root_hash = Message::from_digest_slice(signature_args.root_hash.as_bytes())?;
+    let ethereum_address_from_signature = recover_signer(&signature_bytes.try_into()?, &root_hash)?;
     let verification_successful = ethereum_address_from_signature == ethereum_address_from_quote;
 
     println!(

diff --git a/crates/teepot/src/prover/reportdata.rs b/crates/teepot/src/prover/reportdata.rs
@@ -35,6 +35,11 @@ pub enum ReportData {
 }
 
 fn report_data_to_bytes(data: &[u8], version: u8) -> [u8; REPORT_DATA_LENGTH] {
+    debug_assert!(
+        data.len() < REPORT_DATA_LENGTH, // Ensure there is space for the version byte
+        "Data length exceeds maximum of {} bytes",
+        REPORT_DATA_LENGTH
+    );
     let mut bytes = [0u8; REPORT_DATA_LENGTH];
     bytes[..data.len()].copy_from_slice(data);
     bytes[REPORT_DATA_LENGTH - 1] = version;