Add solana-test-validator --upgradeable-program (#30412)

* Add TestValidator handling for upgradeable programs

* Plumb --upgradeable-program for solana-test-validator

Cargo.lock

@@ -6765,6 +6765,7 @@ name = "solana-test-validator"
 version = "1.16.0"
 dependencies = [
  "base64 0.13.0",
+ "bincode",
  "log",
  "serde_derive",
  "serde_json",

programs/sbf/Cargo.lock

@@ -6042,6 +6042,7 @@ name = "solana-test-validator"
 version = "1.16.0"
 dependencies = [
  "base64 0.13.0",
+ "bincode",
  "log",
  "serde_derive",
  "serde_json",

test-validator/Cargo.toml

@@ -11,6 +11,7 @@ edition = { workspace = true }
 
 [dependencies]
 base64 = { workspace = true }
+bincode = { workspace = true }
 log = { workspace = true }
 serde_derive = { workspace = true }
 serde_json = { workspace = true }

test-validator/src/lib.rs

@@ -76,6 +76,14 @@ pub struct ProgramInfo {
     pub program_path: PathBuf,
 }
 
+#[derive(Clone)]
+pub struct UpgradeableProgramInfo {
+    pub program_id: Pubkey,
+    pub loader: Pubkey,
+    pub upgrade_authority: Pubkey,
+    pub program_path: PathBuf,
+}
+
 #[derive(Debug)]
 pub struct TestValidatorNodeConfig {
     gossip_addr: SocketAddr,
@@ -111,6 +119,7 @@ pub struct TestValidatorGenesis {
     no_bpf_jit: bool,
     accounts: HashMap<Pubkey, AccountSharedData>,
     programs: Vec<ProgramInfo>,
+    upgradeable_programs: Vec<UpgradeableProgramInfo>,
     ticks_per_slot: Option<u64>,
     epoch_schedule: Option<EpochSchedule>,
     node_config: TestValidatorNodeConfig,
@@ -142,6 +151,7 @@ impl Default for TestValidatorGenesis {
             no_bpf_jit: bool::default(),
             accounts: HashMap::<Pubkey, AccountSharedData>::default(),
             programs: Vec::<ProgramInfo>::default(),
+            upgradeable_programs: Vec::<UpgradeableProgramInfo>::default(),
             ticks_per_slot: Option::<u64>::default(),
             epoch_schedule: Option::<EpochSchedule>::default(),
             node_config: TestValidatorNodeConfig::default(),
@@ -488,6 +498,17 @@ impl TestValidatorGenesis {
         self
     }
 
+    /// Add a list of upgradeable programs to the test environment.
+    pub fn add_upgradeable_programs_with_path(
+        &mut self,
+        programs: &[UpgradeableProgramInfo],
+    ) -> &mut Self {
+        for program in programs {
+            self.upgradeable_programs.push(program.clone());
+        }
+        self
+    }
+
     /// Start a test validator with the address of the mint account that will receive tokens
     /// created at genesis.
     ///
@@ -672,6 +693,44 @@ impl TestValidator {
                 }),
             );
         }
+        for upgradeable_program in &config.upgradeable_programs {
+            let data = solana_program_test::read_file(&upgradeable_program.program_path);
+            let (programdata_address, _) = Pubkey::find_program_address(
+                &[upgradeable_program.program_id.as_ref()],
+                &upgradeable_program.loader,
+            );
+            let mut program_data = bincode::serialize(&UpgradeableLoaderState::ProgramData {
+                slot: 0,
+                upgrade_authority_address: Some(upgradeable_program.upgrade_authority),
+            })
+            .unwrap();
+            program_data.extend_from_slice(&data);
+            accounts.insert(
+                programdata_address,
+                AccountSharedData::from(Account {
+                    lamports: Rent::default().minimum_balance(program_data.len()).max(1),
+                    data: program_data,
+                    owner: upgradeable_program.loader,
+                    executable: true,
+                    rent_epoch: 0,
+                }),
+            );
+
+            let data = bincode::serialize(&UpgradeableLoaderState::Program {
+                programdata_address,
+            })
+            .unwrap();
+            accounts.insert(
+                upgradeable_program.program_id,
+                AccountSharedData::from(Account {
+                    lamports: Rent::default().minimum_balance(data.len()).max(1),
+                    data,
+                    owner: upgradeable_program.loader,
+                    executable: true,
+                    rent_epoch: 0,
+                }),
+            );
+        }
 
         let mut genesis_config = create_genesis_config_with_leader_ex(
             mint_lamports,

validator/src/bin/solana-test-validator.rs

@@ -213,6 +213,56 @@ fn main() {
         }
     }
 
+    let mut upgradeable_programs_to_load = vec![];
+    if let Some(values) = matches.values_of("upgradeable_program") {
+        let values: Vec<&str> = values.collect::<Vec<_>>();
+        for address_program_upgrade_authority in values.chunks(3) {
+            match address_program_upgrade_authority {
+                [address, program, upgrade_authority] => {
+                    let address = address
+                        .parse::<Pubkey>()
+                        .or_else(|_| read_keypair_file(address).map(|keypair| keypair.pubkey()))
+                        .unwrap_or_else(|err| {
+                            println!("Error: invalid address {address}: {err}");
+                            exit(1);
+                        });
+                    let upgrade_authority_address = if *upgrade_authority == "none" {
+                        Pubkey::default()
+                    } else {
+                        upgrade_authority
+                            .parse::<Pubkey>()
+                            .or_else(|_| {
+                                read_keypair_file(upgrade_authority).map(|keypair| keypair.pubkey())
+                            })
+                            .unwrap_or_else(|err| {
+                                println!(
+                                    "Error: invalid upgrade_authority {upgrade_authority}: {err}"
+                                );
+                                exit(1);
+                            })
+                    };
+
+                    let program_path = PathBuf::from(program);
+                    if !program_path.exists() {
+                        println!(
+                            "Error: program file does not exist: {}",
+                            program_path.display()
+                        );
+                        exit(1);
+                    }
+
+                    upgradeable_programs_to_load.push(UpgradeableProgramInfo {
+                        program_id: address,
+                        loader: solana_sdk::bpf_loader_upgradeable::id(),
+                        upgrade_authority: upgrade_authority_address,
+                        program_path,
+                    });
+                }
+                _ => unreachable!(),
+            }
+        }
+    }
+
     let mut accounts_to_load = vec![];
     if let Some(values) = matches.values_of("account") {
         let values: Vec<&str> = values.collect::<Vec<_>>();
@@ -409,6 +459,7 @@ fn main() {
         .bpf_jit(!matches.is_present("no_bpf_jit"))
         .rpc_port(rpc_port)
         .add_programs_with_path(&programs_to_load)
+        .add_upgradeable_programs_with_path(&upgradeable_programs_to_load)
         .add_accounts_from_json_files(&accounts_to_load)
         .unwrap_or_else(|e| {
             println!("Error: add_accounts_from_json_files failed: {e}");

validator/src/cli.rs

@@ -2021,6 +2021,20 @@ pub fn test_app<'a>(version: &'a str, default_args: &'a DefaultTestArgs) -> App<
                        First argument can be a pubkey string or path to a keypair",
                 ),
         )
+        .arg(
+            Arg::with_name("upgradeable_program")
+                .long("upgradeable-program")
+                .value_names(&["ADDRESS_OR_KEYPAIR", "SBF_PROGRAM.SO", "UPGRADE_AUTHORITY"])
+                .takes_value(true)
+                .number_of_values(3)
+                .multiple(true)
+                .help(
+                    "Add an upgradeable SBF program to the genesis configuration. \
+                       If the ledger already exists then this parameter is silently ignored. \
+                       First and third arguments can be a pubkey string or path to a keypair. \
+                       Upgrade authority set to \"none\" disables upgrades",
+                ),
+        )
         .arg(
             Arg::with_name("account")
                 .long("account")