From 3d30b1967671a850a89728dec195c04af3b142ea Mon Sep 17 00:00:00 2001 From: Ammar Arif Date: Wed, 12 Feb 2025 23:05:21 -0500 Subject: [PATCH 1/2] prevent paymaster to have similar pubkey and salt combo --- bin/katana/src/cli/init/prompt.rs | 45 ++++++- bin/katana/src/cli/init/slot.rs | 123 +++++++++++++++++- .../primitives/src/genesis/allocation.rs | 15 +++ 3 files changed, 172 insertions(+), 11 deletions(-) diff --git a/bin/katana/src/cli/init/prompt.rs b/bin/katana/src/cli/init/prompt.rs index 1ebec48487..17c3271a42 100644 --- a/bin/katana/src/cli/init/prompt.rs +++ b/bin/katana/src/cli/init/prompt.rs @@ -1,7 +1,10 @@ +use std::cell::RefCell; +use std::rc::Rc; use std::str::FromStr; use std::sync::Arc; use anyhow::{Context, Result}; +use inquire::validator::{ErrorMessage, Validation}; use inquire::{Confirm, CustomType, Select}; use katana_primitives::block::BlockNumber; use katana_primitives::{ContractAddress, Felt}; @@ -15,6 +18,7 @@ use tokio::runtime::Handle; use super::{deployment, Outcome}; use crate::cli::init::deployment::DeploymentOutcome; +use crate::cli::init::slot::{self, PaymasterAccountArgs}; pub const CARTRIDGE_SN_SEPOLIA_PROVIDER: &str = "https://api.cartridge.gg/x/starknet/sepolia"; @@ -125,15 +129,46 @@ pub async fn prompt() -> Result { DeploymentOutcome { contract_address: address, block_number } }; + // It's wrapped like this because the prompt validator requires captured variables to have + // 'static lifetime. + let slot_paymasters: Rc>> = Default::default(); + let mut paymaster_count = 1; + // Prompt for slot paymaster accounts - let mut slot_paymasters = Vec::new(); + while Confirm::new("Add Slot paymaster account?").with_default(true).prompt()? { + let pubkey_prompt_text = format!("Paymaster #{} public key", paymaster_count); + let public_key = CustomType::::new(&pubkey_prompt_text) + .with_formatter(&|input: Felt| format!("{input:#x}")) + .prompt()?; + + // Check if this public_key + salt combo already exists + // This check is necessary to ensure that each paymaster account has a unique addresses + // because the contract address is derived from the public key and salt. So, if + // there multiple paymasters with the same public key and salt pair, then + // the resultant contract address will be the same. + let slot_paymasters_clone = slot_paymasters.clone(); + let unique_salt_validator = move |salt: &Felt| { + let pred = |pm: &PaymasterAccountArgs| pm.public_key == public_key && pm.salt == *salt; + let duplicate = slot_paymasters_clone.borrow().iter().any(pred); + + if !duplicate { + Ok(Validation::Valid) + } else { + Ok(Validation::Invalid(ErrorMessage::Custom( + "Public key and salt combination already exists!".to_string(), + ))) + } + }; - while Confirm::new("Add slot paymaster account?").with_default(true).prompt()? { - let public_key = CustomType::::new("Paymaster public key") + let salt_prompt_text = format!("Paymaster #{} salt", paymaster_count); + let salt = CustomType::::new(&salt_prompt_text) .with_formatter(&|input: Felt| format!("{input:#x}")) + .with_validator(unique_salt_validator) + .with_default(Felt::ONE) .prompt()?; - slot_paymasters.push(super::slot::PaymasterAccountArgs { public_key }); + slot_paymasters.borrow_mut().push(slot::PaymasterAccountArgs { public_key, salt }); + paymaster_count += 1; } Ok(Outcome { @@ -143,6 +178,6 @@ pub async fn prompt() -> Result { account: account_address, settlement_id: parse_cairo_short_string(&l1_chain_id)?, #[cfg(feature = "init-slot")] - slot_paymasters: Some(slot_paymasters), + slot_paymasters: Some(Rc::unwrap_or_clone(slot_paymasters).take()), }) } diff --git a/bin/katana/src/cli/init/slot.rs b/bin/katana/src/cli/init/slot.rs index a32ec35d24..f6119868a4 100644 --- a/bin/katana/src/cli/init/slot.rs +++ b/bin/katana/src/cli/init/slot.rs @@ -1,6 +1,6 @@ use std::str::FromStr; -use anyhow::Result; +use anyhow::{anyhow, Result}; use clap::Args; use katana_primitives::genesis::allocation::{ GenesisAccount, GenesisAccountAlloc, GenesisAllocation, @@ -14,11 +14,24 @@ use katana_primitives::{ContractAddress, Felt, U256}; #[derive(Debug, Args)] #[command(next_help_heading = "Slot options")] pub struct SlotArgs { + /// Enable `slot`-specific features. #[arg(long)] pub slot: bool, + /// Specify the number of paymaster accounts to create. + /// + /// This argument accepts a list of values, where each value is a pair of public key and salt + /// separated by a comma. + /// + /// For example: + /// + /// ``` + /// --slot.paymasters 0x1,0x2 0x3,0x4 0x5,0x6 + /// ``` + /// + /// where the total number of pairs determine how many paymaster accounts will be created. #[arg(requires_all = ["id", "slot"])] - #[arg(long = "slot.paymasters", value_delimiter = ',')] + #[arg(long = "slot.paymasters", value_delimiter = ' ')] pub paymaster_accounts: Option>, } @@ -26,13 +39,23 @@ pub struct SlotArgs { pub struct PaymasterAccountArgs { /// The public key of the paymaster account. pub public_key: Felt, + /// The salt of the paymaster account. + pub salt: Felt, } impl FromStr for PaymasterAccountArgs { type Err = anyhow::Error; fn from_str(s: &str) -> Result { - Ok(PaymasterAccountArgs { public_key: Felt::from_str(s)? }) + let mut parts = s.split(','); + + let public_key = parts.next().ok_or_else(|| anyhow!("missing public key"))?; + let salt = parts.next().ok_or_else(|| anyhow!("missing salt"))?; + + let public_key = Felt::from_str(public_key)?; + let salt = Felt::from_str(salt)?; + + Ok(PaymasterAccountArgs { public_key, salt }) } } @@ -43,11 +66,16 @@ pub fn add_paymasters_to_genesis( let mut accounts = Vec::with_capacity(slot_paymasters.len()); for paymaster in slot_paymasters { - let public_key = paymaster.public_key; let class_hash = DEFAULT_ACCOUNT_CLASS_HASH; let balance = U256::from(DEFAULT_PREFUNDED_ACCOUNT_BALANCE); - let (addr, account) = GenesisAccount::new_with_balance(public_key, class_hash, balance); + let (addr, account) = GenesisAccount::new_with_salt_and_balance( + paymaster.public_key, + class_hash, + paymaster.salt, + balance, + ); + let account = GenesisAllocation::Account(GenesisAccountAlloc::Account(account)); accounts.push((addr, account)); } @@ -60,6 +88,8 @@ pub fn add_paymasters_to_genesis( #[cfg(test)] mod tests { + use clap::Parser; + use super::*; #[test] @@ -68,7 +98,8 @@ mod tests { let mut paymasters = Vec::new(); for i in 0..3 { - paymasters.push(PaymasterAccountArgs { public_key: Felt::from(i) }); + paymasters + .push(PaymasterAccountArgs { public_key: Felt::from(i), salt: Felt::from(i) }); } let expected_addresses = add_paymasters_to_genesis(&mut genesis, &paymasters); @@ -89,4 +120,84 @@ mod tests { } } } + + #[test] + fn test_distinct_paymasters_same_pubkey() { + let mut genesis = Genesis::default(); + let mut paymasters = Vec::new(); + let public_key = Felt::from(1); + + // Add multiple paymasters with same public key + for i in 0..3 { + let salt = Felt::from(i); + paymasters.push(PaymasterAccountArgs { public_key, salt }); + } + + let addresses = add_paymasters_to_genesis(&mut genesis, &paymasters); + + // Verify addresses are unique + let mut unique_addresses = addresses.clone(); + unique_addresses.sort(); + unique_addresses.dedup(); + + assert_eq!(addresses.len(), unique_addresses.len(), "addresses are not unique"); + + // Verify each paymaster has the same public key + for addr in addresses { + let account = genesis.allocations.get(&addr).expect("account missing"); + match account { + GenesisAllocation::Account(GenesisAccountAlloc::Account(account)) => { + assert_eq!(account.public_key, public_key); + } + _ => panic!("Expected GenesisAccountAlloc::Account"), + } + } + } + + #[test] + fn test_parse_no_paymaster_args() { + #[derive(Parser)] + struct Cli { + #[arg(long)] + id: bool, + #[command(flatten)] + slot: SlotArgs, + } + + let Cli { slot, .. } = Cli::parse_from(["cli", "--id", "--slot"]); + assert!(slot.paymaster_accounts.is_none()); + } + + #[test] + fn test_parse_paymaster_args() { + #[derive(Parser)] + struct Cli { + #[arg(long)] + id: bool, + #[command(flatten)] + slot: SlotArgs, + } + + let Cli { slot, .. } = Cli::parse_from([ + "cli", + "--id", + "--slot", + "--slot.paymasters", + "0x1,0x2", + "0x1,0x3", + "0x1,0x4", + ]); + + let paymasters = slot.paymaster_accounts.unwrap(); + assert_eq!(paymasters.len(), 3); + + assert_eq!(paymasters[0].public_key, Felt::from_str("0x1").unwrap()); + assert_eq!(paymasters[0].salt, Felt::from_str("0x2").unwrap()); + + assert_eq!(paymasters[1].public_key, Felt::from_str("0x1").unwrap()); + assert_eq!(paymasters[1].salt, Felt::from_str("0x3").unwrap()); + + assert_eq!(paymasters[2].public_key, Felt::from_str("0x1").unwrap()); + assert_eq!(paymasters[2].salt, Felt::from_str("0x4").unwrap()); + } } diff --git a/crates/katana/primitives/src/genesis/allocation.rs b/crates/katana/primitives/src/genesis/allocation.rs index b63d878bb6..b9159b08ad 100644 --- a/crates/katana/primitives/src/genesis/allocation.rs +++ b/crates/katana/primitives/src/genesis/allocation.rs @@ -210,6 +210,21 @@ impl GenesisAccount { let (address, account) = Self::new(public_key, class_hash); (address, Self { balance: Some(balance), ..account }) } + + pub fn new_with_salt_and_balance( + public_key: Felt, + class_hash: ClassHash, + salt: Felt, + balance: U256, + ) -> (ContractAddress, Self) { + let (address, account) = Self::new_inner(public_key, class_hash, salt); + (address, Self { balance: Some(balance), ..account }) + } + + fn new_inner(public_key: Felt, class_hash: ClassHash, salt: Felt) -> (ContractAddress, Self) { + let address = get_contract_address(salt, class_hash, &[public_key], Felt::ZERO); + (ContractAddress::from(address), Self { public_key, class_hash, ..Default::default() }) + } } impl From for GenesisAllocation { From 88588689c0ff8eecac41d78ebf058cfe2ac4af1a Mon Sep 17 00:00:00 2001 From: glihm Date: Wed, 12 Feb 2025 22:31:17 -0600 Subject: [PATCH 2/2] adjust tests and doc for multiple paymasters --- bin/katana/src/cli/init/slot.rs | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/bin/katana/src/cli/init/slot.rs b/bin/katana/src/cli/init/slot.rs index f6119868a4..573b072a75 100644 --- a/bin/katana/src/cli/init/slot.rs +++ b/bin/katana/src/cli/init/slot.rs @@ -21,12 +21,17 @@ pub struct SlotArgs { /// Specify the number of paymaster accounts to create. /// /// This argument accepts a list of values, where each value is a pair of public key and salt - /// separated by a comma. + /// separated by a comma. If more than one pair is provided, the double quotes are required to + /// prevent the shell from prematurely splitting the argument. /// /// For example: /// /// ``` - /// --slot.paymasters 0x1,0x2 0x3,0x4 0x5,0x6 + /// --slot.paymasters 0x1,0x2 + /// ``` + /// + /// ``` + /// --slot.paymasters "0x1,0x2 0x3,0x4 0x5,0x6" /// ``` /// /// where the total number of pairs determine how many paymaster accounts will be created. @@ -183,9 +188,8 @@ mod tests { "--id", "--slot", "--slot.paymasters", - "0x1,0x2", - "0x1,0x3", - "0x1,0x4", + // Must be in one argument because the shell splits the argument otherwise. + "0x1,0x2 0x1,0x3 0x1,0x4", ]); let paymasters = slot.paymaster_accounts.unwrap();