From a8ea693072a10e2051673d552c49462c8233e2ca Mon Sep 17 00:00:00 2001
From: 0xZensh <txr1883@gmail.com>
Date: Wed, 6 Nov 2024 13:34:50 +0800
Subject: [PATCH] fix: fix TEEAgent

---
 src/ic_tee_agent/src/agent.rs        | 12 ++++++++++--
 src/ic_tee_nitro_gateway/src/main.rs | 19 ++++++++++++-------
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/src/ic_tee_agent/src/agent.rs b/src/ic_tee_agent/src/agent.rs
index 7cac857..68f324e 100644
--- a/src/ic_tee_agent/src/agent.rs
+++ b/src/ic_tee_agent/src/agent.rs
@@ -110,8 +110,16 @@ impl TEEAgent {
         self.sign_in(kind, attestation).await
     }
 
-    pub async fn upgrade_identity_with(&self, id: &BasicIdentity, expires_in_ms: u64) {
-        self.identity.write().await.upgrade_with(id, expires_in_ms);
+    pub async fn upgrade_identity_with(&self, identity: &BasicIdentity, expires_in_ms: u64) {
+        let mut id = {
+            let id = self.identity.read().await;
+            id.clone()
+            // drop read lock
+        };
+        id.upgrade_with(identity, expires_in_ms);
+        self.agent.write().await.set_identity(id.clone());
+        let mut w = self.identity.write().await;
+        *w = id;
     }
 
     pub async fn get_cose_secret(&self, path: SettingPath) -> Result<[u8; 32], String> {
diff --git a/src/ic_tee_nitro_gateway/src/main.rs b/src/ic_tee_nitro_gateway/src/main.rs
index d7047ec..f48b0c1 100644
--- a/src/ic_tee_nitro_gateway/src/main.rs
+++ b/src/ic_tee_nitro_gateway/src/main.rs
@@ -127,7 +127,7 @@ async fn serve() -> Result<()> {
 
     log::info!(target: "server",
         elapsed = start.elapsed().as_millis() as u64;
-       "tee_agent sign_in, principal: {:?}", tee_agent.principal().await.to_text());
+       "sign_in, principal: {:?}", tee_agent.principal().await.to_text());
 
     let upgrade_identity =
         if let Some(v) = cli.configuration_upgrade_identity {
@@ -153,7 +153,7 @@ async fn serve() -> Result<()> {
             .map_err(anyhow::Error::msg)?;
         log::info!(target: "server",
             elapsed = start.elapsed().as_millis() as u64;
-            "tee_agent get_cose_secret for upgrade_identity, principal: {:?}", subject.to_text());
+            "get_cose_secret for upgrade_identity, principal: {:?}", subject.to_text());
 
         let setting = tee_agent
             .get_cose_setting(id_path)
@@ -161,7 +161,7 @@ async fn serve() -> Result<()> {
             .map_err(anyhow::Error::msg)?;
         log::info!(target: "server",
             elapsed = start.elapsed().as_millis() as u64;
-            "tee_agent get_cose_setting for upgrade_identity, principal: {:?}", subject.to_text());
+            "get_cose_setting for upgrade_identity, principal: {:?}", subject.to_text());
 
         let ed25519_secret = decrypt_payload(setting, secret).map_err(anyhow::Error::msg)?;
         let ed25519_secret: [u8; 32] = ed25519_secret.try_into().map_err(|val: Vec<u8>| {
@@ -174,7 +174,7 @@ async fn serve() -> Result<()> {
 
         log::info!(target: "server",
             elapsed = start.elapsed().as_millis() as u64;
-            "tee_agent upgrade_identity, principal: {:?}", tee_agent.principal().await.to_text());
+            "upgrade_identity, principal: {:?}", tee_agent.principal().await.to_text());
         Some(id)
     } else {
         None
@@ -195,6 +195,11 @@ async fn serve() -> Result<()> {
         registration_canister: None,
     };
 
+    log::info!(target: "server",
+        info:serde = info,
+        elapsed = start.elapsed().as_millis() as u64;
+        "TEE app information, principal: {:?}", principal.to_text());
+
     let http_client = Arc::new(handler::new_client());
     let tee_agent = Arc::new(tee_agent);
     let info = Arc::new(info);
@@ -257,7 +262,7 @@ async fn serve() -> Result<()> {
         let listener = tokio::net::TcpListener::bind(&addr)
             .await
             .map_err(anyhow::Error::new)?;
-        log::warn!(target: "local_server", "{}@{} listening on {:?}", APP_NAME, APP_VERSION, addr);
+        log::warn!(target: "server", "local {}@{} listening on {:?}", APP_NAME, APP_VERSION, addr);
         axum::serve(listener, app)
             .with_graceful_shutdown(shutdown_future)
             .await
@@ -277,7 +282,7 @@ async fn serve() -> Result<()> {
             .map_err(anyhow::Error::msg)?;
         log::info!(target: "server",
             elapsed = start.elapsed().as_millis() as u64;
-            "tee_agent get_cose_secret for tls");
+            "get_cose_secret for TLS");
 
         let setting = tee_agent
             .get_cose_setting(SettingPath {
@@ -291,7 +296,7 @@ async fn serve() -> Result<()> {
             .map_err(anyhow::Error::msg)?;
         log::info!(target: "server",
             elapsed = start.elapsed().as_millis() as u64;
-            "tee_agent get_cose_setting for tls");
+            "get_cose_setting for TLS");
 
         let tls = decrypt_tls(setting, secret).map_err(anyhow::Error::msg)?;
         let app = Router::new()