Pierugo/delegation manager/refactor test

Cargo.toml

@@ -153,6 +153,7 @@ members = [
     "rs/http_endpoints/async_utils",
     "rs/http_endpoints/metrics",
     "rs/http_endpoints/nns_delegation_manager",
+    "rs/http_endpoints/nns_delegation_reader",
     "rs/http_endpoints/public",
     "rs/http_endpoints/test_agent",
     "rs/http_endpoints/xnet",

rs/execution_environment/src/query_handler.rs

@@ -22,8 +22,9 @@ use ic_crypto_tree_hash::{Label, LabeledTree, LabeledTree::SubTree, flatmap};
 use ic_cycles_account_manager::CyclesAccountManager;
 use ic_error_types::{ErrorCode, UserError};
 use ic_interfaces::execution_environment::{
-    QueryExecutionError, QueryExecutionInput, QueryExecutionResponse, QueryExecutionService,
-    TransformExecutionInput, TransformExecutionService,
+    CanisterRangesCheck, CanisterRangesFilter, NNSDelegationBuilder, QueryExecutionError,
+    QueryExecutionInput, QueryExecutionResponse, QueryExecutionService, TransformExecutionInput,
+    TransformExecutionService,
 };
 use ic_interfaces_state_manager::{Labeled, StateReader};
 use ic_logger::ReplicaLogger;
@@ -37,7 +38,7 @@ use ic_types::messages::CertificateDelegationMetadata;
 use ic_types::{
     CanisterId, NumInstructions,
     ingress::WasmResult,
-    messages::{Blob, Certificate, CertificateDelegation, Query},
+    messages::{Blob, Certificate, Query},
 };
 use prometheus::{Histogram, histogram_opts, labels};
 use serde::Serialize;
@@ -73,11 +74,29 @@ fn into_cbor<R: Serialize>(r: &R) -> Vec<u8> {
     ser.into_inner()
 }
 
+/// The result of [`get_latest_certified_state_and_data_certificate`]: the certified
+/// state together with the data certificate built from it (with the NNS delegation
+/// embedded) and the metadata of the embedded delegation, if any.
+struct CertifiedStateWithDataCertificate {
+    state: Labeled<Arc<ReplicatedState>>,
+    data_certificate: Vec<u8>,
+    certificate_delegation_metadata: Option<CertificateDelegationMetadata>,
+}
+
+/// Reads the latest certified state and builds the data certificate for the canister
+/// from it, with the NNS delegation (built with `canister_ranges_filter`) embedded.
+/// Also returns the metadata of the embedded delegation, if any.
+///
+/// The delegation is only embedded after having been verified (according to
+/// `canister_ranges_check`) to be consistent with the exact certified state which the
+/// data certificate is built from.
 fn get_latest_certified_state_and_data_certificate(
     state_reader: Arc<dyn StateReader<State = ReplicatedState>>,
-    certificate_delegation: Option<CertificateDelegation>,
+    nns_delegation_builder: Option<Arc<NNSDelegationBuilder>>,
+    canister_ranges_filter: CanisterRangesFilter,
+    canister_ranges_check: CanisterRangesCheck,
     canister_id: CanisterId,
-) -> Option<(Labeled<Arc<ReplicatedState>>, Vec<u8>)> {
+) -> Result<CertifiedStateWithDataCertificate, QueryExecutionError> {
     // The path to fetch the data certificate for the canister.
     let path = SubTree(flatmap! {
         label("canister") => SubTree(
@@ -90,18 +109,44 @@ fn get_latest_certified_state_and_data_certificate(
         label("time") => LabeledTree::Leaf(())
     });
 
-    state_reader
-        .read_certified_state(&path)
-        .map(|(state, tree, cert)| {
-            (
-                Labeled::new(cert.height, state),
-                into_cbor(&Certificate {
-                    tree,
-                    signature: Blob(cert.signed.signature.signature.get().0),
-                    delegation: certificate_delegation,
-                }),
-            )
-        })
+    let Some((state, tree, cert)) = state_reader.read_certified_state(&path) else {
+        return Err(QueryExecutionError::CertifiedStateUnavailable);
+    };
+
+    // Only embed a delegation which is consistent with the certified state which the
+    // data certificate is built from.
+    let (certificate_delegation, certificate_delegation_metadata) = match nns_delegation_builder {
+        None => (None, None),
+        Some(builder) => {
+            let network_topology = &state.metadata.network_topology;
+            let (delegation, metadata) = builder
+                .build_verified(canister_ranges_filter, canister_ranges_check, |subnet_id| {
+                    network_topology
+                        .subnets_for_certification()
+                        .get(&subnet_id)
+                        .map(|subnet_topology| {
+                            (
+                                subnet_topology.public_key.clone(),
+                                network_topology
+                                    .routing_table_for_certification()
+                                    .ranges(subnet_id),
+                            )
+                        })
+                })
+                .map_err(|_err| QueryExecutionError::DelegationInconsistentWithState)?;
+            (Some(delegation), Some(metadata))
+        }
+    };
+
+    Ok(CertifiedStateWithDataCertificate {
+        state: Labeled::new(cert.height, state),
+        data_certificate: into_cbor(&Certificate {
+            tree,
+            signature: Blob(cert.signed.signature.signature.get().0),
+            delegation: certificate_delegation,
+        }),
+        certificate_delegation_metadata,
+    })
 }
 
 fn label<T: Into<Label>>(t: T) -> Label {
@@ -466,7 +511,9 @@ impl Service<QueryExecutionInput> for HttpQueryHandler {
         &mut self,
         QueryExecutionInput {
             query,
-            certificate_delegation_with_metadata,
+            nns_delegation_builder,
+            canister_ranges_filter,
+            canister_ranges_check,
         }: QueryExecutionInput,
     ) -> Self::Future {
         let internal = Arc::clone(&self.internal);
@@ -487,18 +534,18 @@ impl Service<QueryExecutionInput> for HttpQueryHandler {
                 // Otherwise, retrieving the state in the Query service in `http_endpoints` can lead to queries being queued up,
                 // with a reference to older states which can cause out-of-memory crashes.
 
-                let (certificate_delegation, certificate_delegation_metadata) =
-                    match certificate_delegation_with_metadata {
-                        Some((delegation, metadata)) => (Some(delegation), Some(metadata)),
-                        None => (None, None),
-                    };
-
                 let result = match get_latest_certified_state_and_data_certificate(
                     state_reader,
-                    certificate_delegation,
+                    nns_delegation_builder,
+                    canister_ranges_filter,
+                    canister_ranges_check,
                     query.receiver,
                 ) {
-                    Some((state, cert)) => {
+                    Ok(CertifiedStateWithDataCertificate {
+                        state,
+                        data_certificate,
+                        certificate_delegation_metadata,
+                    }) => {
                         let time = state.get_ref().metadata.batch_time;
 
                         let certified_height_used_for_execution = state.height();
@@ -511,7 +558,7 @@ impl Service<QueryExecutionInput> for HttpQueryHandler {
 
                         let data_certificate_with_delegation_metadata =
                             DataCertificateWithDelegationMetadata {
-                                data_certificate: cert,
+                                data_certificate,
                                 certificate_delegation_metadata,
                             };
 
@@ -526,7 +573,7 @@ impl Service<QueryExecutionInput> for HttpQueryHandler {
 
                         Ok((response, time))
                     }
-                    None => Err(QueryExecutionError::CertifiedStateUnavailable),
+                    Err(err) => Err(err),
                 };
 
                 let _ = tx.send(Ok(result));

rs/http_endpoints/nns_delegation_manager/BUILD.bazel

@@ -1,5 +1,4 @@
 load("@rules_rust//rust:defs.bzl", "rust_library", "rust_test")
-load("//bazel:defs.bzl", "rust_ic_bench")
 
 package(default_visibility = [
     "//rs/http_endpoints:__subpackages__",
@@ -19,11 +18,14 @@ rust_library(
         "//rs/crypto/tls_interfaces",
         "//rs/crypto/tree_hash",
         "//rs/crypto/utils/threshold_sig_der",
+        "//rs/http_endpoints/nns_delegation_reader",
         "//rs/interfaces/registry",
+        "//rs/interfaces/state_manager",
         "//rs/monitoring/logger",
         "//rs/monitoring/metrics",
         "//rs/registry/helpers",
         "//rs/registry/subnet_type",
+        "//rs/replicated_state",
         "//rs/types/types",
         "@crate_index//:axum",
         "@crate_index//:futures",
@@ -34,7 +36,6 @@ rust_library(
         "@crate_index//:prometheus",
         "@crate_index//:rand",
         "@crate_index//:rustls",
-        "@crate_index//:serde",
         "@crate_index//:serde_cbor",
         "@crate_index//:slog",
         "@crate_index//:tokio",
@@ -63,13 +64,14 @@ rust_test(
         "//rs/crypto/tls_interfaces/mocks",
         "//rs/crypto/tree_hash",
         "//rs/crypto/utils/threshold_sig_der",
-        "//rs/http_endpoints/nns_delegation_manager/test_utils",
+        "//rs/interfaces/state_manager/mocks",
         "//rs/monitoring/logger",
         "//rs/monitoring/metrics",
         "//rs/protobuf",
         "//rs/registry/fake",
         "//rs/registry/keys",
         "//rs/registry/proto_data_provider",
+        "//rs/registry/routing_table",
         "//rs/test_utilities/registry",
         "//rs/test_utilities/types",
         "@crate_index//:assert_matches",
@@ -78,25 +80,7 @@ rust_test(
         "@crate_index//:hyper",
         "@crate_index//:rand",
         "@crate_index//:rcgen",
-        "@crate_index//:tokio",
-    ],
-)
-
-rust_ic_bench(
-    name = "nns_delegation_reader_bench",
-    testonly = True,
-    srcs = ["benches/nns_delegation_reader.rs"],
-    deps = [
-        # Keep sorted.
-        ":nns_delegation_manager",
-        "//rs/crypto/tree_hash",
-        "//rs/http_endpoints/nns_delegation_manager/test_utils",
-        "//rs/monitoring/logger",
-        "//rs/test_utilities/types",
-        "//rs/types/types",
-        "@crate_index//:criterion",
-        "@crate_index//:pprof",
-        "@crate_index//:serde_cbor",
+        "@crate_index//:rstest",
         "@crate_index//:tokio",
     ],
 )

rs/http_endpoints/nns_delegation_manager/Cargo.toml

@@ -6,10 +6,6 @@ description.workspace = true
 documentation.workspace = true
 edition.workspace = true
 
-[[bench]]
-name = "nns_delegation_reader"
-harness = false
-
 [dependencies]
 axum = { workspace = true }
 futures = { workspace = true }
@@ -23,15 +19,17 @@ ic-crypto-tls-interfaces = { path = "../../crypto/tls_interfaces" }
 ic-crypto-tree-hash = { path = "../../crypto/tree_hash" }
 ic-crypto-utils-threshold-sig-der = { path = "../../crypto/utils/threshold_sig_der" }
 ic-interfaces-registry = { path = "../../interfaces/registry" }
+ic-interfaces-state-manager = { path = "../../interfaces/state_manager" }
 ic-logger = { path = "../../monitoring/logger" }
 ic-metrics = { path = "../../monitoring/metrics" }
+ic-nns-delegation-reader = { path = "../nns_delegation_reader" }
 ic-registry-client-helpers = { path = "../../registry/helpers" }
 ic-registry-subnet-type = { path = "../../registry/subnet_type" }
+ic-replicated-state = { path = "../../replicated_state" }
 ic-types = { path = "../../types/types" }
 prometheus = { workspace = true }
 rand = { workspace = true }
 rustls = { workspace = true }
-serde = { workspace = true }
 serde_cbor = { workspace = true }
 slog = { workspace = true }
 tokio = { workspace = true }
@@ -43,15 +41,15 @@ webpki-roots = { workspace = true }
 [dev-dependencies]
 assert_matches = { workspace = true }
 axum-server = { workspace = true }
-criterion = { workspace = true }
 ic-certification-test-utils = { path = "../../certification/test-utils" }
 ic-crypto-tls-interfaces-mocks = { path = "../../crypto/tls_interfaces/mocks" }
-ic-nns-delegation-manager-test-utils = { path = "test_utils" }
+ic-interfaces-state-manager-mocks = { path = "../../interfaces/state_manager/mocks" }
 ic-protobuf = { path = "../../protobuf" }
 ic-registry-client-fake = { path = "../../registry/fake" }
 ic-registry-keys = { path = "../../registry/keys" }
 ic-registry-proto-data-provider = { path = "../../registry/proto_data_provider" }
+ic-registry-routing-table = { path = "../../registry/routing_table" }
 ic-test-utilities-registry = { path = "../../test_utilities/registry" }
 ic-test-utilities-types = { path = "../../test_utilities/types" }
-pprof = { workspace = true }
 rcgen = { workspace = true }
+rstest = { workspace = true }