From c1fc624131182ee51366d1b6a6e609ed445f74c8 Mon Sep 17 00:00:00 2001
From: 0xZensh <txr1883@gmail.com>
Date: Fri, 22 Nov 2024 09:32:13 +0800
Subject: [PATCH] chore: add public key challenge in `ic_tee_identity`

---
 Cargo.lock                           |  1 +
 src/ic_tee_identity/Cargo.toml       |  1 +
 src/ic_tee_identity/src/api.rs       | 27 ++++++++++++++++++++++-----
 src/ic_tee_nitro_gateway/src/main.rs |  7 ++++---
 4 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 8ae1c0e..02e635f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2516,6 +2516,7 @@ dependencies = [
  "ic-canister-sig-creation",
  "ic-cdk 0.16.0",
  "ic-certification 2.6.0",
+ "ic-crypto-standalone-sig-verifier",
  "ic-stable-structures",
  "ic_tee_cdk",
  "ic_tee_nitro_attestation",
diff --git a/src/ic_tee_identity/Cargo.toml b/src/ic_tee_identity/Cargo.toml
index 1b0302d..5f5f700 100644
--- a/src/ic_tee_identity/Cargo.toml
+++ b/src/ic_tee_identity/Cargo.toml
@@ -23,6 +23,7 @@ serde_bytes = { workspace = true }
 ic-stable-structures = { workspace = true }
 ic-canister-sig-creation = { workspace = true }
 ic-certification = { workspace = true }
+ic-crypto-standalone-sig-verifier = { workspace = true }
 getrandom = { version = "0.2", features = ["custom"] }
 ic_tee_cdk = { path = "../ic_tee_cdk", version = "0.1" }
 ic_tee_nitro_attestation = { path = "../ic_tee_nitro_attestation", version = "0.1" }
diff --git a/src/ic_tee_identity/src/api.rs b/src/ic_tee_identity/src/api.rs
index 69a4c72..80f43c0 100644
--- a/src/ic_tee_identity/src/api.rs
+++ b/src/ic_tee_identity/src/api.rs
@@ -1,6 +1,9 @@
 use candid::Principal;
 use ciborium::from_reader;
 use ic_canister_sig_creation::delegation_signature_msg;
+use ic_crypto_standalone_sig_verifier::{
+    user_public_key_from_bytes, verify_basic_sig_by_public_key,
+};
 use ic_tee_cdk::{
     canister_user_key, AttestationUserRequest, Delegation, SignInParams, SignInResponse,
     SignedDelegation,
@@ -38,11 +41,25 @@ fn sign_in(kind: String, attestation: ByteBuf) -> Result<SignInResponse, String>
     let pubkey: ByteBuf = attestation
         .public_key
         .ok_or_else(|| "missing public key".to_string())?;
+    let user_data: ByteBuf = attestation
+        .user_data
+        .ok_or_else(|| "missing user data".to_string())?;
+    let sig: ByteBuf = attestation
+        .nonce
+        .ok_or_else(|| "missing nonce".to_string())?;
+
+    let (pk, _) = user_public_key_from_bytes(pubkey.as_slice())
+        .map_err(|err| format!("invalid public key: {:?}", err))?;
+    verify_basic_sig_by_public_key(
+        pk.algorithm_id,
+        user_data.as_slice(),
+        sig.as_slice(),
+        &pk.key,
+    )
+    .map_err(|err| format!("challenge verification failed: {:?}", err))?;
 
-    let req: AttestationUserRequest<SignInParams> = attestation.user_data.map_or_else(
-        || Err("missing user data".to_string()),
-        |data| from_reader(data.as_slice()).map_err(|err| format!("invalid user data: {:?}", err)),
-    )?;
+    let req: AttestationUserRequest<SignInParams> =
+        from_reader(user_data.as_slice()).map_err(|err| format!("invalid user data: {:?}", err))?;
     if req.method != "sign_in" {
         return Err("invalid attestation user request method".to_string());
     }
@@ -51,7 +68,7 @@ fn sign_in(kind: String, attestation: ByteBuf) -> Result<SignInResponse, String>
         Some(SignInParams { id_scope }) => {
             if id_scope == "image" {
                 canister_user_key(ic_cdk::id(), &kind, pcr0.as_slice(), None)
-            } else if id_scope == "enclave" {
+            } else if id_scope == "instance" {
                 canister_user_key(
                     ic_cdk::id(),
                     &kind,
diff --git a/src/ic_tee_nitro_gateway/src/main.rs b/src/ic_tee_nitro_gateway/src/main.rs
index 06ff500..f8da10a 100644
--- a/src/ic_tee_nitro_gateway/src/main.rs
+++ b/src/ic_tee_nitro_gateway/src/main.rs
@@ -45,7 +45,7 @@ struct Cli {
     #[clap(long, value_parser)]
     session_expires_in_ms: Option<u64>,
 
-    // id_scope should be "image" or "enclave", default is "image"
+    // id_scope should be "image" or "instance", default is "image"
     #[clap(long, value_parser)]
     id_scope: Option<String>,
 
@@ -117,12 +117,13 @@ async fn serve(cli: Cli) -> Result<()> {
 
     let user_req = to_cbor_bytes(&user_req);
     let session_key = TEEIdentity::new_session();
-    let public_key = session_key.1.clone();
+    let public_key = session_key.1.clone(); // der encoded public key
+    let sig = session_key.0.sign(&user_req);
 
     let doc = sign_attestation(AttestationRequest {
         public_key: Some(public_key.into()),
         user_data: Some(user_req.clone().into()),
-        nonce: None,
+        nonce: Some(sig.to_bytes().to_vec().into()), // use signature as nonce for challenge
     })
     .map_err(anyhow::Error::msg)?;