fix(katana): ensure paymaster accounts are unique

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<Outcome> {
             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<RefCell<Vec<PaymasterAccountArgs>>> = 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::<Felt>::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::<Felt>::new("Paymaster public key")
+        let salt_prompt_text = format!("Paymaster #{} salt", paymaster_count);
+        let salt = CustomType::<Felt>::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<Outcome> {
         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()),
     })
 }

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,25 +14,48 @@ 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<Vec<PaymasterAccountArgs>>,
 }
 
 #[derive(Debug, Clone)]
 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<Self> {
-        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());
+    }
 }

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<DevGenesisAccount> for GenesisAllocation {