From edcb5ab41a0faa2801b93a9560b6f7f9a58b033d Mon Sep 17 00:00:00 2001 From: Leo Eichhorn <99166915+eichhorl@users.noreply.github.com> Date: Tue, 18 Feb 2025 13:20:20 +0100 Subject: [PATCH] test: CON-1422 vetKD payload builder/verifier unit tests (#3886) This PR adds unit tests for the new vetKD payload builder. Some unit tests are still annotated with `#[should_panic(expected = "not yet implemented")]`, as the current implementation doesn't call the correct crypto endpoints to create and validate combined vetKD shares yet. This will be added in a subsequent PR, at which point the annotation will be removed. --- rs/consensus/vetkd/BUILD.bazel | 1 + rs/consensus/vetkd/Cargo.toml | 1 + rs/consensus/vetkd/src/lib.rs | 603 ++++++++++++++++++++++++++- rs/consensus/vetkd/src/test_utils.rs | 212 ++++++++++ rs/consensus/vetkd/src/utils.rs | 25 ++ rs/types/types/src/batch/vetkd.rs | 5 +- 6 files changed, 844 insertions(+), 3 deletions(-) create mode 100644 rs/consensus/vetkd/src/test_utils.rs diff --git a/rs/consensus/vetkd/BUILD.bazel b/rs/consensus/vetkd/BUILD.bazel index bb1719fd620..03d8104c28d 100644 --- a/rs/consensus/vetkd/BUILD.bazel +++ b/rs/consensus/vetkd/BUILD.bazel @@ -28,6 +28,7 @@ DEV_DEPENDENCIES = [ "//rs/artifact_pool", "//rs/consensus/mocks", "//rs/registry/fake", + "//rs/registry/keys", "//rs/test_utilities", "//rs/test_utilities/registry", "//rs/test_utilities/state", diff --git a/rs/consensus/vetkd/Cargo.toml b/rs/consensus/vetkd/Cargo.toml index 1f5a33b3542..1ad0c148579 100644 --- a/rs/consensus/vetkd/Cargo.toml +++ b/rs/consensus/vetkd/Cargo.toml @@ -30,6 +30,7 @@ assert_matches = { workspace = true } ic-artifact-pool = { path = "../artifact_pool" } ic-consensus-mocks = { path = "./mocks" } ic-registry-client-fake = { path = "../registry/fake" } +ic-registry-keys = { path = "../registry/keys" } ic-test-utilities = { path = "../test_utilities" } ic-test-utilities-registry = { path = "../test_utilities/registry" } ic-test-utilities-state = { path = "../test_utilities/state" } diff --git a/rs/consensus/vetkd/src/lib.rs b/rs/consensus/vetkd/src/lib.rs index f767961b016..6c0323ad9db 100644 --- a/rs/consensus/vetkd/src/lib.rs +++ b/rs/consensus/vetkd/src/lib.rs @@ -45,6 +45,8 @@ use std::{ sync::{Arc, RwLock}, }; +#[cfg(test)] +mod test_utils; mod utils; /// Implementation of the [`BatchPayloadBuilder`] for the VetKd feature. @@ -399,7 +401,7 @@ impl IntoMessages> for VetKdPayloadBuilderImpl { ), VetKdErrorCode::InvalidKey => RejectContext::new( RejectCode::CanisterError, - "Invalid key_id in VetKD request", + "Invalid or disabled key_id in VetKD request", ), }) } @@ -426,3 +428,602 @@ fn reject_if_invalid( None } } + +#[cfg(test)] +mod tests { + use assert_matches::assert_matches; + use core::{convert::From, iter::Iterator, time::Duration}; + use ic_consensus_mocks::{ + dependencies_with_subnet_records_with_raw_state_manager, Dependencies, + }; + use ic_interfaces::consensus::{InvalidPayloadReason, PayloadValidationFailure}; + use ic_interfaces::idkg::IDkgChangeAction; + use ic_interfaces::p2p::consensus::MutablePool; + use ic_interfaces::validation::ValidationError; + use ic_interfaces_state_manager::StateManagerError; + use ic_logger::no_op_logger; + use ic_management_canister_types_private::VetKdKeyId; + use ic_registry_subnet_features::ChainKeyConfig; + use ic_registry_subnet_features::KeyConfig; + use ic_test_utilities_registry::SubnetRecordBuilder; + use ic_types::consensus::idkg::IDkgMessage; + use ic_types::subnet_id_into_protobuf; + use ic_types::time::current_time; + use ic_types::time::UNIX_EPOCH; + use ic_types::RegistryVersion; + use ic_types_test_utils::ids::{node_test_id, subnet_test_id}; + use std::str::FromStr; + + use super::*; + use crate::test_utils::*; + + /// The height of the payload to be tested + const HEIGHT: Height = Height::new(124); + + /// The certified state height to be referenced + const CERTIFIED_HEIGHT: Height = Height::new(123); + + /// The validation context to be used during tests + const VALIDATION_CONTEXT: ValidationContext = ValidationContext { + registry_version: RegistryVersion::new(10), + certified_height: CERTIFIED_HEIGHT, + time: UNIX_EPOCH, + }; + + /// The maximum payload size during tests + const MAX_SIZE: NumBytes = NumBytes::new(1024); + + #[test] + fn test_into_messages() { + let agreements = make_vetkd_agreements(0, 1, 2); + let payload = as_bytes(agreements.clone()); + let messages = VetKdPayloadBuilderImpl::into_messages(&payload); + for i in 0..3 { + let id = CallbackId::from(i); + let agreement = agreements.get(&id).unwrap(); + let response = &messages[i as usize]; + assert_eq!(id, response.callback); + match agreement { + VetKdAgreement::Reject(VetKdErrorCode::InvalidKey) => { + let ResponsePayload::Reject(context) = &response.payload else { + panic!("Unexpected response: {response:?}"); + }; + context.assert_contains( + RejectCode::CanisterError, + "Invalid or disabled key_id in VetKD request", + ); + } + VetKdAgreement::Reject(VetKdErrorCode::TimedOut) => { + let ResponsePayload::Reject(context) = &response.payload else { + panic!("Unexpected response: {response:?}"); + }; + context.assert_contains(RejectCode::CanisterError, "VetKD request expired"); + } + VetKdAgreement::Success(data) => { + let ResponsePayload::Data(response_data) = &response.payload else { + panic!("Unexpected response: {response:?}"); + }; + assert_eq!(data, response_data); + } + } + } + } + + /// Run the given function for a payload builder that was setup using the given + /// config, request contexts, and message shares. + fn test_payload_builder( + config: Option, + contexts: BTreeMap, + shares: Vec, + run: impl FnOnce(VetKdPayloadBuilderImpl) -> T, + ) -> T { + test_payload_builder_ext(config, true, contexts, shares, run) + } + + fn test_payload_builder_ext( + config: Option, + keys_enabled: bool, + contexts: BTreeMap, + shares: Vec, + run: impl FnOnce(VetKdPayloadBuilderImpl) -> T, + ) -> T { + let committee = (0..4).map(|id| node_test_id(id as u64)).collect::>(); + ic_test_utilities::artifact_pool_config::with_test_pool_config(|pool_config| { + // Add the config to registry + let subnet_record_builder = SubnetRecordBuilder::from(&committee); + let subnet_record_builder = if let Some(config) = config.clone() { + subnet_record_builder.with_chain_key_config(config) + } else { + subnet_record_builder + }; + let subnet_id = subnet_test_id(0); + + let Dependencies { + crypto, + mut pool, + idkg_pool, + state_manager, + registry, + registry_data_provider, + .. + } = dependencies_with_subnet_records_with_raw_state_manager( + pool_config, + subnet_id, + vec![(1, subnet_record_builder.build())], + ); + + // Enable the configured keys + if let Some(config) = config { + if keys_enabled { + for key_id in config.key_ids() { + registry_data_provider + .add( + &ic_registry_keys::make_chain_key_signing_subnet_list_key(&key_id), + registry.get_latest_version().increment(), + Some( + ic_protobuf::registry::crypto::v1::ChainKeySigningSubnetList { + subnets: vec![subnet_id_into_protobuf(subnet_test_id(0))], + }, + ), + ) + .expect("Could not add chain key signing subnet list"); + } + registry.update_to_latest_version(); + } + } + + // Setup the state manager expectation + let mut state = ic_test_utilities_state::get_initial_state(0, 0); + state + .metadata + .subnet_call_context_manager + .sign_with_threshold_contexts = contexts; + + // We will not return states above the certified height + state_manager + .get_mut() + .expect_get_state_at() + .returning(move |height| { + if height <= CERTIFIED_HEIGHT { + Ok(ic_interfaces_state_manager::Labeled::new( + CERTIFIED_HEIGHT, + Arc::new(state.clone()), + )) + } else { + Err(StateManagerError::StateRemoved(height)) + } + }); + + pool.advance_round_normal_operation_n(CERTIFIED_HEIGHT.get()); + + // Add the message shares + let mutations = shares + .into_iter() + .map(IDkgChangeAction::AddToValidated) + .collect(); + idkg_pool.write().unwrap().apply(mutations); + + let payload_builder = VetKdPayloadBuilderImpl::new( + idkg_pool.clone(), + pool.get_cache(), + crypto, + state_manager, + subnet_id, + registry, + &MetricsRegistry::new(), + no_op_logger(), + ); + + // Run the test + run(payload_builder) + }) + } + + fn build_and_validate( + builder: &VetKdPayloadBuilderImpl, + max_size: NumBytes, + past_payloads: &[PastPayload], + context: &ValidationContext, + ) -> Vec { + let payload = builder.build_payload(HEIGHT, max_size, past_payloads, context); + let context = ProposalContext { + proposer: node_test_id(0), + validation_context: context, + }; + let validation = builder.validate_payload(HEIGHT, &context, &payload, past_payloads); + assert!(validation.is_ok()); + payload + } + + #[test] + #[should_panic(expected = "not yet implemented")] + fn test_build_payload() { + let config = make_chain_key_config(); + let contexts = make_contexts(&config); + let shares = make_shares(&contexts); + let proposal_context = ProposalContext { + proposer: node_test_id(0), + validation_context: &VALIDATION_CONTEXT, + }; + test_payload_builder(Some(config), contexts, shares, |builder| { + let _payload = build_and_validate(&builder, MAX_SIZE, &[], &VALIDATION_CONTEXT); + + // TODO validate payload manually + + // payload that can't be deserialized should be invalid + let validation = builder.validate_payload(HEIGHT, &proposal_context, &[1, 2, 3], &[]); + assert_matches!( + validation.unwrap_err(), + ValidationError::InvalidArtifact(InvalidPayloadReason::InvalidVetKdPayload( + InvalidVetKdPayloadReason::DeserializationFailed(_) + )) + ); + + // payload that rejects valid contexts should be invalid + let payload = as_bytes(make_vetkd_agreements_with_payload( + &[1, 2], + VetKdAgreement::Reject(VetKdErrorCode::TimedOut), + )); + let validation = builder.validate_payload(HEIGHT, &proposal_context, &payload, &[]); + assert_matches!( + validation.unwrap_err(), + ValidationError::InvalidArtifact(InvalidPayloadReason::InvalidVetKdPayload( + InvalidVetKdPayloadReason::MismatchedAgreement { expected, received } + )) if expected.is_none() + && received == Some(VetKdAgreement::Reject(VetKdErrorCode::TimedOut)) + ); + + // Empty payloads should always be valid + let validation = builder.validate_payload(HEIGHT, &proposal_context, &[], &[]); + assert!(validation.is_ok()); + }) + } + + #[test] + fn test_build_empty_payloads_when_feature_disabled() { + // No chain key config is passed + test_payload_builder(None, BTreeMap::new(), vec![], |builder| { + let payload = build_and_validate(&builder, MAX_SIZE, &[], &VALIDATION_CONTEXT); + assert!(payload.is_empty()); + + let proposal_context = ProposalContext { + proposer: node_test_id(0), + validation_context: &VALIDATION_CONTEXT, + }; + + // Non-empty payloads should be rejected + let payload = as_bytes(make_vetkd_agreements(0, 1, 2)); + let validation = builder.validate_payload(HEIGHT, &proposal_context, &payload, &[]); + assert_matches!( + validation.unwrap_err(), + ValidationError::InvalidArtifact(InvalidPayloadReason::InvalidVetKdPayload( + InvalidVetKdPayloadReason::Disabled + )) + ); + }) + } + + #[test] + fn test_build_empty_payload_if_state_doesnt_exist() { + let config = make_chain_key_config(); + let contexts = make_contexts(&config); + let shares = make_shares(&contexts); + let context = ValidationContext { + // There is no state for this certified height yet + certified_height: CERTIFIED_HEIGHT.increment(), + ..VALIDATION_CONTEXT + }; + test_payload_builder(Some(config), contexts, shares, |builder| { + let payload = build_and_validate(&builder, MAX_SIZE, &[], &context); + assert!(payload.is_empty()); + + let proposal_context = ProposalContext { + proposer: node_test_id(0), + validation_context: &context, + }; + + // Non-empty payload validation should be fail if we don't have the state + let payload = as_bytes(make_vetkd_agreements(0, 1, 2)); + let validation = builder.validate_payload(HEIGHT, &proposal_context, &payload, &[]); + assert_matches!( + validation.unwrap_err(), + ValidationError::ValidationFailed( + PayloadValidationFailure::VetKdPayloadValidationFailed( + VetKdPayloadValidationFailure::StateUnavailable(_) + ) + ) + ); + }) + } + + #[test] + fn test_build_empty_payload_if_pool_is_empty() { + let config = make_chain_key_config(); + let contexts = make_contexts(&config); + test_payload_builder(Some(config), contexts, vec![], |builder| { + let payload = build_and_validate(&builder, MAX_SIZE, &[], &VALIDATION_CONTEXT); + assert!(payload.is_empty()); + }) + } + + #[test] + #[should_panic(expected = "not yet implemented")] + fn test_build_empty_payload_max_size_zero() { + let config = make_chain_key_config(); + let contexts = make_contexts(&config); + let shares = make_shares(&contexts); + test_payload_builder(Some(config), contexts, shares, |builder| { + let payload = build_and_validate(&builder, NumBytes::from(0), &[], &VALIDATION_CONTEXT); + assert!(payload.is_empty()); + }) + } + + #[test] + #[should_panic(expected = "not yet implemented")] + fn test_build_payload_respects_max_size() { + let config = make_chain_key_config(); + let contexts = make_contexts(&config); + let shares = make_shares(&contexts); + test_payload_builder(Some(config), contexts, shares, |builder| { + // Use a small maximum size + let payload = + build_and_validate(&builder, NumBytes::from(50), &[], &VALIDATION_CONTEXT); + let payload_deserialized = bytes_to_vetkd_payload(&payload).unwrap(); + assert_eq!(payload_deserialized.vetkd_agreements.len(), 1); + + // TODO validate agreement is success + }) + } + + #[test] + fn test_build_empty_payload_if_all_contexts_answered() { + let config = make_chain_key_config(); + let contexts = make_contexts(&config); + let payloads = [ + as_bytes(make_vetkd_agreements(0, 1, 2)), + as_bytes(make_vetkd_agreements(2, 3, 4)), + ]; + let past_payloads = payloads + .iter() + .map(|bytes| as_past_payload(bytes)) + .collect::>(); + let shares = make_shares(&contexts); + test_payload_builder(Some(config), contexts, shares, |builder| { + let payload = + build_and_validate(&builder, MAX_SIZE, &past_payloads, &VALIDATION_CONTEXT); + assert!(payload.is_empty()); + + // Payload with agreements that are already part of past payloads should be rejected + let payload = as_bytes(make_vetkd_agreements(0, 1, 2)); + let validation = builder.validate_payload( + HEIGHT, + &ProposalContext { + proposer: node_test_id(0), + validation_context: &VALIDATION_CONTEXT, + }, + &payload, + &past_payloads, + ); + assert_matches!( + validation.unwrap_err(), + ValidationError::InvalidArtifact(InvalidPayloadReason::InvalidVetKdPayload( + InvalidVetKdPayloadReason::DuplicateResponse(_) + )) + ); + }) + } + + #[test] + fn test_reject_payloads_for_unknown_contexts() { + let config = make_chain_key_config(); + let contexts = make_contexts(&config); + let shares = make_shares(&contexts); + let proposal_context = ProposalContext { + proposer: node_test_id(0), + validation_context: &VALIDATION_CONTEXT, + }; + test_payload_builder(Some(config), contexts, shares, |builder| { + // Payload with agreements for IDKG contexts should be rejected + let payload = as_bytes(make_vetkd_agreements(0, 1, 2)); + let validation = builder.validate_payload(HEIGHT, &proposal_context, &payload, &[]); + assert_matches!( + validation.unwrap_err(), + ValidationError::InvalidArtifact(InvalidPayloadReason::InvalidVetKdPayload( + InvalidVetKdPayloadReason::UnexpectedIDkgContext(id) + )) if id.get() == 0 + ); + + // Payload with agreements for unknown contexts should be rejected + let payload = as_bytes(make_vetkd_agreements(3, 4, 5)); + let validation = builder.validate_payload(HEIGHT, &proposal_context, &payload, &[]); + assert_matches!( + validation.unwrap_err(), + ValidationError::InvalidArtifact(InvalidPayloadReason::InvalidVetKdPayload( + InvalidVetKdPayloadReason::MissingContext(id) + )) if id.get() == 3 + ); + }) + } + + #[test] + fn test_reject_invalid_keys() { + let config = ChainKeyConfig { + key_configs: vec![KeyConfig { + key_id: MasterPublicKeyId::VetKd( + VetKdKeyId::from_str("bls12_381_g2:unused_key").unwrap(), + ), + pre_signatures_to_create_in_advance: 0, + max_queue_size: 3, + }], + ..ChainKeyConfig::default() + }; + reject_invalid_contexts_test( + config, + true, + VALIDATION_CONTEXT, + VetKdErrorCode::InvalidKey, + VetKdErrorCode::TimedOut, + ); + } + + #[test] + fn test_reject_disabled_keys() { + let config = make_chain_key_config(); + reject_invalid_contexts_test( + config, + false, + VALIDATION_CONTEXT, + VetKdErrorCode::InvalidKey, + VetKdErrorCode::TimedOut, + ); + } + + #[test] + fn test_reject_timed_out_contexts() { + let config = make_chain_key_config(); + let context = ValidationContext { + time: UNIX_EPOCH + Duration::from_secs(2), + ..VALIDATION_CONTEXT + }; + reject_invalid_contexts_test( + config, + true, + context, + VetKdErrorCode::TimedOut, + VetKdErrorCode::InvalidKey, + ); + } + + fn reject_invalid_contexts_test( + config: ChainKeyConfig, + enabled_keys: bool, + validation_context: ValidationContext, + expected_error: VetKdErrorCode, + rejected_error: VetKdErrorCode, + ) { + let contexts = make_contexts(&make_chain_key_config()); + let shares = make_shares(&contexts); + let proposal_context = ProposalContext { + proposer: node_test_id(0), + validation_context: &validation_context, + }; + test_payload_builder_ext( + Some(config), + enabled_keys, + contexts.clone(), + shares, + |builder| { + let serialized_payload = + build_and_validate(&builder, MAX_SIZE, &[], &validation_context); + let payload = bytes_to_vetkd_payload(&serialized_payload).unwrap(); + assert_eq!(payload.vetkd_agreements.len(), 2); + for (id, context) in contexts { + match context.key_id() { + MasterPublicKeyId::Ecdsa(_) | MasterPublicKeyId::Schnorr(_) => { + assert!(!payload.vetkd_agreements.contains_key(&id)); + } + MasterPublicKeyId::VetKd(_) => { + assert_eq!( + payload.vetkd_agreements.get(&id).unwrap(), + &VetKdAgreement::Reject(expected_error) + ); + } + } + } + + // payload with different rejects for the same contexts should be rejected + let payload = as_bytes(make_vetkd_agreements_with_payload( + &[1, 2], + VetKdAgreement::Reject(rejected_error), + )); + let validation = builder.validate_payload(HEIGHT, &proposal_context, &payload, &[]); + assert_matches!( + validation.unwrap_err(), + ValidationError::InvalidArtifact(InvalidPayloadReason::InvalidVetKdPayload( + InvalidVetKdPayloadReason::MismatchedAgreement { expected, received } + )) if expected == Some(VetKdAgreement::Reject(expected_error)) + && received == Some(VetKdAgreement::Reject(rejected_error)) + ); + + // payload with success responses for the same contexts should be rejected + let payload = as_bytes(make_vetkd_agreements_with_payload( + &[1, 2], + VetKdAgreement::Success(vec![1, 1, 1]), + )); + let validation = builder.validate_payload(HEIGHT, &proposal_context, &payload, &[]); + assert_matches!( + validation.unwrap_err(), + ValidationError::InvalidArtifact(InvalidPayloadReason::InvalidVetKdPayload( + InvalidVetKdPayloadReason::MismatchedAgreement { expected, received } + )) if expected == Some(VetKdAgreement::Reject(expected_error)) + && received.is_none() + ); + + // Empty payloads should always be valid + let validation = builder.validate_payload(HEIGHT, &proposal_context, &[], &[]); + assert!(validation.is_ok()); + }, + ) + } + + #[test] + fn test_get_enabled_keys_and_expiry_if_disabled() { + test_payload_builder(None, BTreeMap::new(), vec![], |builder| { + let res = builder + .get_enabled_keys_and_expiry(HEIGHT, UNIX_EPOCH) + .unwrap_err(); + assert_matches!( + res, + ValidationError::InvalidArtifact(InvalidPayloadReason::InvalidVetKdPayload( + InvalidVetKdPayloadReason::Disabled + )) + ) + }) + } + + #[test] + fn test_get_enabled_keys_and_expiry_if_enabled_no_keys() { + test_payload_builder( + Some(ChainKeyConfig::default()), + BTreeMap::new(), + vec![], + |builder| { + let (keys, expiry) = builder + .get_enabled_keys_and_expiry(HEIGHT, UNIX_EPOCH) + .unwrap(); + assert!(keys.is_empty()); + assert!(expiry.is_none()); + }, + ) + } + + #[test] + fn test_get_enabled_keys_and_expiry_if_enabled_multiple_keys() { + let config = make_chain_key_config(); + let timeout = Duration::from_nanos(config.signature_request_timeout_ns.unwrap()); + let now = current_time(); + test_payload_builder(Some(config), BTreeMap::new(), vec![], |builder| { + let (keys, expiry) = builder.get_enabled_keys_and_expiry(HEIGHT, now).unwrap(); + assert_eq!(keys.len(), 2); + assert!(keys.contains(&MasterPublicKeyId::VetKd( + VetKdKeyId::from_str("bls12_381_g2:some_key").unwrap() + ))); + assert!(keys.contains(&MasterPublicKeyId::VetKd( + VetKdKeyId::from_str("bls12_381_g2:some_other_key").unwrap() + ))); + assert_matches!(expiry, Some(time) if time == now.saturating_sub(timeout)); + }) + } + + #[test] + fn test_get_enabled_keys_and_expiry_if_disabled_multiple_keys() { + let config = make_chain_key_config(); + let timeout = Duration::from_nanos(config.signature_request_timeout_ns.unwrap()); + let now = current_time(); + test_payload_builder_ext(Some(config), false, BTreeMap::new(), vec![], |builder| { + let (keys, expiry) = builder.get_enabled_keys_and_expiry(HEIGHT, now).unwrap(); + assert!(keys.is_empty()); + assert_matches!(expiry, Some(time) if time == now.saturating_sub(timeout)); + }) + } +} diff --git a/rs/consensus/vetkd/src/test_utils.rs b/rs/consensus/vetkd/src/test_utils.rs new file mode 100644 index 00000000000..b9d8473e90e --- /dev/null +++ b/rs/consensus/vetkd/src/test_utils.rs @@ -0,0 +1,212 @@ +use core::{convert::From, iter::Iterator}; +use ic_interfaces::batch_payload::PastPayload; +use ic_management_canister_types_private::{EcdsaKeyId, MasterPublicKeyId, VetKdKeyId}; +use ic_registry_subnet_features::{ChainKeyConfig, KeyConfig}; +use ic_replicated_state::metadata_state::subnet_call_context_manager::{ + EcdsaArguments, SchnorrArguments, SignWithThresholdContext, ThresholdArguments, VetKdArguments, +}; +use ic_test_utilities_types::messages::RequestBuilder; +use ic_types::consensus::idkg::{EcdsaSigShare, IDkgMessage, RequestId, SchnorrSigShare}; +use ic_types::crypto::canister_threshold_sig::{ThresholdEcdsaSigShare, ThresholdSchnorrSigShare}; +use ic_types::crypto::threshold_sig::ni_dkg::{ + NiDkgId, NiDkgMasterPublicKeyId, NiDkgTag, NiDkgTargetSubnet, +}; +use ic_types::{ + batch::{vetkd_payload_to_bytes, VetKdAgreement, VetKdErrorCode, VetKdPayload}, + consensus::idkg::VetKdKeyShare, + crypto::vetkd::VetKdEncryptedKeyShare, + crypto::vetkd::VetKdEncryptedKeyShareContent, + crypto::{CryptoHash, CryptoHashOf}, + messages::CallbackId, + time::UNIX_EPOCH, + Height, NumBytes, +}; +use ic_types_test_utils::ids::{node_test_id, subnet_test_id}; +use std::str::FromStr; +use std::{collections::BTreeMap, sync::Arc}; +use strum::EnumCount; + +/// Create a map of agreements with all possible types +pub(super) fn make_vetkd_agreements( + id1: u64, + id2: u64, + id3: u64, +) -> BTreeMap { + assert_eq!(VetKdAgreement::COUNT, 2); + assert_eq!(VetKdErrorCode::COUNT, 2); + BTreeMap::from([ + ( + CallbackId::from(id1), + VetKdAgreement::Success(vec![1, 2, 3, 4]), + ), + ( + CallbackId::from(id2), + VetKdAgreement::Reject(VetKdErrorCode::TimedOut), + ), + ( + CallbackId::from(id3), + VetKdAgreement::Reject(VetKdErrorCode::InvalidKey), + ), + ]) +} + +/// Create a map of agreements with the same, given type +pub(super) fn make_vetkd_agreements_with_payload( + ids: &[u64], + agreement: VetKdAgreement, +) -> BTreeMap { + let mut map = BTreeMap::new(); + for id in ids { + map.insert(CallbackId::new(*id), agreement.clone()); + } + map +} + +/// Convert the given agreements payload to bytes, using a maximum size of 1KiB. +pub(super) fn as_bytes(vetkd_agreements: BTreeMap) -> Vec { + vetkd_payload_to_bytes(VetKdPayload { vetkd_agreements }, NumBytes::new(1024)) +} + +/// Turn the given payload bytes into a generic [`PastPayload`] +pub(super) fn as_past_payload(payload: &[u8]) -> PastPayload { + PastPayload { + height: Height::from(0), + time: UNIX_EPOCH, + block_hash: CryptoHashOf::from(CryptoHash(vec![])), + payload, + } +} + +/// Create a [`ChainKeyConfig`] with one ECDSA and two VetKD key IDs, +/// and 1 second request timeout +pub(super) fn make_chain_key_config() -> ChainKeyConfig { + let key_config = KeyConfig { + key_id: MasterPublicKeyId::Ecdsa(EcdsaKeyId::from_str("Secp256k1:some_key_1").unwrap()), + pre_signatures_to_create_in_advance: 1, + max_queue_size: 3, + }; + let key_config_1 = KeyConfig { + key_id: MasterPublicKeyId::VetKd(VetKdKeyId::from_str("bls12_381_g2:some_key").unwrap()), + pre_signatures_to_create_in_advance: 1, + max_queue_size: 3, + }; + let key_config_2 = KeyConfig { + key_id: MasterPublicKeyId::VetKd( + VetKdKeyId::from_str("bls12_381_g2:some_other_key").unwrap(), + ), + pre_signatures_to_create_in_advance: 1, + max_queue_size: 3, + }; + + ChainKeyConfig { + key_configs: vec![ + key_config.clone(), + key_config_1.clone(), + key_config_2.clone(), + ], + // 1 second timeout + signature_request_timeout_ns: Some(1_000_000_000), + ..ChainKeyConfig::default() + } +} + +pub(super) fn fake_dkg_id(key_id: VetKdKeyId) -> NiDkgId { + NiDkgId { + start_block_height: Height::from(0), + dealer_subnet: subnet_test_id(0), + dkg_tag: NiDkgTag::HighThresholdForKey(NiDkgMasterPublicKeyId::VetKd(key_id)), + target_subnet: NiDkgTargetSubnet::Local, + } +} + +pub(super) fn fake_signature_request_args(key_id: MasterPublicKeyId) -> ThresholdArguments { + match key_id { + MasterPublicKeyId::Ecdsa(key_id) => ThresholdArguments::Ecdsa(EcdsaArguments { + key_id, + message_hash: [0; 32], + }), + MasterPublicKeyId::Schnorr(key_id) => ThresholdArguments::Schnorr(SchnorrArguments { + key_id, + message: Arc::new(vec![1; 48]), + taproot_tree_root: None, + }), + MasterPublicKeyId::VetKd(key_id) => ThresholdArguments::VetKd(VetKdArguments { + key_id: key_id.clone(), + derivation_id: vec![1; 32], + encryption_public_key: vec![1; 32], + ni_dkg_id: fake_dkg_id(key_id), + height: Height::from(0), + }), + } +} + +pub(super) fn fake_signature_request_context( + key_id: MasterPublicKeyId, +) -> SignWithThresholdContext { + SignWithThresholdContext { + request: RequestBuilder::new().build(), + args: fake_signature_request_args(key_id), + derivation_path: vec![], + batch_time: UNIX_EPOCH, + pseudo_random_id: [0; 32], + matched_pre_signature: None, + nonce: None, + } +} + +/// Create a fake request context for each key ID in the given config. +/// Callback IDs are assigned sequentially starting at 0. +pub(super) fn make_contexts( + config: &ChainKeyConfig, +) -> BTreeMap { + let mut map = BTreeMap::new(); + for (i, key_id) in config.key_ids().into_iter().enumerate() { + map.insert( + CallbackId::new(i as u64), + fake_signature_request_context(key_id), + ); + } + map +} + +/// Create four artifact shares for each request context +pub(super) fn make_shares( + contexts: &BTreeMap, +) -> Vec { + let committee = (0..4).map(|id| node_test_id(id as u64)).collect::>(); + let mut messages = vec![]; + for (&callback_id, context) in contexts { + for &signer_id in &committee { + let request_id = RequestId { + callback_id, + height: Height::from(0), + }; + let message = match context.args { + ThresholdArguments::Ecdsa(_) => IDkgMessage::EcdsaSigShare(EcdsaSigShare { + signer_id, + request_id, + share: ThresholdEcdsaSigShare { + sig_share_raw: vec![], + }, + }), + ThresholdArguments::Schnorr(_) => IDkgMessage::SchnorrSigShare(SchnorrSigShare { + signer_id, + request_id, + share: ThresholdSchnorrSigShare { + sig_share_raw: vec![], + }, + }), + ThresholdArguments::VetKd(_) => IDkgMessage::VetKdKeyShare(VetKdKeyShare { + signer_id, + request_id, + share: VetKdEncryptedKeyShare { + encrypted_key_share: VetKdEncryptedKeyShareContent(vec![]), + node_signature: vec![], + }, + }), + }; + messages.push(message); + } + } + messages +} diff --git a/rs/consensus/vetkd/src/utils.rs b/rs/consensus/vetkd/src/utils.rs index 05fc7df2eef..060ad525f26 100644 --- a/rs/consensus/vetkd/src/utils.rs +++ b/rs/consensus/vetkd/src/utils.rs @@ -64,3 +64,28 @@ pub(super) fn parse_past_payload_ids( .map(|msg| CallbackId::new(msg.callback_id)) .collect() } + +#[cfg(test)] +mod tests { + use core::{convert::From, iter::Iterator}; + use ic_logger::no_op_logger; + + use super::*; + use crate::test_utils::*; + + #[test] + fn test_parse_past_payload_ids() { + let payloads = [ + as_bytes(make_vetkd_agreements(0, 1, 2)), + as_bytes(make_vetkd_agreements(2, 3, 4)), + as_bytes(make_vetkd_agreements(4, 4, 5)), + ]; + let past_payloads = payloads + .iter() + .map(|p| as_past_payload(p)) + .collect::>(); + let past_payload_ids = parse_past_payload_ids(&past_payloads, &no_op_logger()); + let expected = HashSet::from_iter((0..=5).map(CallbackId::from)); + assert_eq!(past_payload_ids, expected); + } +} diff --git a/rs/types/types/src/batch/vetkd.rs b/rs/types/types/src/batch/vetkd.rs index 21987ae6eae..6432d4aa644 100644 --- a/rs/types/types/src/batch/vetkd.rs +++ b/rs/types/types/src/batch/vetkd.rs @@ -10,13 +10,14 @@ use ic_protobuf::{ }; use serde::{Deserialize, Serialize}; use std::{collections::BTreeMap, convert::TryFrom}; +use strum_macros::EnumCount; use crate::{messages::CallbackId, CountBytes}; use super::{iterator_to_bytes, slice_to_messages}; /// Errors that may occur when handling a VetKd request. -#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Deserialize, Serialize)] +#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Deserialize, Serialize, EnumCount)] #[cfg_attr(test, derive(ExhaustiveSet))] pub enum VetKdErrorCode { TimedOut = 1, @@ -24,7 +25,7 @@ pub enum VetKdErrorCode { } /// Consensus may either agree on a successful response, or reject the request. -#[derive(Clone, Eq, PartialEq, Hash, Debug, Deserialize, Serialize)] +#[derive(Clone, Eq, PartialEq, Hash, Debug, Deserialize, Serialize, EnumCount)] #[cfg_attr(test, derive(ExhaustiveSet))] pub enum VetKdAgreement { Success(Vec),