fix: fix ic_tee_identity

ldclabs · Nov 6, 2024 · c5fab1e · c5fab1e
1 parent caed515
commit c5fab1e
Show file tree

Hide file tree

Showing 11 changed files with 140 additions and 53 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -17,7 +17,7 @@ strip = true
 opt-level = 's'
 
 [workspace.package]
-version = "0.1.1"
+version = "0.1.2"
 edition = "2021"
 repository = "https://github.com/ldclabs/ic-tee"
 keywords = ["tee", "canister", "icp", "nitro"]

diff --git a/src/ic_tee_agent/Cargo.toml b/src/ic_tee_agent/Cargo.toml
@@ -25,3 +25,6 @@ bytes = { workspace = true }
 mime = { workspace = true }
 ic_cose_types = { workspace = true }
 ic_tee_cdk = { path = "../ic_tee_cdk", version = "0.1" }
+
+[dev-dependencies]
+const-hex = { workspace = true }
diff --git a/src/ic_tee_agent/src/agent.rs b/src/ic_tee_agent/src/agent.rs
@@ -16,6 +16,7 @@ use crate::{setting::get_cose_secret, BasicIdentity, TEEIdentity};
 
 #[derive(Clone)]
 pub struct TEEAgent {
+    is_local: bool,
     identity: Arc<RwLock<TEEIdentity>>,
     agent: Arc<RwLock<Agent>>,
     authentication_canister: Principal,
@@ -36,13 +37,25 @@ impl TEEAgent {
             .build()
             .map_err(format_error)?;
         Ok(Self {
+            is_local: host.starts_with("http://"),
             identity: Arc::new(RwLock::new(identity)),
             agent: Arc::new(RwLock::new(agent)),
             authentication_canister,
             configuration_canister,
         })
     }
 
+    pub async fn init(&self) -> Result<(), String> {
+        if self.is_local {
+            let agent = self.agent.read().await;
+            agent
+                .fetch_root_key()
+                .await
+                .map_err(|err| format!("fetch root key failed: {:?}", err))?;
+        }
+        Ok(())
+    }
+
     pub async fn principal(&self) -> Principal {
         self.identity.read().await.principal()
     }
@@ -75,11 +88,7 @@ impl TEEAgent {
             .query_call(
                 &self.authentication_canister,
                 "get_delegation",
-                (
-                    id.principal(),
-                    ByteBuf::from(id.session_key()),
-                    res.expiration,
-                ),
+                (res.seed, ByteBuf::from(id.session_key()), res.expiration),
             )
             .await?;
         let res = res?;
@@ -219,3 +228,81 @@ where
     let output = Decode!(res.as_slice(), Out).map_err(format_error)?;
     Ok(output)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use ic_cose_types::types::state::StateInfo;
+
+    #[tokio::test(flavor = "current_thread")]
+    #[ignore]
+    async fn test_ic_tee_identity() {
+        let host = "https://icp-api.io";
+        // let host = "http://127.0.0.1:4943";
+        let authentication_canister = Principal::from_text("e7tgb-6aaaa-aaaap-akqfa-cai").unwrap();
+        let configuration_canister = Principal::from_text("53cyg-yyaaa-aaaap-ahpua-cai").unwrap();
+        let tee_agent = TEEAgent::new(host, authentication_canister, configuration_canister)
+            .expect("Failed to create TEEAgent");
+        tee_agent.init().await.unwrap();
+
+        println!(
+            "tee_agent init principal: {:?}",
+            tee_agent.principal().await.to_text()
+        );
+        let pubkey = tee_agent.session_key().await;
+        let res: Result<SignInResponse, String> = tee_agent
+            .update_call(
+                &tee_agent.authentication_canister,
+                "sign_in_debug",
+                (ByteBuf::from(pubkey),),
+            )
+            .await
+            .unwrap();
+        let res = res.unwrap();
+        let mut id = {
+            let id = tee_agent.identity.read().await;
+            id.clone()
+            // drop read lock
+        };
+
+        println!("user_key: {:?}", const_hex::encode(&res.user_key));
+        id.with_user_key(res.user_key.to_vec());
+        let principal = id.principal();
+        let res: Result<SignedDelegation, String> = tee_agent
+            .query_call(
+                &tee_agent.authentication_canister,
+                "get_delegation",
+                (res.seed, ByteBuf::from(id.session_key()), res.expiration),
+            )
+            .await
+            .unwrap();
+        let res = res.unwrap();
+
+        id.with_delegation(res).unwrap();
+        {
+            tee_agent.agent.write().await.set_identity(id.clone());
+            let mut w = tee_agent.identity.write().await;
+            *w = id;
+            // drop write lock
+        }
+        let principal2 = { tee_agent.agent.read().await.get_principal().unwrap() };
+        assert_eq!(principal, principal2);
+        assert_eq!(principal, tee_agent.principal().await);
+
+        println!("tee_agent sign in principal: {:?}", principal.to_text());
+        println!(
+            "tee_agent sign in seed: {:?}",
+            const_hex::encode(principal.as_slice())
+        );
+        let res: Principal = tee_agent
+            .query_call(&tee_agent.authentication_canister, "whoami", ())
+            .await
+            .unwrap();
+        println!("configuration canister whoami: {:?}", res.to_text());
+        let res: Result<StateInfo, String> = tee_agent
+            .query_call(&tee_agent.configuration_canister, "state_get_info", ())
+            .await
+            .unwrap();
+        println!("configuration canister state: {:?}", res.unwrap());
+    }
+}
diff --git a/src/ic_tee_cdk/src/identity.rs b/src/ic_tee_cdk/src/identity.rs
@@ -29,14 +29,16 @@ pub struct SignInResponse {
     pub expiration: u64,
     /// The user canister public key. This key is used to derive the user principal.
     pub user_key: ByteBuf,
+    /// seed is a part of the user_key
+    pub seed: ByteBuf,
 }
 
 pub fn canister_user_key(
     canister: Principal,
     kind: &str, // should be "Nitro"
     seed: &[u8],
     sub_seed: Option<&[u8]>,
-) -> Vec<u8> {
+) -> CanisterSigPublicKey {
     let len = 1 + kind.len() + 32;
     let mut data = Vec::with_capacity(len);
     data.push(kind.len() as u8);
@@ -50,7 +52,7 @@ pub fn canister_user_key(
     }
     let (_, buf) = data.split_last_chunk_mut::<32>().unwrap();
     hasher.finalize_into(buf.into());
-    CanisterSigPublicKey::new(canister, data).to_der()
+    CanisterSigPublicKey::new(canister, data)
 }
 
 #[cfg(test)]
@@ -66,15 +68,17 @@ mod tests {
         let canister = Principal::from_text("e7tgb-6aaaa-aaaap-akqfa-cai").unwrap();
         let kind = "Nitro";
         let seed = [8u8; 48];
-        let user_key = canister_user_key(canister, kind, &seed, None);
+        let user_key = canister_user_key(canister, kind, &seed, None).to_der();
         assert!(is_sub(&user_key, canister.as_slice()));
         assert!(is_sub(&user_key, kind.to_uppercase().as_bytes()));
         assert!(!is_sub(&user_key, seed.as_slice()));
 
-        let user_key2 = canister_user_key(canister, kind, &seed, Some(&[1u8, 2u8, 3u8, 4u8]));
+        let user_key2 =
+            canister_user_key(canister, kind, &seed, Some(&[1u8, 2u8, 3u8, 4u8])).to_der();
         assert_ne!(user_key, user_key2);
 
-        let user_key3 = canister_user_key(canister, kind, &seed, Some(&[1u8, 2u8, 3u8, 5u8]));
+        let user_key3 =
+            canister_user_key(canister, kind, &seed, Some(&[1u8, 2u8, 3u8, 5u8])).to_der();
         assert_ne!(user_key2, user_key3);
     }
 }
diff --git a/src/ic_tee_cli/src/main.rs b/src/ic_tee_cli/src/main.rs
@@ -182,7 +182,7 @@ async fn main() -> Result<()> {
             let canister = Principal::from_text(&cli.canister)
                 .map_err(|err| anyhow::anyhow!("invalid canister: {:?}", err))?;
             let user_key = canister_user_key(canister, kind, &seed, sub_seed.as_deref());
-            let principal = Principal::self_authenticating(&user_key);
+            let principal = Principal::self_authenticating(user_key.to_der());
 
             println!("principal: {}", principal);
         }

diff --git a/src/ic_tee_identity/ic_tee_identity.did b/src/ic_tee_identity/ic_tee_identity.did
@@ -8,7 +8,11 @@ type InitArgs = record { session_expires_in_ms : nat64; name : text };
 type Result = variant { Ok : SignedDelegation; Err : text };
 type Result_1 = variant { Ok : State; Err : text };
 type Result_2 = variant { Ok : SignInResponse; Err : text };
-type SignInResponse = record { user_key : blob; expiration : nat64 };
+type SignInResponse = record {
+  user_key : blob;
+  seed : blob;
+  expiration : nat64;
+};
 type SignedDelegation = record { signature : blob; delegation : Delegation };
 type State = record {
   session_expires_in_ms : nat64;
@@ -20,7 +24,8 @@ type UpgradeArgs = record {
   name : opt text;
 };
 service : (opt ChainArgs) -> {
-  get_delegation : (principal, blob, nat64) -> (Result) query;
+  get_delegation : (blob, blob, nat64) -> (Result) query;
   get_state : () -> (Result_1) query;
   sign_in : (text, blob) -> (Result_2);
+  whoami : () -> (principal) query;
 }
diff --git a/src/ic_tee_identity/src/api.rs b/src/ic_tee_identity/src/api.rs
@@ -18,6 +18,11 @@ fn get_state() -> Result<store::State, String> {
     Ok(store::state::with(|s| s.clone()))
 }
 
+#[ic_cdk::query]
+fn whoami() -> Principal {
+    ic_cdk::caller()
+}
+
 #[ic_cdk::update]
 fn sign_in(kind: String, attestation: ByteBuf) -> Result<SignInResponse, String> {
     let attestation = match kind.as_str() {
@@ -66,24 +71,27 @@ fn sign_in(kind: String, attestation: ByteBuf) -> Result<SignInResponse, String>
     });
     let expiration = (now_ms + session_expires_in_ms) * MILLISECONDS;
 
-    let principal = Principal::self_authenticating(&user_key);
     let delegation_hash = delegation_signature_msg(pubkey.as_slice(), expiration, None);
-    store::state::add_signature(principal.as_slice(), delegation_hash.as_slice());
+    store::state::add_signature(user_key.seed.as_slice(), delegation_hash.as_slice());
 
     Ok(SignInResponse {
         expiration,
-        user_key: user_key.into(),
+        user_key: user_key.to_der().into(),
+        seed: user_key.seed.into(),
     })
 }
 
 #[ic_cdk::query]
 fn get_delegation(
-    principal: Principal,
+    seed: ByteBuf,
     session_key: ByteBuf,
     expiration: u64,
 ) -> Result<SignedDelegation, String> {
+    if seed.len() > 48 {
+        return Err("invalid seed length".to_string());
+    }
     let delegation_hash = delegation_signature_msg(session_key.as_slice(), expiration, None);
-    let signature = store::state::get_signature(principal.as_slice(), delegation_hash.as_slice())?;
+    let signature = store::state::get_signature(seed.as_slice(), delegation_hash.as_slice())?;
     Ok(SignedDelegation {
         delegation: Delegation {
             pubkey: session_key,

diff --git a/src/ic_tee_logtail/Cargo.toml b/src/ic_tee_logtail/Cargo.toml
@@ -10,18 +10,6 @@ categories.workspace = true
 license.workspace = true
 
 [dependencies]
-candid = { workspace = true, features = ["value", "printer"] }
-serde_bytes = { workspace = true }
 tokio = { workspace = true }
-const-hex = { workspace = true }
-ic-agent = { workspace = true }
-ed25519-consensus = { workspace = true }
-ic_cose_types = { workspace = true }
-rand = { workspace = true }
 anyhow = "1"
 clap = { version = "=4.5", features = ["derive"] }
-pkcs8 = { version = "0.10", features = ["pem"] }
-ed25519 = { version = "2.2", features = ["pem", "pkcs8"] }
-ic_tee_cdk = { path = "../ic_tee_cdk", version = "0.1" }
-ic_tee_agent = { path = "../ic_tee_agent", version = "0.1" }
-ic_tee_nitro_attestation = { path = "../ic_tee_nitro_attestation", version = "0.1" }
diff --git a/src/ic_tee_nitro_gateway/README.md b/src/ic_tee_nitro_gateway/README.md
@@ -61,8 +61,8 @@ sudo nitro-cli build-enclave --docker-uri ghcr.io/ldclabs/ic_tee_nitro_gateway_e
 #     "PCR2": "7584fed461361c6e8c4f56e426f46e610b86ce8eae1cc407f221adf8f5a9053f452eefa3fbae5256e0b17e91ecd4cb3f"
 #   }
 # }
-ic_tee_cli -c e7tgb-6aaaa-aaaap-akqfa-cai identity-derive --seed 349166c4a015d98b39d6cd3c3a65a5c2ed11f4414687500dd0a7310f36b31d223d0f09662fa547d449e209bc3f2c15be
-# principal: ciar7-g7nzs-66aea-eu53p-vtwhv-7aoz2-hlmrv-dzhir-gkses-pbeem-pqe
+ic_tee_cli -c e7tgb-6aaaa-aaaap-akqfa-cai identity-derive --seed ecf476c4d2682dc2f57773085058ee0e36b720d623b5cd6e8dc04388342a4ee9df8150b1e3dd9b5e11cbac02153e9faf
+# principal: yn7la-zcbb2-vnt52-dixn7-u6djx-mlrnd-vux2u-6llr3-v5xdd-2vi6h-oqe
 sudo nitro-cli run-enclave --cpu-count 2 --memory 512 --enclave-cid 88 --eif-path ic_tee_nitro_gateway_enclave_amd64.eif
 # Start allocating memory...
 # Started enclave with enclave-cid: 88, memory: 512 MiB, cpu-ids: [1, 3]
@@ -79,7 +79,7 @@ sudo nitro-cli run-enclave --cpu-count 2 --memory 512 --enclave-cid 88 --eif-pat
 #   "MemoryMiB": 512
 # }
 sudo nitro-cli describe-enclaves
-sudo nitro-cli terminate-enclave --enclave-id i-056e1ab9a31cd77a0-enc192fc732d6e4e41
+sudo nitro-cli terminate-enclave --enclave-id i-056e1ab9a31cd77a0-enc192fce92cafbc6f
 ```