fix: fix tee sign in

Cargo.toml

@@ -18,7 +18,7 @@ strip = true
 opt-level = 's'
 
 [workspace.package]
-version = "0.2.7"
+version = "0.2.8"
 edition = "2021"
 repository = "https://github.com/ldclabs/ic-tee"
 keywords = ["tee", "canister", "icp", "nitro"]
@@ -50,6 +50,7 @@ serde = "1"
 serde_json = "1"
 serde_bytes = "0.11"
 sha2 = "0.10"
+sha3 = "0.10"
 ic-cdk = "0.17"
 ic-stable-structures = "0.6"
 ic-canister-sig-creation = "1.1"

src/ic_tee_agent/src/agent.rs

@@ -14,7 +14,7 @@ use crate::{BasicIdentity, TEEIdentity};
 
 #[derive(Clone)]
 pub struct TEEAgent {
-    pub auth_canister: Principal,
+    pub identity_canister: Principal,
     pub cose_canister: Principal,
     pub identity: TEEIdentity,
     pub agent: Agent,
@@ -23,8 +23,8 @@ pub struct TEEAgent {
 impl TEEAgent {
     pub async fn new(
         host: &str,
-        authentication_canister: Principal,
-        configuration_canister: Principal,
+        identity_canister: Principal,
+        cose_canister: Principal,
     ) -> Result<Self, String> {
         let identity = TEEIdentity::new();
         let agent = Agent::builder()
@@ -37,8 +37,8 @@ impl TEEAgent {
             agent.fetch_root_key().await.map_err(format_error)?;
         }
         Ok(Self {
-            auth_canister: authentication_canister,
-            cose_canister: configuration_canister,
+            identity_canister,
+            cose_canister,
             identity,
             agent,
         })
@@ -48,7 +48,7 @@ impl TEEAgent {
         let mut agent = self.agent.clone();
         agent.set_identity(identity.clone());
         Self {
-            auth_canister: self.auth_canister,
+            identity_canister: self.identity_canister,
             cose_canister: self.cose_canister,
             identity,
             agent,
@@ -79,15 +79,15 @@ impl TEEAgent {
         let mut id = self.identity.clone();
         let (user_key, delegation) = {
             let res: Result<SignInResponse, String> = self
-                .canister_update(&self.auth_canister, "sign_in", (kind, attestation))
+                .canister_update(&self.identity_canister, "sign_in", (kind, attestation))
                 .await
                 .map_err(format_error)?;
             let res = res?;
             let user_key = res.user_key.to_vec();
 
             let res: Result<SignedDelegation, String> = self
                 .canister_query(
-                    &self.auth_canister,
+                    &self.identity_canister,
                     "get_delegation",
                     (
                         res.seed,

src/ic_tee_cdk/Cargo.toml

@@ -14,6 +14,7 @@ candid = { workspace = true }
 ciborium = { workspace = true }
 serde = { workspace = true }
 serde_bytes = { workspace = true }
+sha3 = { workspace = true }
 ic-canister-sig-creation = { workspace = true }
 
 [dev-dependencies]

src/ic_tee_cdk/src/identity.rs

@@ -3,7 +3,7 @@ use ic_canister_sig_creation::CanisterSigPublicKey;
 use serde::{Deserialize, Serialize};
 use serde_bytes::ByteBuf;
 
-use crate::to_cbor_bytes;
+use crate::{sha3_256_n, to_cbor_bytes};
 
 pub const SESSION_EXPIRES_IN_MS: u64 = 1000 * 3600 * 24; // 1 day
 
@@ -43,9 +43,9 @@ pub fn canister_user_key(
     sub_seed: Option<&[u8]>,
 ) -> CanisterSigPublicKey {
     let seed = if let Some(sub_seed) = sub_seed {
-        to_cbor_bytes(&(kind, seed, sub_seed))
+        to_cbor_bytes(&(kind, sha3_256_n([seed, sub_seed])))
     } else {
-        to_cbor_bytes(&(kind, seed))
+        to_cbor_bytes(&(kind, sha3_256_n([seed])))
     };
     CanisterSigPublicKey::new(canister, seed)
 }
@@ -67,12 +67,12 @@ mod tests {
         println!("{:?}", const_hex::encode(user_key.as_slice()));
         assert!(is_sub(&user_key, canister.as_slice()));
         assert!(is_sub(&user_key, kind.as_bytes()));
-        assert!(is_sub(&user_key, seed.as_slice()));
+        assert!(!is_sub(&user_key, seed.as_slice()));
 
         let sub_seed = [1u8, 2u8, 3u8, 4u8];
         let user_key2 = canister_user_key(canister, kind, &seed, Some(&sub_seed)).to_der();
         assert_ne!(user_key, user_key2);
-        assert!(is_sub(&user_key2, seed.as_slice()));
-        assert!(is_sub(&user_key2, sub_seed.as_slice()));
+        assert!(!is_sub(&user_key2, seed.as_slice()));
+        assert!(!is_sub(&user_key2, sub_seed.as_slice()));
     }
 }

src/ic_tee_cdk/src/lib.rs

@@ -1,6 +1,7 @@
 use candid::Principal;
 use ciborium::into_writer;
 use serde::{Deserialize, Serialize};
+use sha3::{Digest, Sha3_256};
 
 pub mod identity;
 
@@ -20,6 +21,14 @@ pub fn to_cbor_bytes(obj: &impl Serialize) -> Vec<u8> {
     buf
 }
 
+pub fn sha3_256_n<const N: usize>(array: [&[u8]; N]) -> [u8; 32] {
+    let mut hasher = Sha3_256::new();
+    for data in array {
+        hasher.update(data);
+    }
+    hasher.finalize().into()
+}
+
 #[derive(Clone, Debug, Default, Deserialize, Serialize)]
 pub struct AttestationUserRequest<T> {
     pub method: String,

src/ic_tee_cli/src/main.rs

@@ -1,7 +1,6 @@
 use anyhow::Result;
 use candid::{pretty::candid::value::pp_value, CandidType, IDLValue, Principal};
 use clap::{Parser, Subcommand};
-use ed25519::pkcs8::{DecodePrivateKey, KeypairBytes};
 use ed25519_consensus::SigningKey;
 use ic_agent::{
     identity::{AnonymousIdentity, BasicIdentity},
@@ -82,7 +81,7 @@ pub enum Commands {
     /// verify a TEE attestation document
     TeeVerify {
         /// TEE kind to verify
-        #[arg(long, default_value = "Nitro")]
+        #[arg(long, default_value = "NITRO")]
         kind: String,
 
         /// TEE attestation document
@@ -254,7 +253,7 @@ async fn main() -> Result<()> {
                     .get_cose_encrypted_key(&SettingPath {
                         ns: ns.clone(),
                         user_owned: false,
-                        subject: path.subject.clone(),
+                        subject: path.subject,
                         key: ByteBuf::from(COSE_SECRET_PERMANENT_KEY.as_bytes()),
                         version: 0,
                     })

src/ic_tee_identity/src/api.rs

@@ -29,7 +29,7 @@ fn whoami() -> Principal {
 #[ic_cdk::update]
 fn sign_in(kind: String, attestation: ByteBuf) -> Result<SignInResponse, String> {
     let attestation = match kind.as_str() {
-        "Nitro" => parse_and_verify(attestation.as_slice())?,
+        "NITRO" => parse_and_verify(attestation.as_slice())?,
         _ => Err("unsupported attestation kind".to_string())?,
     };
 
@@ -104,7 +104,7 @@ fn get_delegation(
     session_key: ByteBuf,
     expiration: u64,
 ) -> Result<SignedDelegation, String> {
-    if seed.len() > 48 {
+    if seed.len() > 128 {
         return Err("invalid seed length".to_string());
     }
     let delegation_hash = delegation_signature_msg(session_key.as_slice(), expiration, None);

src/ic_tee_nitro_gateway/src/main.rs

@@ -121,6 +121,7 @@ async fn main() -> Result<(), BoxError> {
         Ok(_) => Ok(()),
         Err(err) => {
             log::error!(target: LOG_TARGET, "server error: {:?}", err);
+            tokio::time::sleep(Duration::from_secs(1)).await;
             Err(err)
         }
     }
@@ -139,6 +140,10 @@ async fn bootstrap(cli: Cli) -> Result<(), BoxError> {
     let cose_canister = Principal::from_text(cli.cose_canister)?;
     let mut tee_agent = TEEAgent::new(&cli.ic_host, identity_canister, cose_canister).await?;
 
+    log::info!(target: LOG_TARGET,
+        elapsed = start.elapsed().as_millis() as u64;
+       "start with principal: {:?}", tee_agent.get_principal().to_text());
+
     let namespace = cli.cose_namespace;
     let session_expires_in_ms = cli.session_expires_in_ms.unwrap_or(SESSION_EXPIRES_IN_MS);