feat: CON-1445 Start remote DKG as soon as request appears in state

rs/consensus/dkg/src/dkg_key_manager.rs

@@ -151,7 +151,7 @@ impl DkgKeyManager {
             }
         };
 
-        if id.start_block_height > last_height && id.target_subnet == NiDkgTargetSubnet::Local {
+        if id.start_block_height > last_height && id.target_subnet.is_local() {
             info!(
                 every_n_seconds => 5,
                 self.logger,

rs/consensus/dkg/src/payload_validator.rs

@@ -1,5 +1,5 @@
-use self::payload_builder::create_early_remote_transcripts;
 use super::{crypto_validate_dealing, payload_builder, utils};
+use crate::remote::{build_callback_id_config_map, merge_configs};
 use ic_consensus_utils::{crypto::ConsensusCrypto, pool_reader::PoolReader};
 use ic_interfaces::{
     dkg::{DkgPayloadValidationError, DkgPool},
@@ -15,7 +15,10 @@ use ic_types::{
     batch::ValidationContext,
     consensus::{
         Block, BlockPayload,
-        dkg::{DkgDataPayload, DkgPayloadValidationFailure, DkgSummary, InvalidDkgPayloadReason},
+        dkg::{
+            DkgDataPayload, DkgPayloadCreationError, DkgPayloadValidationFailure, DkgSummary,
+            InvalidDkgPayloadReason,
+        },
     },
 };
 use prometheus::IntCounterVec;
@@ -117,6 +120,8 @@ pub fn validate_payload(
                 });
 
             validate_dealings_payload(
+                subnet_id,
+                registry_client,
                 crypto,
                 pool_reader,
                 dkg_pool,
@@ -134,8 +139,11 @@ pub fn validate_payload(
 }
 
 // Validates the payload containing dealings.
+#[allow(clippy::too_many_arguments)]
 #[allow(clippy::result_large_err)]
 fn validate_dealings_payload(
+    subnet_id: SubnetId,
+    registry_client: &dyn RegistryClient,
     crypto: &dyn ConsensusCrypto,
     pool_reader: &PoolReader<'_>,
     dkg_pool: &dyn DkgPool,
@@ -181,6 +189,20 @@ fn validate_dealings_payload(
         return Err(InvalidDkgPayloadReason::DealerAlreadyDealt(dealer_id).into());
     }
 
+    let state = state_reader
+        .get_state_at(validation_context.certified_height)
+        .map_err(DkgPayloadCreationError::StateManagerError)?;
+
+    let remote_config_results = build_callback_id_config_map(
+        subnet_id,
+        registry_client,
+        state.get_ref(),
+        validation_context.registry_version,
+        last_summary,
+        log,
+    )?;
+    let configs = merge_configs(&last_summary.configs, &remote_config_results);
+
     // Check that all messages have a valid DKG config from the summary and the
     // dealer is valid, then verify each dealing.
     for message in &dealings.messages {
@@ -192,7 +214,7 @@ fn validate_dealings_payload(
             continue;
         }
 
-        let Some(config) = last_summary.configs.get(&message.content.dkg_id) else {
+        let Some(config) = configs.get(&message.content.dkg_id) else {
             return Err(InvalidDkgPayloadReason::MissingDkgConfigForDealing.into());
         };
 
@@ -202,14 +224,12 @@ fn validate_dealings_payload(
 
     // If we have early transcripts, we compare them
     if !dealings.transcripts_for_remote_subnets.is_empty() {
-        let expected_transcripts = create_early_remote_transcripts(
+        let expected_transcripts = payload_builder::create_early_remote_transcripts(
             pool_reader,
             crypto,
             parent,
-            last_summary,
-            state_reader,
-            validation_context,
-            log.clone(),
+            remote_config_results,
+            log,
         )?;
 
         if dealings.transcripts_for_remote_subnets != expected_transcripts {
@@ -246,11 +266,13 @@ mod tests {
         dkg::ChangeAction,
         p2p::consensus::{MutablePool, PoolMutationsProducer},
     };
+    use ic_interfaces_state_manager::Labeled;
     use ic_logger::no_op_logger;
     use ic_metrics::MetricsRegistry;
     use ic_registry_keys::make_subnet_record_key;
     use ic_test_utilities_consensus::fake::FakeContentSigner;
     use ic_test_utilities_registry::SubnetRecordBuilder;
+    use ic_test_utilities_state::get_initial_state;
     use ic_test_utilities_types::ids::{
         NODE_1, NODE_2, NODE_3, SUBNET_1, SUBNET_2, node_test_id, subnet_test_id,
     };
@@ -707,6 +729,13 @@ mod tests {
                         .build(),
                 )],
             );
+            state_manager
+                .get_mut()
+                .expect_get_latest_certified_state()
+                .return_const(Some(Labeled::new(
+                    Height::new(0),
+                    Arc::new(get_initial_state(0, 0)),
+                )));
 
             // Both summary registry versions should be 1 initially
             let summary_block = pool.as_cache().summary_block();
@@ -758,6 +787,9 @@ mod tests {
             let key_manager = Arc::new(Mutex::new(key_manager));
             let dkg_impl = DkgImpl::new(
                 node_id,
+                subnet_id,
+                registry.clone(),
+                state_manager.clone(),
                 crypto.clone(),
                 pool.get_cache(),
                 key_manager,
@@ -778,11 +810,12 @@ mod tests {
             };
 
             // It should be possible to validate the dealing
+            let configs = dkg_summary.configs.iter().collect();
             let result = dkg_impl.validate_dealings_for_dealer(
                 &dkg_pool,
-                &dkg_summary.configs,
+                &configs,
                 start_height,
-                vec![dealing],
+                &[dealing],
             );
             let first = result.first().unwrap();
             let ChangeAction::MoveToValidated(dealing_validated) = first else {

rs/consensus/dkg/src/remote.rs

@@ -294,7 +294,6 @@ mod tests {
             vec![],
             current_transcripts,
             next_transcripts,
-            vec![],
             RegistryVersion::from(1),
             Height::from(10),
             Height::from(10),

rs/consensus/dkg/src/test_utils.rs

@@ -73,13 +73,17 @@ pub(super) fn complement_state_manager_with_dkg_contexts(
             .subnet_call_context_manager
             .push_context(context);
     }
+    let state = Arc::new(state);
     let mut mock = state_manager.get_mut();
     let expectation = mock
         .expect_get_state_at()
-        .return_const(Ok(Labeled::new(Height::new(0), Arc::new(state))));
+        .return_const(Ok(Labeled::new(Height::new(0), state.clone())));
     if let Some(times) = times {
         expectation.times(times);
     }
+
+    mock.expect_get_latest_certified_state()
+        .return_const(Some(Labeled::new(Height::new(0), state.clone())));
 }
 
 pub(super) fn complement_state_manager_with_setup_initial_dkg_request(
@@ -96,23 +100,6 @@ pub(super) fn complement_state_manager_with_setup_initial_dkg_request(
     complement_state_manager_with_dkg_contexts(state_manager, contexts, times);
 }
 
-pub(super) fn complement_state_manager_with_reshare_chain_key_request(
-    state_manager: Arc<RefMockStateManager>,
-    registry_version: RegistryVersion,
-    key_id: VetKdKeyId,
-    node_ids: Vec<u64>,
-    times: Option<usize>,
-    target: Option<NiDkgTargetId>,
-) {
-    let contexts = target
-        .into_iter()
-        .map(|t| {
-            make_reshare_chain_key_context(registry_version, key_id.clone(), node_ids.clone(), t)
-        })
-        .collect();
-    complement_state_manager_with_dkg_contexts(state_manager, contexts, times);
-}
-
 /// Extract the remote dkg transcripts from the current highest validated block
 pub(super) fn extract_remote_dkgs_from_highest_block(
     pool: &TestConsensusPool,
@@ -147,19 +134,6 @@ pub(super) fn extract_dealings_from_highest_block(pool: &TestConsensusPool) -> D
     }
 }
 
-/// Extract the remote dkg IDs from the current highest validated block
-pub(super) fn extract_remote_dkg_ids_from_highest_block(
-    pool: &TestConsensusPool,
-    target_id: NiDkgTargetId,
-) -> Vec<NiDkgId> {
-    extract_dkg_configs_from_highest_block(pool)
-        .iter()
-        .filter(|(id, _)| id.target_subnet == NiDkgTargetSubnet::Remote(target_id))
-        .map(|(id, _)| id)
-        .cloned()
-        .collect()
-}
-
 /// Extract the DKG configs from the current highest validated block
 pub(super) fn extract_dkg_configs_from_highest_block(
     pool: &TestConsensusPool,

rs/consensus/idkg/src/payload_builder.rs

@@ -781,7 +781,6 @@ mod tests {
                 vec![],
                 BTreeMap::new(),
                 BTreeMap::new(),
-                Vec::new(),
                 RegistryVersion::from(0),
                 Height::from(100),
                 Height::from(100),

rs/consensus/src/consensus/batch_delivery.rs

@@ -524,7 +524,16 @@ fn generate_dkg_response_payload(
         (Some(Err(err_str)), _) | (_, Some(Err(err_str))) => Some(Payload::Reject(
             RejectContext::new(RejectCode::CanisterReject, err_str),
         )),
-        _ => None,
+        (Some(Ok(transcript)), None) | (None, Some(Ok(transcript))) => {
+            Some(Payload::Reject(RejectContext::new(
+                RejectCode::CanisterReject,
+                format!(
+                    "Data payload contains only the {:?} transcript for SetupInitialDKG request",
+                    transcript.dkg_id.dkg_tag
+                ),
+            )))
+        }
+        (None, None) => None,
     }
 }
 
@@ -639,6 +648,28 @@ mod tests {
         };
     }
 
+    #[test]
+    fn test_generate_setup_initial_dkg_response_rejects_if_only_one_transcript_present() {
+        let transcripts_for_remote_subnets = [(
+            ni_dkg_id(NiDkgTag::LowThreshold),
+            CallbackId::from(1),
+            Ok(dummy_transcript_for_tests()),
+        )];
+
+        let result =
+            generate_responses_to_remote_dkgs(&transcripts_for_remote_subnets[..], &no_op_logger());
+        assert_eq!(result.len(), 1);
+
+        let Payload::Reject(reject) = &result[0].payload else {
+            panic!("Expected reject payload when only one SetupInitialDKG transcript exists");
+        };
+        assert_eq!(reject.code(), RejectCode::CanisterReject);
+        assert_eq!(
+            reject.message(),
+            "Data payload contains only the LowThreshold transcript for SetupInitialDKG request"
+        );
+    }
+
     #[test]
     fn test_generate_responses_for_early_remote_dkg_transcripts() {
         let key_id: NiDkgMasterPublicKeyId = NiDkgMasterPublicKeyId::VetKd(VetKdKeyId {

rs/consensus/tests/framework/runner.rs

@@ -171,6 +171,9 @@ impl<'a> ConsensusRunner<'a> {
         );
         let dkg = ic_consensus_dkg::DkgImpl::new(
             deps.replica_config.node_id,
+            deps.replica_config.subnet_id,
+            Arc::clone(&deps.registry_client),
+            deps.state_manager.clone(),
             Arc::clone(&consensus_crypto),
             deps.consensus_pool.read().unwrap().get_cache(),
             dkg_key_manager,

rs/consensus/tests/payload.rs

@@ -90,6 +90,12 @@ fn consensus_produces_expected_batches() {
                 Height::new(0),
                 Arc::new(get_initial_state(0, 0)),
             )));
+        state_manager
+            .expect_get_latest_certified_state()
+            .return_const(Some(Labeled::new(
+                Height::new(0),
+                Arc::new(get_initial_state(0, 0)),
+            )));
         state_manager
             .expect_get_certified_state_snapshot()
             .returning(|| None);
@@ -182,6 +188,9 @@ fn consensus_produces_expected_batches() {
             ic_consensus::consensus::ConsensusBouncer::new(&metrics_registry, router.clone());
         let dkg = ic_consensus_dkg::DkgImpl::new(
             replica_config.node_id,
+            replica_config.subnet_id,
+            Arc::clone(&registry_client) as Arc<_>,
+            Arc::clone(&state_manager) as Arc<_>,
             Arc::clone(&fake_crypto) as Arc<_>,
             Arc::clone(&consensus_cache),
             dkg_key_manager,

rs/replica/setup_ic_network/src/lib.rs

@@ -612,6 +612,9 @@ fn start_consensus(
         abortable_broadcast_channels.dkg,
         ic_consensus_dkg::DkgImpl::new(
             node_id,
+            subnet_id,
+            Arc::clone(&registry_client),
+            Arc::clone(&state_manager) as Arc<_>,
             Arc::clone(&consensus_crypto),
             Arc::clone(&consensus_pool_cache),
             dkg_key_manager,

rs/test_utilities/consensus/src/fake.rs

@@ -61,7 +61,6 @@ impl Fake for DkgSummary {
             /*current_transcripts=*/
             empty_ni_dkg_transcripts_with_committee(registry_version),
             /*next_transcripts=*/ BTreeMap::default(),
-            /*transcript_for_new_subnets=*/ Vec::default(),
             RegistryVersion::from(registry_version),
             /*interval_length=*/ Height::new(59),
             /*next_interval_length=*/ Height::new(59),

rs/tests/consensus/subnet_splitting_test.rs

@@ -55,7 +55,7 @@ use candid::Principal;
 use slog::{Logger, info};
 use std::{thread, time::Duration};
 
-const DKG_INTERVAL: u64 = 9;
+const DKG_INTERVAL: u64 = 29;
 const APP_NODES: usize = 1;
 
 const MESSAGE_IN_THE_CANISTER_TO_BE_MIGRATED: &str =

rs/tests/consensus/subnet_splitting_v2_test.rs

@@ -47,7 +47,6 @@ use ic_system_test_driver::{
     },
     util::runtime_from_url,
 };
-use ic_types::crypto::threshold_sig::ni_dkg::NiDkgTargetSubnet;
 use ic_types::{CanisterId, Height, NodeId, PrincipalId, RegistryVersion, SubnetId};
 use registry_canister::mutations::do_split_subnet::SplitSubnetPayload;
 use slog::info;
@@ -837,7 +836,7 @@ async fn wait_for_cup_with_subnet_id(
                     )
                 }
                 Ok(cup)
-                    if cup.signature.signer.target_subnet == NiDkgTargetSubnet::Local
+                    if cup.signature.signer.target_subnet.is_local()
                         && cup.signature.signer.dealer_subnet == subnet_id =>
                 {
                     Ok(())

rs/types/types/src/consensus/dkg.rs

@@ -186,12 +186,10 @@ pub struct DkgSummary {
 
 impl DkgSummary {
     /// Create a new Summary
-    #[allow(clippy::too_many_arguments)]
     pub fn new(
         configs: Vec<NiDkgConfig>,
         current_transcripts: BTreeMap<NiDkgTag, NiDkgTranscript>,
         next_transcripts: BTreeMap<NiDkgTag, NiDkgTranscript>,
-        transcripts_for_remote_subnets: Vec<(NiDkgId, CallbackId, Result<NiDkgTranscript, String>)>,
         registry_version: RegistryVersion,
         interval_length: Height,
         next_interval_length: Height,
@@ -205,7 +203,7 @@ impl DkgSummary {
                 .collect(),
             current_transcripts,
             next_transcripts,
-            transcripts_for_remote_subnets,
+            transcripts_for_remote_subnets: vec![],
             registry_version,
             interval_length,
             next_interval_length,