add tests for VK change

aptos-move/aptos-vm/src/aptos_vm.rs

@@ -268,7 +268,10 @@ impl AptosVM {
         // via governance (although, currently, we do check for that in `keyless_account.move`).
         let pvk = keyless_validation::get_groth16_vk_onchain(resolver)
             .ok()
-            .and_then(|vk| vk.try_into().ok());
+            .and_then(|vk| {
+                // println!("[aptos-vm][groth16] PVK cached in VM: {}", vk.hash());
+                vk.try_into().ok()
+            });
 
         Self {
             is_simulation: false,
@@ -1581,6 +1584,7 @@ impl AptosVM {
             // Or, if we spawn a network without initializing the VK in genesis. Either way, it must
             // be handled here.
             if self.pvk.is_none() {
+                // println!("[aptos-vm][groth16] PVK has not been set on-chain");
                 return Err(invalid_signature!("Groth16 VK has not been set on-chain"));
             }
 

aptos-move/aptos-vm/src/keyless_validation.rs

@@ -148,14 +148,18 @@ pub(crate) fn validate_authenticators(
 
     let config = &get_configs_onchain(resolver)?;
     if authenticators.len() > config.max_signatures_per_txn as usize {
+        // println!("[aptos-vm][groth16] Too many keyless authenticators");
         return Err(invalid_signature!("Too many keyless authenticators"));
     }
 
     let onchain_timestamp_obj = get_current_time_onchain(resolver)?;
     // Check the expiry timestamp on all authenticators first to fail fast
     for (_, sig) in authenticators {
-        sig.verify_expiry(&onchain_timestamp_obj)
-            .map_err(|_| invalid_signature!("The ephemeral keypair has expired"))?;
+        sig.verify_expiry(&onchain_timestamp_obj).map_err(|_| {
+            // println!("[aptos-vm][groth16] ZKP expired");
+
+            invalid_signature!("The ephemeral keypair has expired")
+        })?;
     }
 
     let patched_jwks = get_jwks_onchain(resolver)?;
@@ -166,6 +170,8 @@ pub(crate) fn validate_authenticators(
         // (Run `cargo bench -- ed25519/pk_deserialize` in `crates/aptos-crypto`.)
         Some(bytes) => Some(EphemeralPublicKey::ed25519(
             Ed25519PublicKey::try_from(bytes.as_slice()).map_err(|_| {
+                // println!("[aptos-vm][groth16] On chain TW PK is invalid");
+
                 invalid_signature!("The training wheels PK set on chain is not a valid PK")
             })?,
         )),
@@ -178,6 +184,7 @@ pub(crate) fn validate_authenticators(
             EphemeralCertificate::ZeroKnowledgeSig(zksig) => match jwk {
                 JWK::RSA(rsa_jwk) => {
                     if zksig.exp_horizon_secs > config.max_exp_horizon_secs {
+                        // println!("[aptos-vm][groth16] Expiration horizon is too long");
                         return Err(invalid_signature!("The expiration horizon is too long"));
                     }
 
@@ -190,10 +197,13 @@ pub(crate) fn validate_authenticators(
                     match &zksig.proof {
                         ZKP::Groth16(groth16proof) => {
                             // let start = std::time::Instant::now();
-                            let public_inputs_hash =
-                                get_public_inputs_hash(sig, pk, &rsa_jwk, config).map_err(
-                                    |_| invalid_signature!("Could not compute public inputs hash"),
-                                )?;
+                            let public_inputs_hash = get_public_inputs_hash(
+                                sig, pk, &rsa_jwk, config,
+                            )
+                            .map_err(|_| {
+                                // println!("[aptos-vm][groth16] PIH computation failed");
+                                invalid_signature!("Could not compute public inputs hash")
+                            })?;
                             // println!("Public inputs hash time: {:?}", start.elapsed());
 
                             let groth16_and_stmt =
@@ -209,22 +219,37 @@ pub(crate) fn validate_authenticators(
                                                 training_wheels_pk.as_ref().unwrap(),
                                             )
                                             .map_err(|_| {
+                                                // println!("[aptos-vm][groth16] TW sig verification failed");
                                                 invalid_signature!(
                                                     "Could not verify training wheels signature"
                                                 )
                                             })?;
                                     },
                                     None => {
+                                        // println!("[aptos-vm][groth16] Expected TW sig to be set");
                                         return Err(invalid_signature!(
                                             "Training wheels signature expected but it is missing"
-                                        ))
+                                        ));
                                     },
                                 }
                             }
 
                             let result = zksig.verify_groth16_proof(public_inputs_hash, pvk);
 
-                            result.map_err(|_| invalid_signature!("Proof verification failed"))?;
+                            result.map_err(|_| {
+                                // println!("[aptos-vm][groth16] ZKP verification failed");
+                                // println!("[aptos-vm][groth16] PIH: {}", public_inputs_hash);
+                                // match zksig.proof {
+                                //     ZKP::Groth16(proof) => {
+                                //         println!("[aptos-vm][groth16] ZKP: {}", proof.hash());
+                                //     },
+                                // }
+                                // println!(
+                                //     "[aptos-vm][groth16] PVK: {}",
+                                //     Groth16VerificationKey::from(pvk).hash()
+                                // );
+                                invalid_signature!("Proof verification failed")
+                            })?;
                         },
                     }
                 },

aptos-move/e2e-move-tests/src/tests/keyless_feature_gating.rs

@@ -8,11 +8,11 @@ use aptos_language_e2e_tests::account::{Account, AccountPublicKey, TransactionBu
 use aptos_types::{
     keyless::{
         test_utils::{
-            get_sample_esk, get_sample_groth16_sig_and_pk, get_sample_iss, get_sample_jwk,
-            get_sample_openid_sig_and_pk,
+            get_groth16_sig_and_pk_for_upgraded_vk, get_sample_esk, get_sample_groth16_sig_and_pk,
+            get_sample_iss, get_sample_jwk, get_sample_openid_sig_and_pk, get_upgraded_vk,
         },
-        Configuration, EphemeralCertificate, KeylessPublicKey, KeylessSignature,
-        TransactionAndProof,
+        Configuration, EphemeralCertificate, Groth16VerificationKey, KeylessPublicKey,
+        KeylessSignature, TransactionAndProof,
     },
     on_chain_config::FeatureFlag,
     transaction::{
@@ -21,22 +21,23 @@ use aptos_types::{
     },
 };
 use move_core_types::{
-    account_address::AccountAddress, transaction_argument::TransactionArgument,
-    vm_status::StatusCode::FEATURE_UNDER_GATING,
+    account_address::AccountAddress,
+    transaction_argument::TransactionArgument,
+    vm_status::{StatusCode, StatusCode::FEATURE_UNDER_GATING},
 };
 
 fn init_feature_gating(
     enabled_features: Vec<FeatureFlag>,
     disabled_features: Vec<FeatureFlag>,
-) -> (MoveHarness, Account) {
+) -> (MoveHarness, Account, Account) {
     let mut h = MoveHarness::new_with_features(enabled_features, disabled_features);
 
     let recipient = h.new_account_at(AccountAddress::from_hex_literal("0xb0b").unwrap());
 
     // initialize JWKs
-    run_setup_script(&mut h);
+    let core_resources = run_jwk_and_config_script(&mut h);
 
-    (h, recipient)
+    (h, recipient, core_resources)
 }
 
 fn test_feature_gating(
@@ -75,11 +76,69 @@ fn test_feature_gating(
     }
 }
 
+#[test]
+fn test_rotate_vk() {
+    let (mut h, recipient, core_resources) = init_feature_gating(
+        vec![
+            FeatureFlag::CRYPTOGRAPHY_ALGEBRA_NATIVES,
+            FeatureFlag::BN254_STRUCTURES,
+            FeatureFlag::KEYLESS_ACCOUNTS,
+        ],
+        vec![],
+    );
+
+    // Old proof for old VK
+    let (old_sig, pk) = get_sample_groth16_sig_and_pk();
+    let account = create_keyless_account(&mut h, pk.clone());
+    let transaction =
+        spend_keyless_account(&mut h, old_sig.clone(), &account, *recipient.address());
+    let output = h.run_raw(transaction);
+    assert_success!(output.status().clone());
+
+    // New proof for old VK
+    let (new_sig, _) = get_groth16_sig_and_pk_for_upgraded_vk();
+    let transaction =
+        spend_keyless_account(&mut h, new_sig.clone(), &account, *recipient.address());
+    let output = h.run_raw(transaction);
+    //println!("TXN status: {:?}", output.status());
+    match output.status() {
+        TransactionStatus::Discard(sc) => assert_eq!(*sc, StatusCode::INVALID_SIGNATURE),
+        TransactionStatus::Keep(es) => {
+            panic!("Expected TransactionStatus::Discard, got Keep({:?})", es)
+        },
+        TransactionStatus::Retry => panic!("Expected TransactionStatus::Discard, got Retry"),
+    }
+
+    // Upgrade the VK
+    run_upgrade_vk_script(
+        &mut h,
+        core_resources,
+        Groth16VerificationKey::from(get_upgraded_vk()),
+    );
+
+    // New proof for new VK
+    let transaction = spend_keyless_account(&mut h, new_sig, &account, *recipient.address());
+    let output = h.run_raw(transaction);
+    assert_success!(output.status().clone());
+
+    // Old proof for old VK
+    let transaction = spend_keyless_account(&mut h, old_sig, &account, *recipient.address());
+    let output = h.run_raw(transaction);
+    // println!("TXN status: {:?}", output.status());
+    match output.status() {
+        TransactionStatus::Discard(sc) => assert_eq!(*sc, StatusCode::INVALID_SIGNATURE),
+        TransactionStatus::Keep(es) => {
+            panic!("Expected TransactionStatus::Discard, got Keep({:?})", es)
+        },
+        TransactionStatus::Retry => panic!("Expected TransactionStatus::Discard, got Retry"),
+    }
+}
+
 #[test]
 fn test_feature_gating_with_zk_on() {
     //
     // ZK & ZKless
-    let (mut h, recipient) = init_feature_gating(
+    let (mut h, recipient, _) = init_feature_gating(
         vec![
             FeatureFlag::CRYPTOGRAPHY_ALGEBRA_NATIVES,
             FeatureFlag::BN254_STRUCTURES,
@@ -95,7 +154,7 @@ fn test_feature_gating_with_zk_on() {
 
     //
     // ZK & !ZKless
-    let (mut h, recipient) = init_feature_gating(
+    let (mut h, recipient, _) = init_feature_gating(
         vec![
             FeatureFlag::CRYPTOGRAPHY_ALGEBRA_NATIVES,
             FeatureFlag::BN254_STRUCTURES,
@@ -113,7 +172,7 @@ fn test_feature_gating_with_zk_on() {
 fn test_feature_gating_with_zk_off() {
     //
     // !ZK & ZKless
-    let (mut h, recipient) = init_feature_gating(
+    let (mut h, recipient, _) = init_feature_gating(
         vec![
             FeatureFlag::CRYPTOGRAPHY_ALGEBRA_NATIVES,
             FeatureFlag::BN254_STRUCTURES,
@@ -128,7 +187,7 @@ fn test_feature_gating_with_zk_off() {
 
     //
     // !ZK & !ZKless
-    let (mut h, recipient) = init_feature_gating(
+    let (mut h, recipient, _) = init_feature_gating(
         vec![
             FeatureFlag::CRYPTOGRAPHY_ALGEBRA_NATIVES,
             FeatureFlag::BN254_STRUCTURES,
@@ -213,7 +272,7 @@ fn create_and_spend_keyless_account(
     spend_keyless_account(h, sig, &account, recipient)
 }
 
-fn run_setup_script(h: &mut MoveHarness) {
+fn run_jwk_and_config_script(h: &mut MoveHarness) -> Account {
     let core_resources = h.new_account_at(AccountAddress::from_hex_literal("0xA550C18").unwrap());
 
     let package = build_package(
@@ -249,4 +308,38 @@ fn run_setup_script(h: &mut MoveHarness) {
     // because it does not (yet) work with resource groups.
 
     assert_success!(h.run(txn));
+
+    core_resources
+}
+
+fn run_upgrade_vk_script(h: &mut MoveHarness, core_resources: Account, vk: Groth16VerificationKey) {
+    let package = build_package(
+        common::test_dir_path("keyless_new_vk.data/pack"),
+        aptos_framework::BuildOptions::default(),
+    )
+    .expect("building package must succeed");
+
+    let txn = h.create_publish_built_package(&core_resources, &package, |_| {});
+    assert_success!(h.run(txn));
+
+    let script = package.extract_script_code()[0].clone();
+
+    let txn = TransactionBuilder::new(core_resources.clone())
+        .script(Script::new(script, vec![], vec![
+            TransactionArgument::U8Vector(vk.alpha_g1),
+            TransactionArgument::U8Vector(vk.beta_g2),
+            TransactionArgument::U8Vector(vk.gamma_g2),
+            TransactionArgument::U8Vector(vk.delta_g2),
+            TransactionArgument::U8Vector(vk.gamma_abc_g1[0].clone()),
+            TransactionArgument::U8Vector(vk.gamma_abc_g1[1].clone()),
+        ]))
+        .sequence_number(h.sequence_number(core_resources.address()))
+        .max_gas_amount(1_000_000)
+        .gas_unit_price(1)
+        .sign();
+
+    // NOTE: We cannot write the Groth16Verification key via MoveHarness::set_resource
+    // because it does not (yet) work with resource groups.
+
+    assert_success!(h.run(txn));
 }

aptos-move/e2e-move-tests/src/tests/keyless_new_vk.data/pack/Move.toml

@@ -0,0 +1,6 @@
+[package]
+name = 'UpdateGroth16VerificationKey'
+version = "0.0.0"
+
+[dependencies]
+AptosFramework = { local = "../../../../../framework/aptos-framework" }

aptos-move/e2e-move-tests/src/tests/keyless_new_vk.data/pack/sources/main.move

@@ -0,0 +1,20 @@
+script {
+    use aptos_framework::aptos_governance;
+    use aptos_framework::keyless_account;
+
+    fun main(
+        core_resources: &signer,
+         alpha_g1: vector<u8>,
+         beta_g2: vector<u8>,
+         gamma_g2: vector<u8>,
+         delta_g2: vector<u8>,
+         gamma_abc_g1_0: vector<u8>,
+         gamma_abc_g1_1: vector<u8>
+    ) {
+        let vk = keyless_account::new_groth16_verification_key(alpha_g1, beta_g2, gamma_g2, delta_g2, vector[gamma_abc_g1_0, gamma_abc_g1_1]);
+        let fx = aptos_governance::get_signer_testnet_only(core_resources, @aptos_framework);
+        keyless_account::update_groth16_verification_key(&fx, vk);
+        // sets the pending Configuration change to the max expiration horizon from above
+        aptos_governance::force_end_epoch_test_only(core_resources);
+    }
+}

aptos-move/e2e-move-tests/src/tests/keyless_setup.data/pack/Move.toml

@@ -1,5 +1,5 @@
 [package]
-name = 'InsertJwk'
+name = 'InsertJwkAndUpdateMaxExpHorizon'
 version = "0.0.0"
 
 [dependencies]

aptos-move/e2e-move-tests/src/tests/keyless_setup.data/pack/sources/main.move

@@ -19,11 +19,8 @@ script {
         ];
         jwks::set_patches(&fx, patches);
 
-        // sets the pending VK and Configuration changes
-        keyless_account::on_new_epoch_for_testing(&fx);
-
         keyless_account::update_max_exp_horizon(&fx, max_exp_horizon_secs);
         // sets the pending Configuration change to the max expiration horizon from above
-        keyless_account::on_new_epoch_for_testing(&fx);
+        aptos_governance::force_end_epoch_test_only(core_resources);
     }
 }