handle case where a node is added (#491)

dfinity · Nov 16, 2023 · a0e851f · a0e851f
1 parent 5aa3ba7
commit a0e851f
Showing 1 changed file with 37 additions and 23 deletions.
diff --git a/ic-agent/src/agent/mod.rs b/ic-agent/src/agent/mod.rs
@@ -474,7 +474,7 @@ impl Agent {
         request_id: RequestId,
     ) -> Result<Vec<u8>, AgentError> {
         let response = if self.verify_query_signatures {
-            let (response, subnet) = futures_util::try_join!(
+            let (response, mut subnet) = futures_util::try_join!(
                 self.query_endpoint::<QueryResponse>(effective_canister_id, signed_query),
                 self.get_subnet_by_canister(&effective_canister_id)
             )?;
@@ -494,10 +494,17 @@ impl Agent {
                     return Err(AgentError::CertificateOutdated(self.ingress_expiry));
                 }
                 let signable = response.signable(request_id, signature.timestamp);
-                let node_key = subnet
-                    .node_keys
-                    .get(&signature.identity)
-                    .ok_or(AgentError::CertificateNotAuthorized())?;
+                let node_key = if let Some(node_key) = subnet.node_keys.get(&signature.identity) {
+                    node_key
+                } else {
+                    subnet = self
+                        .fetch_subnet_by_canister(&effective_canister_id)
+                        .await?;
+                    subnet
+                        .node_keys
+                        .get(&signature.identity)
+                        .ok_or(AgentError::CertificateNotAuthorized())?
+                };
                 if node_key.len() != 44 {
                     return Err(AgentError::DerKeyLengthMismatch {
                         expected: 44,
@@ -999,24 +1006,31 @@ impl Agent {
         if let Some(subnet) = subnet {
             Ok(subnet)
         } else {
-            let cert = self
-                .read_state_raw(vec![vec!["subnet".into()]], *canister)
-                .await?;
-            let time = leb128::read::unsigned(&mut lookup_value(&cert.tree, [b"time".as_ref()])?)?;
-            if (OffsetDateTime::now_utc()
-                - OffsetDateTime::from_unix_timestamp_nanos(time as _).unwrap())
-                > self.ingress_expiry
-            {
-                Err(AgentError::CertificateOutdated(self.ingress_expiry))
-            } else {
-                let (subnet_id, subnet) = lookup_subnet(&cert, &self.root_key.read().unwrap())?;
-                let subnet = Arc::new(subnet);
-                self.subnet_key_cache
-                    .lock()
-                    .unwrap()
-                    .insert_subnet(subnet_id, subnet.clone());
-                Ok(subnet)
-            }
+            self.fetch_subnet_by_canister(canister).await
+        }
+    }
+
+    async fn fetch_subnet_by_canister(
+        &self,
+        canister: &Principal,
+    ) -> Result<Arc<Subnet>, AgentError> {
+        let cert = self
+            .read_state_raw(vec![vec!["subnet".into()]], *canister)
+            .await?;
+        let time = leb128::read::unsigned(&mut lookup_value(&cert.tree, [b"time".as_ref()])?)?;
+        if (OffsetDateTime::now_utc()
+            - OffsetDateTime::from_unix_timestamp_nanos(time as _).unwrap())
+            > self.ingress_expiry
+        {
+            Err(AgentError::CertificateOutdated(self.ingress_expiry))
+        } else {
+            let (subnet_id, subnet) = lookup_subnet(&cert, &self.root_key.read().unwrap())?;
+            let subnet = Arc::new(subnet);
+            self.subnet_key_cache
+                .lock()
+                .unwrap()
+                .insert_subnet(subnet_id, subnet.clone());
+            Ok(subnet)
         }
     }
 }