chore: add more test

src/ic_tee_agent/Cargo.toml

@@ -37,3 +37,4 @@ ic-crypto-ed25519 = { workspace = true }
 
 [dev-dependencies]
 const-hex = { workspace = true }
+structured-logger = { workspace = true }

src/ic_tee_agent/src/http/authentication.rs

@@ -121,11 +121,14 @@ impl UserSignature {
             }
             if let Some(signature) = get_data(headers, &HEADER_IC_TEE_SIGNATURE) {
                 sig.signature = signature;
-                if let Some(data) = get_data(headers, &HEADER_IC_TEE_DELEGATION) {
-                    if let Ok(delegation) = from_reader(&data[..]) {
-                        sig.delegation = delegation;
-                        return Some(sig);
+                match get_data(headers, &HEADER_IC_TEE_DELEGATION) {
+                    Some(data) => {
+                        if let Ok(delegation) = from_reader(&data[..]) {
+                            sig.delegation = delegation;
+                            return Some(sig);
+                        }
                     }
+                    None => return Some(sig),
                 }
             }
         }
@@ -135,7 +138,7 @@ impl UserSignature {
 
     /// Validation Rules
     /// - Rejects anonymous users
-    /// - Delegation chain length ≤ 10
+    /// - Delegation chain length ≤ 3
     /// - Delegations must not be expired
     /// - Signature must verify against the public key
     /// - Canister must be in delegation targets (if specified)
@@ -148,7 +151,7 @@ impl UserSignature {
             return Err(AuthenticationError::AnonymousSignatureNotAllowed);
         }
 
-        if self.delegation.len() > 10 {
+        if self.delegation.len() > 3 {
             return Err(AuthenticationError::DelegationTooLongError {
                 length: self.delegation.len(),
                 maximum: 5,
@@ -245,3 +248,64 @@ fn get_data(headers: &HeaderMap, key: &HeaderName) -> Option<Vec<u8>> {
     }
     None
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use ed25519_consensus::SigningKey;
+    use ic_agent::{identity::BasicIdentity, Identity};
+    use ic_cose_types::{cose::sha3_256, to_cbor_bytes};
+    use structured_logger::unix_ms;
+
+    #[test]
+    fn test_user_signature() {
+        let secret = [8u8; 32];
+        let sk = SigningKey::from(secret);
+        let id = BasicIdentity::from_signing_key(sk);
+        println!("id: {:?}", id.sender().unwrap().to_text());
+        // jjn6g-sh75l-r3cxb-wxrkl-frqld-6p6qq-d4ato-wske5-op7s5-n566f-bqe
+
+        let msg = b"hello world";
+        let digest = sha3_256(msg);
+        let sig = id.sign_arbitrary(digest.as_slice()).unwrap();
+        println!("{:?}", sig);
+        let mut headers = HeaderMap::new();
+        headers.insert(
+            &HEADER_IC_TEE_PUBKEY,
+            URL_SAFE_NO_PAD
+                .encode(sig.public_key.unwrap())
+                .parse()
+                .unwrap(),
+        );
+        headers.insert(
+            &HEADER_IC_TEE_CONTENT_DIGEST,
+            URL_SAFE_NO_PAD.encode(digest).parse().unwrap(),
+        );
+        headers.insert(
+            &HEADER_IC_TEE_SIGNATURE,
+            URL_SAFE_NO_PAD
+                .encode(sig.signature.unwrap())
+                .parse()
+                .unwrap(),
+        );
+        if let Some(delegations) = sig.delegations {
+            headers.insert(
+                &HEADER_IC_TEE_DELEGATION,
+                URL_SAFE_NO_PAD
+                    .encode(to_cbor_bytes(&delegations))
+                    .parse()
+                    .unwrap(),
+            );
+        }
+
+        let mut us = UserSignature::try_from(&headers).unwrap();
+        assert!(us
+            .validate_request(unix_ms(), Principal::anonymous())
+            .is_ok());
+
+        us.digest = sha3_256(b"hello world 2").to_vec();
+        assert!(us
+            .validate_request(unix_ms(), Principal::anonymous())
+            .is_err());
+    }
+}

src/ic_tee_nitro_gateway/Cargo.toml

@@ -34,3 +34,6 @@ xid = { workspace = true }
 ic_tee_cdk = { path = "../ic_tee_cdk", version = "0.2" }
 ic_tee_agent = { path = "../ic_tee_agent", version = "0.2" }
 ic_tee_nitro_attestation = { path = "../ic_tee_nitro_attestation", version = "0.2" }
+
+[dev-dependencies]
+reqwest = { workspace = true }

src/ic_tee_nitro_gateway/src/handler.rs

@@ -127,7 +127,7 @@ pub async fn get_attestation(State(app): State<AppState>, req: Request) -> impl
 }
 
 /// local_server: POST /attestation
-pub async fn post_attestation(
+pub async fn local_sign_attestation(
     State(_app): State<AppState>,
     ct: Content<AttestationRequest>,
 ) -> impl IntoResponse {
@@ -163,7 +163,7 @@ pub async fn post_attestation(
 }
 
 /// local_server: POST /canister/query
-pub async fn query_canister(
+pub async fn local_query_canister(
     State(app): State<AppState>,
     ct: Content<CanisterRequest>,
 ) -> impl IntoResponse {
@@ -183,7 +183,7 @@ pub async fn query_canister(
 }
 
 /// local_server: POST /canister/update
-pub async fn update_canister(
+pub async fn local_update_canister(
     State(app): State<AppState>,
     ct: Content<CanisterRequest>,
 ) -> impl IntoResponse {
@@ -203,7 +203,10 @@ pub async fn update_canister(
 }
 
 /// local_server: POST /keys
-pub async fn call_keys(State(app): State<AppState>, ct: Content<RPCRequest>) -> impl IntoResponse {
+pub async fn local_call_keys(
+    State(app): State<AppState>,
+    ct: Content<RPCRequest>,
+) -> impl IntoResponse {
     match ct {
         Content::CBOR(req, _) => {
             let res = handle_keys_request(&req, app.tee_agent.as_ref());
@@ -367,3 +370,140 @@ fn forbid_canister_request(req: &CanisterRequest, info: &TEEAppInformation) -> b
 
     false
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candid::{decode_args, encode_args, Principal};
+    use ic_cose::rand_bytes;
+    use ic_cose_types::{
+        cose::ecdh,
+        types::{state::StateInfo, ECDHOutput, SettingPath},
+    };
+    use ic_tee_agent::{crypto::decrypt_ecdh, http::CONTENT_TYPE_CBOR};
+    use ic_tee_cdk::CanisterResponse;
+
+    static TEE_HOST: &str = "http://127.0.0.1:8080";
+    // static TEE_ID: &str = "m6a24-ioo3h-wtn6z-rntjm-rkzgw-24nrf-2x6jb-znzpt-7uctp-akavf-yqe";
+    static COSE_CANISTER: &str = "53cyg-yyaaa-aaaap-ahpua-cai";
+
+    #[tokio::test(flavor = "current_thread")]
+    #[ignore]
+    async fn test_local_call_canister() {
+        let client = reqwest::Client::new();
+
+        // local_query_canister
+        {
+            let params = encode_args(()).unwrap();
+            let req = CanisterRequest {
+                canister: Principal::from_text(COSE_CANISTER).unwrap(),
+                method: "state_get_info".to_string(),
+                params: params.into(),
+            };
+            let res = client
+                .post(format!("{}/canister/query", TEE_HOST))
+                .header(&header::CONTENT_TYPE, CONTENT_TYPE_CBOR)
+                .body(to_cbor_bytes(&req))
+                .send()
+                .await
+                .unwrap();
+            assert!(res.status().is_success());
+            assert_eq!(
+                res.headers().get(header::CONTENT_TYPE).unwrap(),
+                CONTENT_TYPE_CBOR
+            );
+
+            let data = res.bytes().await.unwrap();
+            let res: CanisterResponse = from_reader(&data[..]).unwrap();
+            assert!(res.is_ok());
+            let res: (Result<StateInfo, String>,) = decode_args(&res.unwrap()).unwrap();
+            let res = res.0.unwrap();
+            assert_eq!(res.name, "LDC Labs");
+            assert_eq!(res.schnorr_key_name, "dfx_test_key");
+        }
+
+        // local_update_canister
+        {
+            let nonce: [u8; 12] = rand_bytes();
+            let secret: [u8; 32] = rand_bytes();
+            let secret = ecdh::StaticSecret::from(secret);
+            let public = ecdh::PublicKey::from(&secret);
+            let params = encode_args((
+                SettingPath {
+                    ns: "_".to_string(),
+                    key: "v1".as_bytes().to_vec().into(),
+                    ..Default::default()
+                },
+                ECDHInput {
+                    nonce: nonce.into(),
+                    public_key: public.to_bytes().into(),
+                },
+            ))
+            .unwrap();
+            let req = CanisterRequest {
+                canister: Principal::from_text(COSE_CANISTER).unwrap(),
+                method: "ecdh_cose_encrypted_key".to_string(),
+                params: params.into(),
+            };
+            let res = client
+                .post(format!("{}/canister/update", TEE_HOST))
+                .header(&header::CONTENT_TYPE, CONTENT_TYPE_CBOR)
+                .body(to_cbor_bytes(&req))
+                .send()
+                .await
+                .unwrap();
+            assert!(res.status().is_success());
+            assert_eq!(
+                res.headers().get(header::CONTENT_TYPE).unwrap(),
+                CONTENT_TYPE_CBOR
+            );
+
+            let data = res.bytes().await.unwrap();
+            let res: CanisterResponse = from_reader(&data[..]).unwrap();
+            assert!(res.is_ok());
+            let res: (Result<ECDHOutput<ByteBuf>, String>,) = decode_args(&res.unwrap()).unwrap();
+            assert!(res.0.is_ok());
+        }
+    }
+
+    #[tokio::test(flavor = "current_thread")]
+    #[ignore]
+    async fn test_local_call_keys() {
+        let client = reqwest::Client::new();
+
+        let nonce: [u8; 12] = rand_bytes();
+        let secret: [u8; 32] = rand_bytes();
+        let secret = ecdh::StaticSecret::from(secret);
+        let public = ecdh::PublicKey::from(&secret);
+        let params = to_cbor_bytes(&(
+            &[0u8; 0],
+            &ECDHInput {
+                nonce: nonce.into(),
+                public_key: public.to_bytes().into(),
+            },
+        ));
+        let req = RPCRequest {
+            method: "a256gcm_ecdh_key".to_string(),
+            params: params.into(),
+        };
+        let res = client
+            .post(format!("{}/keys", TEE_HOST))
+            .header(&header::CONTENT_TYPE, CONTENT_TYPE_CBOR)
+            .body(to_cbor_bytes(&req))
+            .send()
+            .await
+            .unwrap();
+        assert!(res.status().is_success());
+        assert_eq!(
+            res.headers().get(header::CONTENT_TYPE).unwrap(),
+            CONTENT_TYPE_CBOR
+        );
+
+        let data = res.bytes().await.unwrap();
+        let res: RPCResponse = from_reader(&data[..]).unwrap();
+        assert!(res.is_ok());
+        let res: ECDHOutput<ByteBuf> = from_reader(res.unwrap().as_slice()).unwrap();
+        let key = decrypt_ecdh(secret.to_bytes(), &res).unwrap();
+        assert_eq!(key.len(), 32);
+    }
+}

src/ic_tee_nitro_gateway/src/main.rs

@@ -214,10 +214,8 @@ async fn bootstrap(cli: Cli) -> Result<()> {
     let master_secret = tee_agent
         .cose_get_secret(&SettingPath {
             ns: namespace.clone(),
-            user_owned: false,
-            subject: Some(principal),
             key: COSE_SECRET_PERMANENT_KEY.as_bytes().to_vec().into(),
-            version: 0,
+            ..Default::default()
         })
         .await
         .map_err(anyhow::Error::msg)?;
@@ -301,11 +299,17 @@ async fn bootstrap(cli: Cli) -> Result<()> {
             .route("/information", routing::get(handler::get_information))
             .route(
                 "/attestation",
-                routing::get(handler::get_attestation).post(handler::post_attestation),
+                routing::get(handler::get_attestation).post(handler::local_sign_attestation),
+            )
+            .route(
+                "/canister/query",
+                routing::post(handler::local_query_canister),
+            )
+            .route(
+                "/canister/update",
+                routing::post(handler::local_update_canister),
             )
-            .route("/canister/query", routing::post(handler::query_canister))
-            .route("/canister/update", routing::post(handler::update_canister))
-            .route("/keys", routing::post(handler::call_keys))
+            .route("/keys", routing::post(handler::local_call_keys))
             .with_state(handler::AppState {
                 info: info.clone(),
                 http_client: http_client.clone(),