WIP: implement ic_tee_nitro_gateway

Cargo.toml

@@ -51,10 +51,12 @@ ic-stable-structures = "0.6"
 ic-canister-sig-creation = "1.1"
 ic-certification = "2.6"
 ic-agent = "0.39"
+ic_cose_types = "0.3"
 getrandom = { version = "0.2", features = ["custom"] }
 coset = "0.3.8"
 x509-parser = { version = "0.16.0" }
 ed25519-consensus = "2.1"
+x25519-dalek = { version = "2", features = ["static_secrets"] }
 rand = "0.8"
 tokio = { version = "1", features = ["full"] }
 tokio-util = "0.7"

src/ic_tee_agent/Cargo.toml

@@ -12,6 +12,7 @@ license.workspace = true
 [dependencies]
 candid = { workspace = true }
 ed25519-consensus = { workspace = true }
+x25519-dalek = { workspace = true }
 ciborium = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
@@ -22,4 +23,5 @@ tokio = { workspace = true }
 axum = { workspace = true }
 bytes = { workspace = true }
 mime = { workspace = true }
+ic_cose_types = { workspace = true }
 ic_tee_cdk = { path = "../ic_tee_cdk", version = "0.1" }

src/ic_tee_agent/src/agent.rs

@@ -3,12 +3,20 @@ use candid::{
     CandidType, Decode, Principal,
 };
 use ic_agent::Agent;
-use ic_tee_cdk::{format_error, SignInResponse, SignedDelegation};
+use ic_cose_types::{
+    cose::{
+        ecdh::ecdh_x25519, encrypt0::cose_decrypt0, format_error, get_cose_key_secret,
+        CborSerializable, CoseKey,
+    },
+    types::{setting::SettingInfo, ECDHInput, ECDHOutput, SettingPath},
+};
+use ic_tee_cdk::{SignInResponse, SignedDelegation};
 use serde_bytes::ByteBuf;
 use std::sync::Arc;
 use tokio::sync::RwLock;
+use x25519_dalek::{PublicKey, StaticSecret};
 
-use crate::TEEIdentity;
+use crate::{rand_bytes, BasicIdentity, TEEIdentity};
 
 #[derive(Clone)]
 pub struct TEEAgent {
@@ -97,6 +105,53 @@ 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 get_cose_secret(&self, path: SettingPath) -> Result<[u8; 32], String> {
+        let nonce: [u8; 12] = rand_bytes();
+        let secret: [u8; 32] = rand_bytes();
+        let secret = StaticSecret::from(secret);
+        let public = PublicKey::from(&secret);
+
+        let subject = if let Some(subject) = path.subject {
+            subject
+        } else {
+            self.principal().await
+        };
+        let res: Result<ECDHOutput<ByteBuf>, String> = self
+            .update_call(
+                &self.configuration_canister,
+                "ecdh_cose_encrypted_key",
+                (
+                    path,
+                    ECDHInput {
+                        nonce: nonce.into(),
+                        public_key: public.to_bytes().into(),
+                    },
+                ),
+            )
+            .await;
+        let res = res?;
+        let (shared_secret, _) = ecdh_x25519(secret.to_bytes(), *res.public_key);
+        let add = subject.as_slice();
+        let kek = cose_decrypt0(&res.payload, &shared_secret.to_bytes(), add)?;
+        let key =
+            CoseKey::from_slice(&kek).map_err(|err| format!("invalid COSE key: {:?}", err))?;
+        let secret = get_cose_key_secret(key)?;
+        secret.try_into().map_err(|val: Vec<u8>| {
+            format!("invalid COSE secret, expected 32 bytes, got {}", val.len())
+        })
+    }
+
+    pub async fn get_cose_setting(&self, path: SettingPath) -> Result<SettingInfo, String> {
+        let res: Result<SettingInfo, String> = self
+            .update_call(&self.configuration_canister, "setting_get", (path,))
+            .await;
+        res
+    }
+
     pub async fn update_call<In, Out>(
         &self,
         canister_id: &Principal,
@@ -144,4 +199,36 @@ impl TEEAgent {
         let output = Decode!(res.as_slice(), Out).map_err(format_error)?;
         Ok(output)
     }
+
+    pub async fn update_call_raw(
+        &self,
+        canister_id: &Principal,
+        method_name: &str,
+        input: Vec<u8>,
+    ) -> Result<Vec<u8>, String> {
+        self.agent
+            .read()
+            .await
+            .update(canister_id, method_name)
+            .with_arg(input)
+            .call_and_wait()
+            .await
+            .map_err(format_error)
+    }
+
+    pub async fn query_call_raw(
+        &self,
+        canister_id: &Principal,
+        method_name: &str,
+        input: Vec<u8>,
+    ) -> Result<Vec<u8>, String> {
+        self.agent
+            .read()
+            .await
+            .query(canister_id, method_name)
+            .with_arg(input)
+            .call()
+            .await
+            .map_err(format_error)
+    }
 }

src/ic_tee_agent/src/http.rs

@@ -1,6 +1,6 @@
 use axum::{
     async_trait,
-    body::{Body, Bytes},
+    body::Bytes,
     extract::{FromRequest, Request},
     http::{
         header::{self, HeaderMap, HeaderValue},

src/ic_tee_agent/src/identity.rs

@@ -1,15 +1,15 @@
 use candid::Principal;
 use ed25519_consensus::SigningKey;
 use ic_agent::{
-    identity::{
-        AnonymousIdentity, BasicIdentity, DelegatedIdentity, Delegation, Identity, SignedDelegation,
-    },
+    identity::{DelegatedIdentity, Delegation, SignedDelegation},
     {agent::EnvelopeContent, Signature},
 };
 use ic_tee_cdk::identity;
 use rand::thread_rng;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
+pub use ic_agent::identity::{AnonymousIdentity, BasicIdentity, Identity};
+
 enum InnerIdentity {
     Anonymous(AnonymousIdentity),
     Delegated(DelegatedIdentity),
@@ -35,7 +35,7 @@ impl Clone for TEEIdentity {
     fn clone(&self) -> Self {
         Self {
             identity: match &self.identity {
-                InnerIdentity::Anonymous(id) => InnerIdentity::Anonymous(id.clone()),
+                InnerIdentity::Anonymous(id) => InnerIdentity::Anonymous(*id),
                 InnerIdentity::Delegated(_) => {
                     InnerIdentity::Delegated(DelegatedIdentity::new_unchecked(
                         self.user_key.clone(),
@@ -48,7 +48,7 @@ impl Clone for TEEIdentity {
             delegation: self.delegation.clone(),
             user_key: self.user_key.clone(),
             session_key: self.session_key.clone(),
-            principal: self.principal.clone(),
+            principal: self.principal,
             expiration: self.expiration,
         }
     }
@@ -90,6 +90,32 @@ impl TEEIdentity {
         self.principal
     }
 
+    pub fn upgrade_with(&mut self, identity: &BasicIdentity, expires_in_ms: u64) {
+        let expiration = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap_or_default()
+            .saturating_add(Duration::from_millis(expires_in_ms));
+        let delegation = Delegation {
+            pubkey: self.session_key.clone(),
+            expiration: expiration.as_nanos() as u64,
+            targets: None,
+        };
+        let signature = identity.sign_delegation(&delegation).unwrap();
+        self.principal = identity.sender().unwrap();
+        self.user_key = identity.public_key().unwrap();
+        self.expiration = delegation.expiration;
+        self.delegation = vec![SignedDelegation {
+            delegation,
+            signature: signature.signature.unwrap(),
+        }];
+        let id = DelegatedIdentity::new_unchecked(
+            self.user_key.clone(),
+            Box::new(BasicIdentity::from_signing_key(self.signing_key.clone())),
+            self.delegation.clone(),
+        );
+        self.identity = InnerIdentity::Delegated(id);
+    }
+
     pub fn with_user_key(&mut self, user_key: Vec<u8>) {
         self.principal = Principal::self_authenticating(&user_key);
         self.user_key = user_key;
@@ -160,3 +186,8 @@ impl Identity for TEEIdentity {
         }
     }
 }
+
+pub fn identity_from(secret: [u8; 32]) -> BasicIdentity {
+    let key = SigningKey::from(secret);
+    BasicIdentity::from_signing_key(key)
+}

src/ic_tee_agent/src/lib.rs

@@ -1,5 +1,16 @@
+use rand::thread_rng;
+use rand::RngCore;
+
 pub mod agent;
 pub mod http;
 pub mod identity;
+pub mod setting;
 
 pub use identity::*;
+
+pub fn rand_bytes<const N: usize>() -> [u8; N] {
+    let mut rng = thread_rng();
+    let mut bytes = [0u8; N];
+    rng.fill_bytes(&mut bytes);
+    bytes
+}

src/ic_tee_agent/src/setting.rs

@@ -0,0 +1,41 @@
+use ciborium::from_reader;
+use ic_cose_types::{
+    cose::{encrypt0::cose_decrypt0, get_cose_key_secret, CborSerializable, CoseKey},
+    types::setting::SettingInfo,
+};
+use serde::{Deserialize, Serialize};
+use serde_bytes::ByteBuf;
+
+pub fn decrypt_payload(info: SettingInfo, mut secret: [u8; 32]) -> Result<Vec<u8>, String> {
+    let aad: &[u8] = &[];
+    if let Some(dek) = info.dek {
+        let key = cose_decrypt0(dek.as_slice(), &secret, aad)?;
+        let key =
+            CoseKey::from_slice(&key).map_err(|err| format!("invalid COSE key: {:?}", err))?;
+        let secret2 = get_cose_key_secret(key)?;
+        if info.payload.is_none() {
+            // payload in dek
+            return Ok(secret2);
+        }
+        // get dek
+        secret = secret2.try_into().map_err(|val: Vec<u8>| {
+            format!("invalid COSE secret, expected 32 bytes, got {}", val.len())
+        })?;
+    }
+
+    match info.payload {
+        None => Err("no payload".to_string()),
+        Some(payload) => cose_decrypt0(payload.as_slice(), &secret, aad),
+    }
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct TLSPayload {
+    pub crt: ByteBuf, // PEM-encoded certificate
+    pub key: ByteBuf, // PEM-encoded private key
+}
+
+pub fn decrypt_tls(info: SettingInfo, secret: [u8; 32]) -> Result<TLSPayload, String> {
+    let data = decrypt_payload(info, secret)?;
+    from_reader(&data[..]).map_err(|err| format!("failed to decode TLS payload: {:?}", err))
+}

src/ic_tee_cdk/Cargo.toml

@@ -16,3 +16,6 @@ serde = { workspace = true }
 serde_bytes = { workspace = true }
 sha3 = { workspace = true }
 ic-canister-sig-creation = { workspace = true }
+
+[dev-dependencies]
+const-hex = { workspace = true }

src/ic_tee_cdk/src/lib.rs

@@ -68,3 +68,10 @@ pub struct TEEAttestationJSON {
     pub kind: String,
     pub document: String,
 }
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct CanisterRequest {
+    pub canister: Principal,
+    pub method: String,
+    pub params: ByteBuf,
+}

src/ic_tee_nitro_gateway/Cargo.toml

@@ -32,6 +32,7 @@ aws-nitro-enclaves-nsm-api = { workspace = true }
 log = { workspace = true }
 structured-logger = { workspace = true }
 hyper-util = { workspace = true }
+ic_cose_types = { workspace = true }
 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" }

src/ic_tee_nitro_gateway/src/handler.rs

@@ -6,7 +6,9 @@ use axum::{
 };
 use hyper_util::{client::legacy::connect::HttpConnector, rt::TokioExecutor};
 use ic_tee_agent::{agent::TEEAgent, http::ContentType};
-use ic_tee_cdk::{TEEAppInformation, TEEAppInformationJSON, TEEAttestation, TEEAttestationJSON};
+use ic_tee_cdk::{
+    CanisterRequest, TEEAppInformation, TEEAppInformationJSON, TEEAttestation, TEEAttestationJSON,
+};
 use ic_tee_nitro_attestation::AttestationRequest;
 use std::sync::Arc;
 
@@ -123,6 +125,38 @@ pub async fn post_attestation(
     }
 }
 
+pub async fn query_canister(
+    State(app): State<AppState>,
+    ct: ContentType<CanisterRequest>,
+) -> impl IntoResponse {
+    match ct {
+        ContentType::CBOR(req, _) => {
+            let res = app
+                .tee_agent
+                .query_call_raw(&req.canister, &req.method, req.params.to_vec())
+                .await;
+            ContentType::CBOR(res, None).into_response()
+        }
+        _ => StatusCode::UNSUPPORTED_MEDIA_TYPE.into_response(),
+    }
+}
+
+pub async fn update_canister(
+    State(app): State<AppState>,
+    ct: ContentType<CanisterRequest>,
+) -> impl IntoResponse {
+    match ct {
+        ContentType::CBOR(req, _) => {
+            let res = app
+                .tee_agent
+                .update_call_raw(&req.canister, &req.method, req.params.to_vec())
+                .await;
+            ContentType::CBOR(res, None).into_response()
+        }
+        _ => StatusCode::UNSUPPORTED_MEDIA_TYPE.into_response(),
+    }
+}
+
 pub async fn proxy(State(app): State<AppState>, mut req: Request) -> impl IntoResponse {
     let port = if let Some(port) = app.upstream_port {
         port