feat!: add support for fetching root keys automatically

CHANGELOG.md

@@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 * Removed the Bitcoin query methods from `ManagementCanister`. Users should use `BitcoinCanister` for that.
 * Added `BitcoinCanister` to `ic-utils`.
+* Added support for automatically fetching the root key in `ic-agent`
+* `Agent::verify` and `Agent::verify_for_subnet` are now `async`.
 
 ## [0.36.0] - 2024-06-04
 

ic-agent/src/agent/agent_config.rs

@@ -18,6 +18,8 @@ pub struct AgentConfig {
     pub verify_query_signatures: bool,
     /// See [`with_max_concurrent_requests`](super::AgentBuilder::with_max_concurrent_requests).
     pub max_concurrent_requests: usize,
+    /// See [`with_auto_fetch_root_key`](super::AgentBuilder::with_auto_fetch_root_key).
+    pub auto_fetch_root_key: bool,
 }
 
 impl Default for AgentConfig {
@@ -29,6 +31,7 @@ impl Default for AgentConfig {
             transport: None,
             verify_query_signatures: true,
             max_concurrent_requests: 50,
+            auto_fetch_root_key: false,
         }
     }
 }

ic-agent/src/agent/agent_test.rs

@@ -659,7 +659,7 @@ async fn too_many_delegations() {
         .expect("read state failed");
     let new_cert = self_delegate_cert(subnet_id, &cert, 1);
     assert!(matches!(
-        agent.verify(&new_cert, canister_id).unwrap_err(),
+        agent.verify(&new_cert, canister_id).await.unwrap_err(),
         AgentError::CertificateHasTooManyDelegations
     ));
 }

ic-agent/src/agent/builder.rs

@@ -106,4 +106,18 @@ impl AgentBuilder {
         self.config.max_concurrent_requests = max_concurrent_requests;
         self
     }
+
+    /// By default, the agent is configured to talk to the main Internet Computer, and verifies
+    /// responses using a hard-coded public key.
+    ///
+    /// This flag will instruct the agent to ask the endpoint for its public key automatically on
+    /// first request and use that instead. This is required when talking to a local test instance,
+    /// for example.
+    ///
+    /// *Only use this when you are  _not_ talking to the main Internet Computer, otherwise
+    /// you are prone to man-in-the-middle attacks! Do not enable this flag by default.*
+    pub fn with_auto_fetch_root_key(mut self, auto_fetch_root_key: bool) -> Self {
+        self.config.auto_fetch_root_key = auto_fetch_root_key;
+        self
+    }
 }

ic-agent/src/agent/mod.rs

@@ -50,7 +50,10 @@ use std::{
     fmt,
     future::{Future, IntoFuture},
     pin::Pin,
-    sync::{Arc, Mutex, RwLock},
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc, Mutex, RwLock,
+    },
     task::{Context, Poll},
     time::Duration,
 };
@@ -218,13 +221,13 @@ pub enum PollResult {
 ///   let agent = Agent::builder()
 ///     .with_url(url)
 ///     .with_identity(create_identity())
+///     // Only enable the following flag when not contacting the IC main net (e.g. a local emulator).
+///     // This is important as the main net public key is static and a rogue network could return
+///     // a different key.
+///     // If you know the root key ahead of time, you can use `agent.set_root_key(root_key);`.
+///     .with_auto_fetch_root_key(true)
 ///     .build()?;
 ///
-///   // Only do the following call when not contacting the IC main net (e.g. a local emulator).
-///   // This is important as the main net public key is static and a rogue network could return
-///   // a different key.
-///   // If you know the root key ahead of time, you can use `agent.set_root_key(root_key);`.
-///   agent.fetch_root_key().await?;
 ///   let management_canister_id = Principal::from_text("aaaaa-aa")?;
 ///
 ///   // Create a call to the management canister to create a new canister ID,
@@ -259,6 +262,8 @@ pub struct Agent {
     subnet_key_cache: Arc<Mutex<SubnetCache>>,
     concurrent_requests_semaphore: Arc<Semaphore>,
     verify_query_signatures: bool,
+    fetch_root_key_fuse: Arc<AtomicBool>,
+    auto_fetch_root_key: bool,
 }
 
 impl fmt::Debug for Agent {
@@ -289,6 +294,8 @@ impl Agent {
             subnet_key_cache: Arc::new(Mutex::new(SubnetCache::new())),
             verify_query_signatures: config.verify_query_signatures,
             concurrent_requests_semaphore: Arc::new(Semaphore::new(config.max_concurrent_requests)),
+            fetch_root_key_fuse: Arc::new(AtomicBool::new(false)),
+            auto_fetch_root_key: config.auto_fetch_root_key,
         })
     }
 
@@ -326,7 +333,7 @@ impl Agent {
     /// *Only use this when you are  _not_ talking to the main Internet Computer, otherwise
     /// you are prone to man-in-the-middle attacks! Do not call this function by default.*
     pub async fn fetch_root_key(&self) -> Result<(), AgentError> {
-        if self.read_root_key()[..] != IC_ROOT_KEY[..] {
+        if self.fetch_root_key_fuse.load(Ordering::Acquire) {
             // already fetched the root key
             return Ok(());
         }
@@ -336,6 +343,7 @@ impl Agent {
             None => return Err(AgentError::NoRootKeyInStatus(status)),
         };
         self.set_root_key(root_key);
+        self.fetch_root_key_fuse.store(true, Ordering::Release);
         Ok(())
     }
 
@@ -348,6 +356,18 @@ impl Agent {
     }
 
     /// Return the root key currently in use.
+    /// this may perform a network request if automatic fetching
+    /// of the root key is enabled.
+    pub async fn read_or_fetch_root_key(&self) -> Result<Vec<u8>, AgentError> {
+        if self.auto_fetch_root_key {
+            self.fetch_root_key().await?;
+        }
+        Ok(self.read_root_key())
+    }
+
+    /// Return the root key currently stored in the agent.
+    /// This key may change if the agent is configured to automatically
+    /// fetch the root key upon the first request.
     pub fn read_root_key(&self) -> Vec<u8> {
         self.root_key.read().unwrap().clone()
     }
@@ -778,7 +798,7 @@ impl Agent {
             .await?;
         let cert: Certificate = serde_cbor::from_slice(&read_state_response.certificate)
             .map_err(AgentError::InvalidCborData)?;
-        self.verify(&cert, effective_canister_id)?;
+        self.verify(&cert, effective_canister_id).await?;
         Ok(cert)
     }
 
@@ -797,7 +817,7 @@ impl Agent {
             .await?;
         let cert: Certificate = serde_cbor::from_slice(&read_state_response.certificate)
             .map_err(AgentError::InvalidCborData)?;
-        self.verify_for_subnet(&cert, subnet_id)?;
+        self.verify_for_subnet(&cert, subnet_id).await?;
         Ok(cert)
     }
 
@@ -811,12 +831,13 @@ impl Agent {
 
     /// Verify a certificate, checking delegation if present.
     /// Only passes if the certificate also has authority over the canister.
-    pub fn verify(
+    pub async fn verify(
         &self,
         cert: &Certificate,
         effective_canister_id: Principal,
     ) -> Result<(), AgentError> {
-        self.verify_cert(cert, effective_canister_id)?;
+        let root_key = self.read_or_fetch_root_key().await?;
+        self.verify_cert(cert, effective_canister_id, root_key)?;
         self.verify_cert_timestamp(cert)?;
         Ok(())
     }
@@ -825,6 +846,7 @@ impl Agent {
         &self,
         cert: &Certificate,
         effective_canister_id: Principal,
+        root_key: Vec<u8>,
     ) -> Result<(), AgentError> {
         let sig = &cert.signature;
 
@@ -833,7 +855,7 @@ impl Agent {
         msg.extend_from_slice(IC_STATE_ROOT_DOMAIN_SEPARATOR);
         msg.extend_from_slice(&root_hash);
 
-        let der_key = self.check_delegation(&cert.delegation, effective_canister_id)?;
+        let der_key = self.check_delegation(&cert.delegation, effective_canister_id, root_key)?;
         let key = extract_der(der_key)?;
 
         ic_verify_bls_signature::verify_bls_signature(sig, &msg, &key)
@@ -843,12 +865,13 @@ impl Agent {
 
     /// Verify a certificate, checking delegation if present.
     /// Only passes if the certificate is for the specified subnet.
-    pub fn verify_for_subnet(
+    pub async fn verify_for_subnet(
         &self,
         cert: &Certificate,
         subnet_id: Principal,
     ) -> Result<(), AgentError> {
-        self.verify_cert_for_subnet(cert, subnet_id)?;
+        let root_key = self.read_or_fetch_root_key().await?;
+        self.verify_cert_for_subnet(cert, subnet_id, root_key)?;
         self.verify_cert_timestamp(cert)?;
         Ok(())
     }
@@ -857,6 +880,7 @@ impl Agent {
         &self,
         cert: &Certificate,
         subnet_id: Principal,
+        root_key: Vec<u8>,
     ) -> Result<(), AgentError> {
         let sig = &cert.signature;
 
@@ -865,7 +889,7 @@ impl Agent {
         msg.extend_from_slice(IC_STATE_ROOT_DOMAIN_SEPARATOR);
         msg.extend_from_slice(&root_hash);
 
-        let der_key = self.check_delegation_for_subnet(&cert.delegation, subnet_id)?;
+        let der_key = self.check_delegation_for_subnet(&cert.delegation, subnet_id, root_key)?;
         let key = extract_der(der_key)?;
 
         ic_verify_bls_signature::verify_bls_signature(sig, &msg, &key)
@@ -890,16 +914,17 @@ impl Agent {
         &self,
         delegation: &Option<Delegation>,
         effective_canister_id: Principal,
+        root_key: Vec<u8>,
     ) -> Result<Vec<u8>, AgentError> {
         match delegation {
-            None => Ok(self.read_root_key()),
+            None => Ok(root_key),
             Some(delegation) => {
                 let cert: Certificate = serde_cbor::from_slice(&delegation.certificate)
                     .map_err(AgentError::InvalidCborData)?;
                 if cert.delegation.is_some() {
                     return Err(AgentError::CertificateHasTooManyDelegations);
                 }
-                self.verify_cert(&cert, effective_canister_id)?;
+                self.verify_cert(&cert, effective_canister_id, root_key)?;
                 let canister_range_lookup = [
                     "subnet".as_bytes(),
                     delegation.subnet_id.as_ref(),
@@ -927,16 +952,17 @@ impl Agent {
         &self,
         delegation: &Option<Delegation>,
         subnet_id: Principal,
+        root_key: Vec<u8>,
     ) -> Result<Vec<u8>, AgentError> {
         match delegation {
-            None => Ok(self.read_root_key()),
+            None => Ok(root_key),
             Some(delegation) => {
                 let cert: Certificate = serde_cbor::from_slice(&delegation.certificate)
                     .map_err(AgentError::InvalidCborData)?;
                 if cert.delegation.is_some() {
                     return Err(AgentError::CertificateHasTooManyDelegations);
                 }
-                self.verify_cert_for_subnet(&cert, subnet_id)?;
+                self.verify_cert_for_subnet(&cert, subnet_id, root_key)?;
                 let public_key_path = [
                     "subnet".as_bytes(),
                     delegation.subnet_id.as_ref(),
@@ -1031,7 +1057,7 @@ impl Agent {
 
         let cert: Certificate = serde_cbor::from_slice(&read_state_response.certificate)
             .map_err(AgentError::InvalidCborData)?;
-        self.verify(&cert, effective_canister_id)?;
+        self.verify(&cert, effective_canister_id).await?;
         lookup_request_status(cert, request_id)
     }
 

ref-tests/src/utils.rs

@@ -102,6 +102,7 @@ pub async fn create_agent(identity: impl Identity + 'static) -> Result<Agent, St
     Agent::builder()
         .with_transport(ReqwestTransport::create(format!("http://127.0.0.1:{}", port)).unwrap())
         .with_identity(identity)
+        .with_auto_fetch_root_key(true)
         .build()
         .map_err(|e| format!("{:?}", e))
 }
@@ -126,10 +127,6 @@ where
         let agent = create_agent(agent_identity)
             .await
             .expect("Could not create an agent.");
-        agent
-            .fetch_root_key()
-            .await
-            .expect("could not fetch root key");
         match f(agent).await {
             Ok(_) => {}
             Err(e) => panic!("{:?}", e),

ref-tests/tests/ic-ref.rs

@@ -148,7 +148,6 @@ mod management_canister {
             let other_agent_identity = create_basic_identity()?;
             let other_agent_principal = other_agent_identity.sender()?;
             let other_agent = create_agent(other_agent_identity).await?;
-            other_agent.fetch_root_key().await?;
             let other_ic00 = ManagementCanister::create(&other_agent);
 
             // Reinstall with another agent should fail.
@@ -265,19 +264,16 @@ mod management_canister {
             let other_agent_identity = create_basic_identity()?;
             let other_agent_principal = other_agent_identity.sender()?;
             let other_agent = create_agent(other_agent_identity).await?;
-            other_agent.fetch_root_key().await?;
             let other_ic00 = ManagementCanister::create(&other_agent);
 
             let secp256k1_identity = create_secp256k1_identity()?;
             let secp256k1_principal = secp256k1_identity.sender()?;
             let secp256k1_agent = create_agent(secp256k1_identity).await?;
-            secp256k1_agent.fetch_root_key().await?;
             let secp256k1_ic00 = ManagementCanister::create(&secp256k1_agent);
 
             let prime256v1_identity = create_prime256v1_identity()?;
             let prime256v1_principal = prime256v1_identity.sender()?;
             let prime256v1_agent = create_agent(prime256v1_identity).await?;
-            prime256v1_agent.fetch_root_key().await?;
             let prime256v1_ic00 = ManagementCanister::create(&prime256v1_agent);
 
             let ic00 = ManagementCanister::create(&agent);
@@ -625,7 +621,6 @@ mod management_canister {
             // Create another agent with different identity.
             let other_agent_identity = create_basic_identity()?;
             let other_agent = create_agent(other_agent_identity).await?;
-            other_agent.fetch_root_key().await?;
             let other_ic00 = ManagementCanister::create(&other_agent);
 
             // Start as a wrong controller should fail.