exebench for keyless

aptos-move/vm-genesis/src/lib.rs

@@ -22,7 +22,8 @@ use aptos_types::{
     contract_event::{ContractEvent, ContractEventV1},
     executable::ModulePath,
     jwks::{
-        patch::{PatchJWKMoveStruct, PatchUpsertJWK},
+        jwk::{JWKMoveStruct, JWK},
+        patch::{IssuerJWK, PatchJWKMoveStruct, PatchUpsertJWK},
         secure_test_rsa_jwk,
     },
     keyless::{
@@ -109,6 +110,8 @@ pub struct GenesisConfiguration {
     pub initial_features_override: Option<Features>,
     pub randomness_config_override: Option<OnChainRandomnessConfig>,
     pub jwk_consensus_config_override: Option<OnChainJWKConsensusConfig>,
+    pub initial_jwks: Vec<IssuerJWK>,
+    pub keyless_groth16_vk_override: Option<Groth16VerificationKey>,
 }
 
 pub static GENESIS_KEYPAIR: Lazy<(Ed25519PrivateKey, Ed25519PublicKey)> = Lazy::new(|| {
@@ -304,7 +307,13 @@ pub fn encode_genesis_change_set(
         .unwrap_or_else(OnChainJWKConsensusConfig::default_for_genesis);
     initialize_jwk_consensus_config(&mut session, &module_storage, &jwk_consensus_config);
     initialize_jwks_resources(&mut session, &module_storage);
-    initialize_keyless_accounts(&mut session, &module_storage, chain_id);
+    initialize_keyless_accounts(
+        &mut session,
+        &module_storage,
+        chain_id,
+        genesis_config.initial_jwks.clone(),
+        genesis_config.keyless_groth16_vk_override.clone(),
+    );
     set_genesis_end(&mut session, &module_storage);
 
     // Reconfiguration should happen after all on-chain invocations.
@@ -676,6 +685,8 @@ fn initialize_keyless_accounts(
     session: &mut SessionExt,
     module_storage: &impl AptosModuleStorage,
     chain_id: ChainId,
+    mut initial_jwks: Vec<IssuerJWK>,
+    vk_override: Option<Groth16VerificationKey>,
 ) {
     let config = keyless::Configuration::new_for_devnet();
     exec_function(
@@ -690,7 +701,8 @@ fn initialize_keyless_accounts(
         ]),
     );
     if !chain_id.is_mainnet() {
-        let vk = Groth16VerificationKey::from(&*DEVNET_VERIFICATION_KEY);
+        let vk =
+            vk_override.unwrap_or_else(|| Groth16VerificationKey::from(&*DEVNET_VERIFICATION_KEY));
         exec_function(
             session,
             module_storage,
@@ -703,11 +715,24 @@ fn initialize_keyless_accounts(
             ]),
         );
 
-        let patch: PatchJWKMoveStruct = PatchUpsertJWK {
+        let additional_jwk_patch = IssuerJWK {
             issuer: get_sample_iss(),
-            jwk: secure_test_rsa_jwk().into(),
-        }
-        .into();
+            jwk: JWK::RSA(secure_test_rsa_jwk()),
+        };
+        initial_jwks.push(additional_jwk_patch);
+
+        let jwk_patches: Vec<PatchJWKMoveStruct> = initial_jwks
+            .into_iter()
+            .map(|issuer_jwk| {
+                let IssuerJWK { issuer, jwk } = issuer_jwk;
+                let upsert_patch = PatchUpsertJWK {
+                    issuer,
+                    jwk: JWKMoveStruct::from(jwk),
+                };
+                PatchJWKMoveStruct::from(upsert_patch)
+            })
+            .collect();
+
         exec_function(
             session,
             module_storage,
@@ -716,7 +741,7 @@ fn initialize_keyless_accounts(
             vec![],
             serialize_values(&vec![
                 MoveValue::Signer(CORE_CODE_ADDRESS),
-                MoveValue::Vector(vec![patch.as_move_value()]),
+                jwk_patches.as_move_value(),
             ]),
         );
     }
@@ -1229,6 +1254,8 @@ pub fn generate_test_genesis(
             initial_features_override: None,
             randomness_config_override: None,
             jwk_consensus_config_override: None,
+            initial_jwks: vec![],
+            keyless_groth16_vk_override: None,
         },
         &OnChainConsensusConfig::default_for_genesis(),
         &OnChainExecutionConfig::default_for_genesis(),
@@ -1279,6 +1306,8 @@ fn mainnet_genesis_config() -> GenesisConfiguration {
         initial_features_override: None,
         randomness_config_override: None,
         jwk_consensus_config_override: None,
+        initial_jwks: vec![],
+        keyless_groth16_vk_override: None,
     }
 }
 

crates/aptos-genesis/src/builder.rs

@@ -27,6 +27,8 @@ use aptos_keygen::KeyGen;
 use aptos_logger::prelude::*;
 use aptos_types::{
     chain_id::ChainId,
+    jwks::patch::IssuerJWK,
+    keyless::Groth16VerificationKey,
     on_chain_config::{
         Features, GasScheduleV2, OnChainConsensusConfig, OnChainExecutionConfig,
         OnChainJWKConsensusConfig, OnChainRandomnessConfig,
@@ -441,6 +443,8 @@ pub struct GenesisConfiguration {
     pub initial_features_override: Option<Features>,
     pub randomness_config_override: Option<OnChainRandomnessConfig>,
     pub jwk_consensus_config_override: Option<OnChainJWKConsensusConfig>,
+    pub initial_jwks: Vec<IssuerJWK>,
+    pub keyless_groth16_vk_override: Option<Groth16VerificationKey>,
 }
 
 pub type InitConfigFn = Arc<dyn Fn(usize, &mut NodeConfig, &mut NodeConfig) + Send + Sync>;
@@ -662,6 +666,8 @@ impl Builder {
             initial_features_override: None,
             randomness_config_override: None,
             jwk_consensus_config_override: None,
+            initial_jwks: vec![],
+            keyless_groth16_vk_override: None,
         };
         if let Some(init_genesis_config) = &self.init_genesis_config {
             (init_genesis_config)(&mut genesis_config);

crates/aptos-genesis/src/config.rs

@@ -6,6 +6,8 @@ use aptos_crypto::{bls12381, ed25519::Ed25519PublicKey, x25519};
 use aptos_types::{
     account_address::{AccountAddress, AccountAddressWithChecks},
     chain_id::ChainId,
+    jwks::patch::IssuerJWK,
+    keyless::Groth16VerificationKey,
     network_address::{DnsName, NetworkAddress, Protocol},
     on_chain_config::{OnChainConsensusConfig, OnChainExecutionConfig, OnChainJWKConsensusConfig},
     transaction::authenticator::AuthenticationKey,
@@ -75,7 +77,16 @@ pub struct Layout {
     pub on_chain_execution_config: OnChainExecutionConfig,
 
     /// An optional JWK consensus config to use, instead of `default_for_genesis()`.
+    #[serde(default)]
     pub jwk_consensus_config_override: Option<OnChainJWKConsensusConfig>,
+
+    /// JWKs to patch in genesis.
+    #[serde(default)]
+    pub initial_jwks: Vec<IssuerJWK>,
+
+    /// Keyless Groth16 verification key to install in genesis.
+    #[serde(default)]
+    pub keyless_groth16_vk_override: Option<Groth16VerificationKey>,
 }
 
 impl Layout {
@@ -116,6 +127,8 @@ impl Default for Layout {
             on_chain_consensus_config: OnChainConsensusConfig::default(),
             on_chain_execution_config: OnChainExecutionConfig::default_for_genesis(),
             jwk_consensus_config_override: None,
+            initial_jwks: vec![],
+            keyless_groth16_vk_override: None,
         }
     }
 }

crates/aptos-genesis/src/lib.rs

@@ -23,6 +23,8 @@ use aptos_storage_interface::DbReaderWriter;
 use aptos_temppath::TempPath;
 use aptos_types::{
     chain_id::ChainId,
+    jwks::patch::IssuerJWK,
+    keyless::Groth16VerificationKey,
     on_chain_config::{
         Features, GasScheduleV2, OnChainConsensusConfig, OnChainExecutionConfig,
         OnChainJWKConsensusConfig, OnChainRandomnessConfig,
@@ -76,6 +78,8 @@ pub struct GenesisInfo {
     pub initial_features_override: Option<Features>,
     pub randomness_config_override: Option<OnChainRandomnessConfig>,
     pub jwk_consensus_config_override: Option<OnChainJWKConsensusConfig>,
+    pub initial_jwks: Vec<IssuerJWK>,
+    pub keyless_groth16_vk_override: Option<Groth16VerificationKey>,
 }
 
 impl GenesisInfo {
@@ -115,6 +119,8 @@ impl GenesisInfo {
             initial_features_override: genesis_config.initial_features_override.clone(),
             randomness_config_override: genesis_config.randomness_config_override.clone(),
             jwk_consensus_config_override: genesis_config.jwk_consensus_config_override.clone(),
+            initial_jwks: genesis_config.initial_jwks.clone(),
+            keyless_groth16_vk_override: genesis_config.keyless_groth16_vk_override.clone(),
         })
     }
 
@@ -150,6 +156,8 @@ impl GenesisInfo {
                 initial_features_override: self.initial_features_override.clone(),
                 randomness_config_override: self.randomness_config_override.clone(),
                 jwk_consensus_config_override: self.jwk_consensus_config_override.clone(),
+                initial_jwks: self.initial_jwks.clone(),
+                keyless_groth16_vk_override: self.keyless_groth16_vk_override.clone(),
             },
             &self.consensus_config,
             &self.execution_config,

crates/aptos-genesis/src/mainnet.rs

@@ -143,6 +143,8 @@ impl MainnetGenesisInfo {
                 initial_features_override: self.initial_features_override.clone(),
                 randomness_config_override: self.randomness_config_override.clone(),
                 jwk_consensus_config_override: self.jwk_consensus_config_override.clone(),
+                initial_jwks: vec![],
+                keyless_groth16_vk_override: None,
             },
         )
     }

crates/aptos/src/genesis/mod.rs

@@ -260,6 +260,8 @@ pub fn fetch_mainnet_genesis_info(git_options: GitOptions) -> CliTypedResult<Mai
             initial_features_override: None,
             randomness_config_override: None,
             jwk_consensus_config_override: None,
+            initial_jwks: vec![],
+            keyless_groth16_vk_override: None,
         },
     )?)
 }
@@ -303,6 +305,8 @@ pub fn fetch_genesis_info(git_options: GitOptions) -> CliTypedResult<GenesisInfo
             initial_features_override: None,
             randomness_config_override: None,
             jwk_consensus_config_override: layout.jwk_consensus_config_override.clone(),
+            initial_jwks: layout.initial_jwks.clone(),
+            keyless_groth16_vk_override: layout.keyless_groth16_vk_override.clone(),
         },
     )?)
 }

crates/transaction-emitter-lib/src/emitter/local_account_generator.rs

@@ -6,7 +6,7 @@ use aptos_sdk::types::{
     AccountKey, EphemeralKeyPair, EphemeralPrivateKey, KeylessAccount, LocalAccount,
 };
 use aptos_transaction_generator_lib::ReliableTransactionSubmitter;
-use aptos_types::keyless::{Claims, OpenIdSig, Pepper, ZeroKnowledgeSig};
+use aptos_types::keyless::{Claims, Configuration, OpenIdSig, Pepper, ZeroKnowledgeSig};
 use async_trait::async_trait;
 use futures::StreamExt;
 use rand::rngs::StdRng;
@@ -136,6 +136,7 @@ impl LocalAccountGenerator for KeylessAccountGenerator {
             },
         };
 
+        let config = Configuration::new_for_devnet();
         for line in lines {
             let serialized_proof = line?;
             let zk_sig_bytes = hex::decode(serialized_proof)?;
@@ -154,6 +155,7 @@ impl LocalAccountGenerator for KeylessAccountGenerator {
                 &self.uid_val,
                 &self.jwt_header_json,
                 EphemeralKeyPair::new(
+                    &config,
                     esk,
                     self.epk_expiry_date_secs,
                     vec![0; OpenIdSig::EPK_BLINDER_NUM_BYTES],

execution/executor-benchmark/src/account_generator.rs

@@ -17,15 +17,15 @@ impl AccountGenerator {
     const SEED_ACCOUNTS_ROOT_SEED: u64 = u64::max_value();
     const USER_ACCOUNTS_ROOT_SEED: u64 = 0;
 
-    pub fn new_for_seed_accounts() -> Self {
-        Self::new(Self::SEED_ACCOUNTS_ROOT_SEED, 0)
+    pub fn new_for_seed_accounts(is_keyless: bool) -> Self {
+        Self::new(Self::SEED_ACCOUNTS_ROOT_SEED, 0, is_keyless)
     }
 
-    pub fn new_for_user_accounts(num_to_skip: u64) -> Self {
-        Self::new(Self::USER_ACCOUNTS_ROOT_SEED, num_to_skip)
+    pub fn new_for_user_accounts(num_to_skip: u64, is_keyless: bool) -> Self {
+        Self::new(Self::USER_ACCOUNTS_ROOT_SEED, num_to_skip, is_keyless)
     }
 
-    fn new(root_seed: u64, num_to_skip: u64) -> Self {
+    fn new(root_seed: u64, num_to_skip: u64, is_keyless: bool) -> Self {
         let mut root_rng = StdRng::seed_from_u64(root_seed);
         let num_rngs_to_skip = num_to_skip / Self::MAX_ACCOUNT_GEN_PER_RNG;
         for _ in 0..num_rngs_to_skip {
@@ -35,14 +35,20 @@ impl AccountGenerator {
         let mut active_rng_quota = Self::MAX_ACCOUNT_GEN_PER_RNG - active_rng_to_skip;
         let mut active_rng = StdRng::seed_from_u64(root_rng.next_u64());
         for _ in 0..active_rng_to_skip {
-            LocalAccount::generate(&mut active_rng);
+            LocalAccount::generate_for_testing(&mut active_rng, is_keyless);
         }
         let (sender, receiver) = mpsc::sync_channel(100 /* bound */);
 
         std::thread::Builder::new()
             .name("account_generator".to_string())
             .spawn(move || {
-                while sender.send(LocalAccount::generate(&mut active_rng)).is_ok() {
+                while sender
+                    .send(LocalAccount::generate_for_testing(
+                        &mut active_rng,
+                        is_keyless,
+                    ))
+                    .is_ok()
+                {
                     active_rng_quota -= 1;
                     if active_rng_quota == 0 {
                         active_rng = StdRng::seed_from_u64(root_rng.next_u64());

execution/executor-benchmark/src/db_generator.rs

@@ -16,9 +16,17 @@ use aptos_executor::{
     db_bootstrapper::{generate_waypoint, maybe_bootstrap},
 };
 use aptos_storage_interface::DbReaderWriter;
-use aptos_types::on_chain_config::Features;
+use aptos_types::{
+    jwks::{jwk::JWK, patch::IssuerJWK},
+    keyless::{
+        circuit_constants::TEST_GROTH16_KEYS,
+        test_utils::{get_sample_iss, get_sample_jwk},
+        Groth16VerificationKey,
+    },
+    on_chain_config::Features,
+};
 use aptos_vm::AptosVM;
-use std::{fs, path::Path};
+use std::{fs, path::Path, sync::Arc};
 
 pub fn create_db_with_accounts<V>(
     num_accounts: usize,
@@ -30,6 +38,7 @@ pub fn create_db_with_accounts<V>(
     enable_storage_sharding: bool,
     pipeline_config: PipelineConfig,
     init_features: Features,
+    is_keyless: bool,
 ) where
     V: TransactionBlockExecutor + 'static,
 {
@@ -59,6 +68,7 @@ pub fn create_db_with_accounts<V>(
         enable_storage_sharding,
         pipeline_config,
         init_features,
+        is_keyless,
     );
 }
 
@@ -68,7 +78,15 @@ pub(crate) fn bootstrap_with_genesis(
     init_features: Features,
 ) {
     let (config, _genesis_key) =
-        aptos_genesis::test_utils::test_config_with_custom_features(init_features);
+        aptos_genesis::test_utils::test_config_with_custom_onchain(Some(Arc::new(move |config| {
+            config.initial_features_override = Some(init_features.clone());
+            config.initial_jwks = vec![IssuerJWK {
+                issuer: get_sample_iss(),
+                jwk: JWK::RSA(get_sample_jwk()),
+            }];
+            config.keyless_groth16_vk_override =
+                Some(Groth16VerificationKey::from(&TEST_GROTH16_KEYS.prepared_vk));
+        })));
 
     let mut rocksdb_configs = RocksdbConfigs::default();
     rocksdb_configs.state_merkle_db_config.max_open_files = -1;