feat: add identity API for client app

src/ic_tee_agent/src/agent.rs

@@ -1,7 +1,7 @@
 use candid::Principal;
 use ciborium::into_writer;
 use ed25519_consensus::SigningKey;
-use ic_agent::Agent;
+use ic_agent::{Agent, Identity, Signature};
 use ic_cose::{
     agent::{query_call, update_call},
     client::Client,
@@ -77,6 +77,10 @@ impl TEEAgent {
         self.agents.read().await.identity.is_authenticated()
     }
 
+    pub async fn with_identity<R>(&self, f: impl FnOnce(&TEEIdentity) -> R) -> R {
+        f(&self.agents.read().await.identity)
+    }
+
     pub async fn set_identity(&self, identity: &BasicIdentity, expires_in_ms: u64) {
         let mut id = {
             self.agents.read().await.identity.clone()

src/ic_tee_agent/src/http/authentication.rs

@@ -32,6 +32,9 @@ pub static HEADER_IC_TEE_INSTANCE: HeaderName = HeaderName::from_static("ic-tee-
 /// Authenticated caller principal (or anonymous principal)
 pub static HEADER_IC_TEE_CALLER: HeaderName = HeaderName::from_static("ic-tee-caller");
 
+/// TEE session ID that client should provide in subsequent requests
+pub static HEADER_IC_TEE_SESSION: HeaderName = HeaderName::from_static("ic-tee-session");
+
 /// The `UserSignature` struct represents an end user's signature and provides methods to
 /// parse and validate the signature from HTTP headers.
 ///

src/ic_tee_agent/src/identity.rs

@@ -143,6 +143,45 @@ impl TEEIdentity {
     }
 }
 
+impl Identity for &TEEIdentity {
+    fn sender(&self) -> Result<Principal, String> {
+        match &self.identity {
+            InnerIdentity::Anonymous(id) => id.sender(),
+            InnerIdentity::Delegated(id) => id.sender(),
+        }
+    }
+    fn public_key(&self) -> Option<Vec<u8>> {
+        match &self.identity {
+            InnerIdentity::Anonymous(id) => id.public_key(),
+            InnerIdentity::Delegated(id) => id.public_key(),
+        }
+    }
+    fn sign(&self, content: &EnvelopeContent) -> Result<Signature, String> {
+        match &self.identity {
+            InnerIdentity::Anonymous(id) => id.sign(content),
+            InnerIdentity::Delegated(id) => id.sign(content),
+        }
+    }
+    fn sign_delegation(&self, content: &Delegation) -> Result<Signature, String> {
+        match &self.identity {
+            InnerIdentity::Anonymous(id) => id.sign_delegation(content),
+            InnerIdentity::Delegated(id) => id.sign_delegation(content),
+        }
+    }
+    fn sign_arbitrary(&self, content: &[u8]) -> Result<Signature, String> {
+        match &self.identity {
+            InnerIdentity::Anonymous(id) => id.sign_arbitrary(content),
+            InnerIdentity::Delegated(id) => id.sign_arbitrary(content),
+        }
+    }
+    fn delegation_chain(&self) -> Vec<SignedDelegation> {
+        match &self.identity {
+            InnerIdentity::Anonymous(id) => id.delegation_chain(),
+            InnerIdentity::Delegated(id) => id.delegation_chain(),
+        }
+    }
+}
+
 impl Identity for TEEIdentity {
     fn sender(&self) -> Result<Principal, String> {
         match &self.identity {

src/ic_tee_agent/src/lib.rs

@@ -20,5 +20,11 @@ pub struct RPCRequest {
     pub params: ByteBuf, // params should be encoded in CBOR format
 }
 
+#[derive(Clone, Debug, Serialize)]
+pub struct RPCRequestRef<'a> {
+    pub method: &'a str,
+    pub params: &'a ByteBuf,
+}
+
 // result should be encoded in CBOR format
 pub type RPCResponse = Result<ByteBuf, String>;

src/ic_tee_nitro_attestation/src/types.rs

@@ -3,7 +3,7 @@ use serde::{Deserialize, Serialize};
 use serde_bytes::ByteBuf;
 use std::collections::BTreeMap;
 
-#[derive(CandidType, Debug, Default, Serialize, Deserialize, Copy, Clone, PartialEq)]
+#[derive(CandidType, Serialize, Deserialize, Default, Debug, Clone, PartialEq)]
 pub enum Digest {
     /// SHA256
     SHA256,
@@ -15,7 +15,7 @@ pub enum Digest {
 }
 
 /// An attestation response.  This is also used for sealing data.
-#[derive(CandidType, Serialize, Default, Deserialize, Debug, Clone, PartialEq)]
+#[derive(CandidType, Serialize, Deserialize, Default, Debug, Clone, PartialEq)]
 pub struct Attestation {
     /// Issuing NSM ID
     pub module_id: String,
@@ -46,7 +46,7 @@ pub struct Attestation {
     pub nonce: Option<ByteBuf>,
 }
 
-#[derive(CandidType, Serialize, Deserialize, Debug, Clone, PartialEq)]
+#[derive(CandidType, Serialize, Deserialize, Default, Debug, Clone, PartialEq)]
 pub struct AttestationRequest {
     pub public_key: Option<ByteBuf>,
     pub user_data: Option<ByteBuf>,

src/ic_tee_nitro_gateway/src/handler.rs

@@ -1,7 +1,7 @@
 use axum::{
     body::Body,
     extract::{Request, State},
-    http::{header, uri::Uri, StatusCode},
+    http::{header, uri::Uri, HeaderMap, StatusCode},
     response::IntoResponse,
 };
 use ciborium::from_reader;
@@ -13,20 +13,26 @@ use ic_cose_types::{
 use ic_tee_agent::{
     agent::TEEAgent,
     http::{
-        Content, UserSignature, ANONYMOUS_PRINCIPAL, HEADER_IC_TEE_CALLER,
+        sign_digest_to_headers, Content, UserSignature, ANONYMOUS_PRINCIPAL, HEADER_IC_TEE_CALLER,
         HEADER_IC_TEE_CONTENT_DIGEST, HEADER_IC_TEE_DELEGATION, HEADER_IC_TEE_ID,
-        HEADER_IC_TEE_INSTANCE, HEADER_IC_TEE_PUBKEY, HEADER_IC_TEE_SIGNATURE,
-        HEADER_X_FORWARDED_FOR, HEADER_X_FORWARDED_HOST, HEADER_X_FORWARDED_PROTO,
+        HEADER_IC_TEE_INSTANCE, HEADER_IC_TEE_PUBKEY, HEADER_IC_TEE_SESSION,
+        HEADER_IC_TEE_SIGNATURE, HEADER_X_FORWARDED_FOR, HEADER_X_FORWARDED_HOST,
+        HEADER_X_FORWARDED_PROTO,
     },
     RPCRequest, RPCResponse,
 };
 use ic_tee_cdk::{
-    CanisterRequest, TEEAppInformation, TEEAppInformationJSON, TEEAttestation, TEEAttestationJSON,
+    AttestationUserRequest, CanisterRequest, TEEAppInformation, TEEAppInformationJSON,
+    TEEAttestation, TEEAttestationJSON,
 };
 use ic_tee_nitro_attestation::AttestationRequest;
 use serde_bytes::{ByteArray, ByteBuf};
-use std::sync::Arc;
+use std::{
+    collections::{BTreeMap, HashMap},
+    sync::Arc,
+};
 use structured_logger::unix_ms;
+use tokio::sync::RwLock;
 
 use crate::{attestation::sign_attestation, crypto, ic_sig_verifier::verify_sig, TEE_KIND};
 
@@ -39,6 +45,7 @@ pub struct AppState {
     tee_agent: Arc<TEEAgent>,
     root_secret: [u8; 48],
     upstream_port: Option<u16>,
+    sessions: Arc<RwLock<BTreeMap<String, Option<String>>>>,
 }
 
 impl AppState {
@@ -47,18 +54,62 @@ impl AppState {
         tee_agent: Arc<TEEAgent>,
         root_secret: [u8; 48],
         upstream_port: Option<u16>,
+        apps: Vec<String>,
     ) -> Self {
         let http_client = Arc::new(
             hyper_util::client::legacy::Client::<(), ()>::builder(TokioExecutor::new())
                 .build(HttpConnector::new()),
         );
+        let mut sessions = BTreeMap::new();
+
+        // initialize apps session with empty values
+        for app in apps {
+            for s in app.split(&[',', ' ', ';'][..]) {
+                sessions.insert(s.trim().to_ascii_lowercase(), None);
+            }
+        }
         Self {
             info,
             http_client,
             tee_agent,
             root_secret,
             upstream_port,
+            sessions: Arc::new(RwLock::new(sessions)),
+        }
+    }
+
+    pub async fn valid_session(&self, header: &HeaderMap) -> bool {
+        let sessions = self.sessions.read().await;
+        if sessions.is_empty() {
+            return true;
+        }
+
+        if let Some(sess) = header.get(&HEADER_IC_TEE_SESSION) {
+            if let Ok(sess) = sess.to_str() {
+                let sess: Vec<&str> = sess.split("-").collect();
+                return sess.len() == 2
+                    && sessions
+                        .get(sess[0])
+                        .is_some_and(|v| v.as_ref().is_some_and(|v| v == sess[1]));
+            }
         }
+
+        false
+    }
+
+    pub async fn register_session(&self, app: String) -> Option<String> {
+        let mut sessions = self.sessions.write().await;
+        // app should be registered at bootstrap
+        if let Some(sess) = sessions.get_mut(&app) {
+            // app session should not be registered
+            if sess.is_none() {
+                let session = format!("{}-{}", app, unix_ms());
+                *sess = Some(session.clone());
+                return Some(session);
+            }
+        }
+        // register session failed
+        None
     }
 
     pub fn a256gcm_key(&self, derivation_path: Vec<ByteBuf>) -> ByteArray<32> {
@@ -219,11 +270,19 @@ pub async fn get_attestation(State(app): State<AppState>, req: Request) -> impl
 
 /// local_server: POST /attestation
 pub async fn local_sign_attestation(
-    State(_app): State<AppState>,
-    ct: Content<AttestationRequest>,
+    State(app): State<AppState>,
+    headers: HeaderMap,
+    cr: Content<AttestationUserRequest<ByteBuf>>,
 ) -> impl IntoResponse {
-    match ct {
-        Content::CBOR(req, _) => match sign_attestation(req) {
+    if !app.valid_session(&headers).await {
+        return StatusCode::UNAUTHORIZED.into_response();
+    }
+
+    match cr {
+        Content::CBOR(req, _) => match sign_attestation(AttestationRequest {
+            user_data: Some(ByteBuf::from(to_cbor_bytes(&req))),
+            ..Default::default()
+        }) {
             Ok(doc) => Content::CBOR(
                 TEEAttestation {
                     kind: TEE_KIND.to_string(),
@@ -236,7 +295,10 @@ pub async fn local_sign_attestation(
                 Content::Text::<()>(err, Some(StatusCode::INTERNAL_SERVER_ERROR)).into_response()
             }
         },
-        Content::JSON(req, _) => match sign_attestation(req) {
+        Content::JSON(req, _) => match sign_attestation(AttestationRequest {
+            user_data: Some(ByteBuf::from(to_cbor_bytes(&req))),
+            ..Default::default()
+        }) {
             Ok(doc) => Content::JSON(
                 TEEAttestationJSON {
                     kind: TEE_KIND.to_string(),
@@ -256,8 +318,12 @@ pub async fn local_sign_attestation(
 /// local_server: POST /canister/query
 pub async fn local_query_canister(
     State(app): State<AppState>,
+    headers: HeaderMap,
     ct: Content<CanisterRequest>,
 ) -> impl IntoResponse {
+    if !app.valid_session(&headers).await {
+        return StatusCode::UNAUTHORIZED.into_response();
+    }
     match ct {
         Content::CBOR(req, _) => {
             if forbid_canister_request(&req, app.info.as_ref()) {
@@ -276,8 +342,13 @@ pub async fn local_query_canister(
 /// local_server: POST /canister/update
 pub async fn local_update_canister(
     State(app): State<AppState>,
+    headers: HeaderMap,
     ct: Content<CanisterRequest>,
 ) -> impl IntoResponse {
+    if !app.valid_session(&headers).await {
+        return StatusCode::UNAUTHORIZED.into_response();
+    }
+
     match ct {
         Content::CBOR(req, _) => {
             if forbid_canister_request(&req, app.info.as_ref()) {
@@ -296,8 +367,13 @@ pub async fn local_update_canister(
 /// local_server: POST /keys
 pub async fn local_call_keys(
     State(app): State<AppState>,
+    headers: HeaderMap,
     ct: Content<RPCRequest>,
 ) -> impl IntoResponse {
+    if !app.valid_session(&headers).await {
+        return StatusCode::UNAUTHORIZED.into_response();
+    }
+
     match ct {
         Content::CBOR(req, _) => {
             let res = handle_keys_request(&req, &app);
@@ -307,6 +383,25 @@ pub async fn local_call_keys(
     }
 }
 
+/// local_server: POST /identity
+pub async fn local_call_identity(
+    State(app): State<AppState>,
+    headers: HeaderMap,
+    ct: Content<RPCRequest>,
+) -> impl IntoResponse {
+    match ct {
+        Content::CBOR(req, _) => {
+            if req.method != "register_session" && !app.valid_session(&headers).await {
+                return StatusCode::UNAUTHORIZED.into_response();
+            }
+
+            let res = handle_identity_request(&req, &app).await;
+            Content::CBOR(res, None).into_response()
+        }
+        _ => StatusCode::UNSUPPORTED_MEDIA_TYPE.into_response(),
+    }
+}
+
 /// Handles all other public server requests and proxies them to the upstream service.
 ///
 /// If the request contains valid headers for `ic-tee-pubkey`, `ic-tee-content-digest`,
@@ -399,6 +494,32 @@ pub async fn proxy(
     }
 }
 
+async fn handle_identity_request(req: &RPCRequest, app: &AppState) -> RPCResponse {
+    match req.method.as_str() {
+        "sign_http" => {
+            let digest: ByteArray<32> = from_reader(req.params.as_slice()).map_err(format_error)?;
+            let mut headers = HeaderMap::new();
+            app.tee_agent
+                .with_identity(|id| sign_digest_to_headers(id, &mut headers, digest.as_slice()))
+                .await?;
+            let headers: HashMap<&str, &str> = headers
+                .iter()
+                .map(|(k, v)| (k.as_str(), v.to_str().unwrap()))
+                .collect();
+            Ok(to_cbor_bytes(&headers).into())
+        }
+        "register_session" => {
+            let name: String = from_reader(req.params.as_slice()).map_err(format_error)?;
+            if let Some(sess) = app.register_session(name).await {
+                Ok(to_cbor_bytes(&sess).into())
+            } else {
+                Err("register session failed".to_string())
+            }
+        }
+        _ => Err(format!("unsupported method {}", req.method)),
+    }
+}
+
 fn handle_keys_request(req: &RPCRequest, app: &AppState) -> RPCResponse {
     match req.method.as_str() {
         "a256gcm_key" => {

src/ic_tee_nitro_gateway/src/main.rs

@@ -77,6 +77,9 @@ struct Cli {
     #[clap(long)]
     upstream_port: Option<u16>,
 
+    #[clap(long, default_value = "dTEE")]
+    apps: Vec<String>,
+
     /// where the logtail server is running on host (e.g. 127.0.0.1:9999)
     #[clap(long)]
     bootstrap_logtail: Option<String>,
@@ -309,11 +312,13 @@ async fn bootstrap(cli: Cli) -> Result<()> {
                 routing::post(handler::local_update_canister),
             )
             .route("/keys", routing::post(handler::local_call_keys))
+            .route("/identity", routing::post(handler::local_call_identity))
             .with_state(handler::AppState::new(
                 info.clone(),
                 tee_agent.clone(),
                 root_secret,
                 None,
+                cli.apps,
             ));
         let addr: SocketAddr = LOCAL_HTTP_ADDR.parse().map_err(anyhow::Error::new)?;
         let listener = tokio::net::TcpListener::bind(&addr)
@@ -344,6 +349,7 @@ async fn bootstrap(cli: Cli) -> Result<()> {
                 tee_agent.clone(),
                 [0u8; 48],
                 None,
+                Vec::new(),
             ));
         let addr: SocketAddr = PUBLIC_HTTP_ADDR.parse().map_err(anyhow::Error::new)?;