From 959c076380a1f853322d85dfa1fc937841d21d6c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Patryk=20B=C4=99za?= <patryk.beza@gmail.com>
Date: Fri, 20 Dec 2024 12:19:19 +0100
Subject: [PATCH] feat(tee-key-preexec): add support for Solidity-compatible
 pubkey in report_data

---
 Cargo.lock                      |  2 ++
 Cargo.toml                      |  1 +
 bin/tee-key-preexec/Cargo.toml  |  2 ++
 bin/tee-key-preexec/src/main.rs | 44 +++++++++++++++++++++++++++------
 4 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index e46dde7..7c92cf0 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5106,8 +5106,10 @@ version = "0.3.0"
 dependencies = [
  "anyhow",
  "clap 4.5.23",
+ "hex",
  "rand",
  "secp256k1 0.29.1",
+ "sha3",
  "teepot",
  "tracing",
  "tracing-log 0.2.0",
diff --git a/Cargo.toml b/Cargo.toml
index 113ab1e..eb19d4b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -47,6 +47,7 @@ serde = { version = "1", features = ["derive", "rc"] }
 serde_json = "1"
 serde_with = { version = "3.8", features = ["base64", "hex"] }
 sha2 = "0.10.8"
+sha3 = "0.10.8"
 signature = "2.2.0"
 tdx-attest-rs = { version = "0.1.2", git = "https://github.com/intel/SGXDataCenterAttestationPrimitives.git", rev = "aa239d25a437a28f3f4de92c38f5b6809faac842" }
 teepot = { path = "crates/teepot" }
diff --git a/bin/tee-key-preexec/Cargo.toml b/bin/tee-key-preexec/Cargo.toml
index e5a62cc..7732d68 100644
--- a/bin/tee-key-preexec/Cargo.toml
+++ b/bin/tee-key-preexec/Cargo.toml
@@ -12,8 +12,10 @@ repository.workspace = true
 [dependencies]
 anyhow.workspace = true
 clap.workspace = true
+hex.workspace = true
 rand.workspace = true
 secp256k1.workspace = true
+sha3.workspace = true
 teepot.workspace = true
 tracing.workspace = true
 tracing-log.workspace = true
diff --git a/bin/tee-key-preexec/src/main.rs b/bin/tee-key-preexec/src/main.rs
index 70be9c4..9b35b2e 100644
--- a/bin/tee-key-preexec/src/main.rs
+++ b/bin/tee-key-preexec/src/main.rs
@@ -8,7 +8,8 @@
 
 use anyhow::{Context, Result};
 use clap::Parser;
-use secp256k1::{rand, Keypair, PublicKey, Secp256k1, SecretKey};
+use secp256k1::{rand, PublicKey, Secp256k1, SecretKey};
+use sha3::{Digest, Keccak256};
 use std::{ffi::OsString, os::unix::process::CommandExt, process::Command};
 use teepot::quote::get_quote;
 use tracing::error;
@@ -28,6 +29,19 @@ struct Args {
     cmd_args: Vec<OsString>,
 }
 
+/// Converts a public key into an Ethereum address by hashing the encoded public key with Keccak256.
+pub fn public_key_to_address(public: &PublicKey) -> [u8; 20] {
+    let public_key_bytes = public.serialize_uncompressed();
+
+    // Skip the first byte (0x04) which indicates uncompressed key
+    let hash: [u8; 32] = Keccak256::digest(&public_key_bytes[1..]).into();
+
+    // Take the last 20 bytes of the hash to get the Ethereum address
+    let mut address = [0u8; 20];
+    address.copy_from_slice(&hash[12..]);
+    address
+}
+
 fn main_with_error() -> Result<()> {
     LogTracer::init().context("Failed to set logger")?;
 
@@ -37,14 +51,11 @@ fn main_with_error() -> Result<()> {
     tracing::subscriber::set_global_default(subscriber).context("Failed to set logger")?;
 
     let args = Args::parse();
-
     let mut rng = rand::thread_rng();
     let secp = Secp256k1::new();
-    let keypair = Keypair::new(&secp, &mut rng);
-    let signing_key = SecretKey::from_keypair(&keypair);
-    let verifying_key = PublicKey::from_keypair(&keypair);
-    let verifying_key_bytes = verifying_key.serialize();
-    let tee_type = match get_quote(verifying_key_bytes.as_ref()) {
+    let (signing_key, verifying_key) = secp.generate_keypair(&mut rng);
+    let ethereum_address = public_key_to_address(&verifying_key);
+    let tee_type = match get_quote(ethereum_address.as_ref()) {
         Ok((tee_type, quote)) => {
             // save quote to file
             std::fs::write(TEE_QUOTE_FILE, quote)?;
@@ -85,3 +96,22 @@ fn main() -> Result<()> {
     }
     ret
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_public_key_to_address() {
+        let secp = Secp256k1::new();
+        let secret_key_bytes =
+            hex::decode("4c0883a69102937d6231471b5dbb6204fe512961708279e0b4aeb6a364ee18bb")
+                .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_address(&public_key);
+
+        assert_eq!(address, expected_address.as_slice());
+    }
+}