chore(crypto): CRP-2629 Add VetKD to the state machine tests (#4625)

Adds the ability to write state machine tests that perform VetKD calls,
and adds some tests in execution which make use of the state machine
test framework.

Author: randombit

Cargo.lock

@@ -112,6 +112,7 @@ members = [
     "rs/crypto/test_utils/metrics",
     "rs/crypto/test_utils/ni-dkg",
     "rs/crypto/test_utils/tls",
+    "rs/crypto/test_utils/vetkd",
     "rs/crypto/tls_interfaces",
     "rs/crypto/tls_interfaces/mocks",
     "rs/crypto/tree_hash",

Cargo.toml

@@ -267,7 +267,7 @@ impl EncryptedVetKey {
         &self,
         tsk: &TransportSecretKey,
         derived_public_key: &DerivedPublicKey,
-        context: &[u8],
+        input: &[u8],
     ) -> Result<VetKey, String> {
         // Check that c1 and c2 have the same discrete logarithm
 
@@ -286,7 +286,7 @@ impl EncryptedVetKey {
         let k = G1Affine::from(G1Projective::from(&self.c3) - self.c1 * tsk.secret_key);
 
         // Check that the VetKey is a valid BLS signature
-        let msg = augmented_hash_to_g1(&derived_public_key.point, context);
+        let msg = augmented_hash_to_g1(&derived_public_key.point, input);
         let dpk_prep = G2Prepared::from(derived_public_key.point);
 
         use pairing::group::Group;

packages/ic-vetkd-utils/src/lib.rs

@@ -0,0 +1,29 @@
+load("@rules_rust//rust:defs.bzl", "rust_library", "rust_test_suite")
+
+package(default_visibility = ["//visibility:public"])
+
+rust_library(
+    name = "vetkd",
+    testonly = True,
+    srcs = glob(["src/**"]),
+    crate_name = "ic_crypto_test_utils_vetkd",
+    version = "0.1.0",
+    deps = [
+        "//rs/crypto/internal/crypto_lib/bls12_381/type",
+        "//rs/crypto/internal/crypto_lib/bls12_381/vetkd",
+        "@crate_index//:rand_chacha",
+    ],
+)
+
+rust_test_suite(
+    name = "vetkd_tests",
+    srcs = glob(["tests/**/*.rs"]),
+    deps = [
+        ":vetkd",
+        "//packages/ic-vetkd-utils",
+        "//rs/crypto/internal/crypto_lib/bls12_381/type",
+        "//rs/crypto/internal/crypto_lib/bls12_381/vetkd",
+        "@crate_index//:rand",
+        "@crate_index//:rand_chacha",
+    ],
+)

rs/crypto/test_utils/vetkd/BUILD.bazel

@@ -0,0 +1,13 @@
+[package]
+name = "ic-crypto-test-utils-vetkd"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+ic-crypto-internal-bls12-381-type = { path = "../../internal/crypto_lib/bls12_381/type" }
+ic-crypto-internal-bls12-381-vetkd = { path = "../../internal/crypto_lib/bls12_381/vetkd" }
+rand_chacha = { workspace = true }
+
+[dev-dependencies]
+ic-vetkd-utils = { path = "../../../../packages/ic-vetkd-utils" }
+rand = { workspace = true }

rs/crypto/test_utils/vetkd/Cargo.toml

@@ -0,0 +1,66 @@
+use ic_crypto_internal_bls12_381_type::{G1Affine, G2Affine, Scalar};
+use ic_crypto_internal_bls12_381_vetkd::{
+    DerivationContext, EncryptedKey, EncryptedKeyShare, TransportPublicKey,
+};
+use rand_chacha::rand_core::SeedableRng;
+
+pub fn dummy_transport_public_key() -> [u8; 48] {
+    G1Affine::generator().serialize()
+}
+
+pub struct PrivateKey {
+    secret_key: Scalar,
+    public_point: G2Affine,
+    pk_bytes: Vec<u8>,
+}
+
+impl PrivateKey {
+    pub fn generate(seed: &[u8]) -> Self {
+        let secret_key = Scalar::hash(b"ic-crypto-vetkd-test-utils-generate-test-key", seed);
+        Self::from_scalar(secret_key)
+    }
+
+    fn from_scalar(secret_key: Scalar) -> Self {
+        let public_point = G2Affine::from(G2Affine::generator() * &secret_key);
+        let pk_bytes = public_point.serialize().to_vec();
+        Self {
+            secret_key,
+            public_point,
+            pk_bytes,
+        }
+    }
+
+    pub fn public_key_bytes(&self) -> Vec<u8> {
+        self.pk_bytes.clone()
+    }
+
+    pub fn vetkd_protocol(
+        &self,
+        canister_id: &[u8],
+        context: &[u8],
+        input: &[u8],
+        tpk: &[u8],
+        seed: &[u8; 32],
+    ) -> Vec<u8> {
+        let dc = DerivationContext::new(canister_id, context);
+
+        let tpk =
+            TransportPublicKey::deserialize(tpk).expect("Failed to deserialize TransportPublicKey");
+
+        let mut rng = rand_chacha::ChaCha20Rng::from_seed(*seed);
+
+        let eks = EncryptedKeyShare::create(
+            &mut rng,
+            &self.public_point,
+            &self.secret_key,
+            &tpk,
+            &dc,
+            input,
+        );
+
+        let ek = EncryptedKey::combine_all(&[(0, eks)], 1, &self.public_point, &tpk, &dc, input)
+            .expect("Failed to combine single EncryptedKeyShare to an EncryptedKey");
+
+        ek.serialize().to_vec()
+    }
+}

rs/crypto/test_utils/vetkd/src/lib.rs

@@ -0,0 +1,34 @@
+use ic_crypto_test_utils_vetkd::*;
+use ic_vetkd_utils::{DerivedPublicKey, EncryptedVetKey, TransportSecretKey};
+use rand::Rng;
+use rand_chacha::rand_core::SeedableRng;
+
+#[test]
+fn should_generate_valid_bls_signature() {
+    let mut rng = rand_chacha::ChaCha20Rng::seed_from_u64(42);
+
+    let pk = PrivateKey::generate(&rng.gen::<[u8; 32]>());
+
+    let canister_id = rng.gen::<[u8; 32]>();
+    let context = rng.gen::<[u8; 32]>();
+    let input = rng.gen::<[u8; 32]>();
+
+    let tsk = TransportSecretKey::from_seed(rng.gen::<[u8; 32]>().to_vec()).unwrap();
+
+    let ek_bytes = pk.vetkd_protocol(
+        &canister_id,
+        &context,
+        &input,
+        &tsk.public_key(),
+        &rng.gen::<[u8; 32]>(),
+    );
+
+    let ek = EncryptedVetKey::deserialize(&ek_bytes).unwrap();
+
+    let dpk = DerivedPublicKey::deserialize(&pk.public_key_bytes())
+        .unwrap()
+        .derive_sub_key(&canister_id)
+        .derive_sub_key(&context);
+
+    assert!(ek.decrypt_and_verify(&tsk, &dpk, &input).is_ok());
+}

rs/crypto/test_utils/vetkd/tests/tests.rs

@@ -63,6 +63,7 @@ MACRO_DEPENDENCIES = []
 
 DEV_DEPENDENCIES = [
     # Keep sorted.
+    "//rs/crypto/test_utils/vetkd",
     "//rs/interfaces/state_manager/mocks",
     "//rs/rust_canisters/canister_test",
     "//rs/state_machine_tests",

rs/execution_environment/BUILD.bazel

@@ -66,6 +66,7 @@ assert_matches = { workspace = true }
 canister-test = { path = "../rust_canisters/canister_test" }
 criterion = { workspace = true }
 execution-environment-bench = { path = "benches/lib" }
+ic-crypto-test-utils-vetkd = { path = "../crypto/test_utils/vetkd" }
 ic-interfaces-state-manager-mocks = { path = "../interfaces/state_manager/mocks" }
 ic-management-canister-types-private = { path = "../types/management_canister_types" }
 ic-state-machine-tests = { path = "../state_machine_tests" }

rs/execution_environment/Cargo.toml

@@ -4,7 +4,7 @@ use ic_management_canister_types_private::{
     self as ic00, CanisterInstallMode, DerivationPath, ECDSAPublicKeyResponse, EcdsaCurve,
     EcdsaKeyId, MasterPublicKeyId, Method, Payload as Ic00Payload, SchnorrAlgorithm, SchnorrKeyId,
     SchnorrPublicKeyResponse, SignWithBip341Aux, SignWithECDSAReply, SignWithSchnorrAux,
-    SignWithSchnorrReply,
+    SignWithSchnorrReply, VetKdCurve, VetKdDeriveKeyResult, VetKdKeyId, VetKdPublicKeyResult,
 };
 use ic_registry_subnet_type::SubnetType;
 use ic_state_machine_tests::{StateMachine, StateMachineBuilder, UserError};
@@ -48,6 +48,13 @@ fn make_bip340_key(name: &str) -> MasterPublicKeyId {
     })
 }
 
+fn make_vetkd_key(name: &str) -> MasterPublicKeyId {
+    MasterPublicKeyId::VetKd(VetKdKeyId {
+        curve: VetKdCurve::Bls12_381_G2,
+        name: name.to_string(),
+    })
+}
+
 fn into_inner_ecdsa(key_id: MasterPublicKeyId) -> EcdsaKeyId {
     match key_id {
         MasterPublicKeyId::Ecdsa(key) => key,
@@ -62,6 +69,13 @@ fn into_inner_schnorr(key_id: MasterPublicKeyId) -> SchnorrKeyId {
     }
 }
 
+fn into_inner_vetkd(key_id: MasterPublicKeyId) -> VetKdKeyId {
+    match key_id {
+        MasterPublicKeyId::VetKd(key) => key,
+        _ => panic!("unexpected key_id type"),
+    }
+}
+
 fn compute_initial_threshold_key_dealings_payload(
     method: Method,
     key_id: MasterPublicKeyId,
@@ -107,6 +121,17 @@ fn sign_with_threshold_key_payload(method: Method, key_id: MasterPublicKeyId) ->
             }
         }
         .encode(),
+        Method::VetKdDeriveKey => {
+            let key_id = into_inner_vetkd(key_id);
+
+            ic00::VetKdDeriveKeyArgs {
+                context: vec![],
+                input: vec![],
+                key_id,
+                transport_public_key: ic_crypto_test_utils_vetkd::dummy_transport_public_key(),
+            }
+        }
+        .encode(),
         _ => panic!("unexpected method"),
     }
 }
@@ -125,6 +150,12 @@ fn threshold_public_key_payload(method: Method, key_id: MasterPublicKeyId) -> Ve
             key_id: into_inner_schnorr(key_id),
         }
         .encode(),
+        Method::VetKdPublicKey => ic00::VetKdPublicKeyArgs {
+            canister_id: None,
+            context: vec![],
+            key_id: into_inner_vetkd(key_id),
+        }
+        .encode(),
         _ => panic!("unexpected method"),
     }
 }
@@ -422,6 +453,7 @@ fn test_sign_with_threshold_key_fee_charged() {
         let contexts = match method {
             Method::SignWithECDSA => env.sign_with_ecdsa_contexts(),
             Method::SignWithSchnorr => env.sign_with_schnorr_contexts(),
+            Method::VetKdDeriveKey => env.vetkd_derive_key_contexts(),
             _ => panic!("Unexpected method"),
         };
         let (_, context) = contexts.iter().next().unwrap();
@@ -435,6 +467,7 @@ fn test_sign_with_threshold_key_fee_charged() {
         let signature = match method {
             Method::SignWithECDSA => expect_reply::<SignWithECDSAReply>(result).signature,
             Method::SignWithSchnorr => expect_reply::<SignWithSchnorrReply>(result).signature,
+            Method::VetKdDeriveKey => expect_reply::<VetKdDeriveKeyResult>(result).encrypted_key,
             _ => panic!("Unexpected method"),
         };
         // Expect non-empty signature.
@@ -512,6 +545,11 @@ fn test_sign_with_threshold_key_unknown_key_rejected() {
             make_bip340_key("correct_key"),
             make_bip340_key("wrong_key"),
         ),
+        (
+            Method::VetKdDeriveKey,
+            make_vetkd_key("correct_key"),
+            make_vetkd_key("wrong_key"),
+        ),
     ];
     for (method, correct_key, wrong_key) in test_cases {
         let own_subnet = subnet_test_id(1);
@@ -624,6 +662,12 @@ fn test_signing_disabled_vs_unknown_key_on_public_key_and_signing_requests() {
             make_bip340_key("signing_disabled_key"),
             make_bip340_key("unknown_key"),
         ),
+        (
+            Method::VetKdPublicKey,
+            Method::VetKdDeriveKey,
+            make_vetkd_key("signing_disabled_key"),
+            make_vetkd_key("unknown_key"),
+        ),
     ];
     for (public_key_method, sign_with_method, signing_disabled_key, unknown_key) in test_cases {
         let own_subnet = subnet_test_id(1);
@@ -654,6 +698,10 @@ fn test_signing_disabled_vs_unknown_key_on_public_key_and_signing_requests() {
                 let response = expect_reply::<SchnorrPublicKeyResponse>(result);
                 assert!(!response.public_key.is_empty() && !response.chain_code.is_empty());
             }
+            Method::VetKdPublicKey => {
+                let response = expect_reply::<VetKdPublicKeyResult>(result);
+                assert!(!response.public_key.is_empty());
+            }
             _ => panic!("Unexpected method"),
         }
 
@@ -710,6 +758,11 @@ fn test_threshold_key_public_key_req_with_unknown_key_rejected() {
             make_bip340_key("correct_key"),
             make_bip340_key("wrong_key"),
         ),
+        (
+            Method::VetKdPublicKey,
+            make_vetkd_key("correct_key"),
+            make_vetkd_key("wrong_key"),
+        ),
     ];
     for (method, correct_key, wrong_key) in test_cases {
         let own_subnet = subnet_test_id(1);
@@ -742,6 +795,7 @@ fn test_sign_with_threshold_key_fee_ignored_for_nns() {
         (Method::SignWithECDSA, make_ecdsa_key("some_key")),
         (Method::SignWithSchnorr, make_ed25519_key("some_key")),
         (Method::SignWithSchnorr, make_bip340_key("some_key")),
+        (Method::VetKdDeriveKey, make_vetkd_key("some_key")),
     ];
     for (method, key_id) in test_cases {
         let fee = 1_000_000;
@@ -778,6 +832,7 @@ fn test_sign_with_threshold_key_fee_ignored_for_nns() {
         let contexts = match method {
             Method::SignWithECDSA => env.sign_with_ecdsa_contexts(),
             Method::SignWithSchnorr => env.sign_with_schnorr_contexts(),
+            Method::VetKdDeriveKey => env.vetkd_derive_key_contexts(),
             _ => panic!("Unexpected method"),
         };
         let (_, context) = contexts.iter().next().unwrap();
@@ -791,6 +846,7 @@ fn test_sign_with_threshold_key_queue_fills_up() {
         (Method::SignWithECDSA, make_ecdsa_key("some_key"), 20),
         (Method::SignWithSchnorr, make_ed25519_key("some_key"), 20),
         (Method::SignWithSchnorr, make_bip340_key("some_key"), 20),
+        (Method::VetKdDeriveKey, make_vetkd_key("some_key"), 20),
     ];
     for (method, key_id, max_queue_size) in test_cases {
         let fee = 1_000_000;
@@ -804,11 +860,14 @@ fn test_sign_with_threshold_key_queue_fills_up() {
             .with_nns_subnet_id(nns_subnet)
             .with_ecdsa_signature_fee(fee)
             .with_schnorr_signature_fee(fee)
+            .with_vetkd_derive_key_fee(fee)
             .with_chain_key(key_id.clone())
             // Turn off automatic ECDSA signatures to fill up the queue.
             .with_ecdsa_signing_enabled(false)
             // Turn off automatic Schnorr signatures to fill up the queue.
             .with_schnorr_signing_enabled(false)
+            // Turn off automatic VetKey derivation to fill up the queue.
+            .with_vetkd_enabled(false)
             .build();
 
         let canister_id = create_universal_canister(&env);