solana-genesis: add support for initializing upgradeable programs

fetch-spl.sh

@@ -6,6 +6,8 @@
 
 set -e
 
+upgradeableLoader=BPFLoaderUpgradeab1e11111111111111111111111
+
 fetch_program() {
   declare name=$1
   declare version=$2
@@ -14,7 +16,11 @@ fetch_program() {
 
   declare so=spl_$name-$version.so
 
-  genesis_args+=(--bpf-program "$address" "$loader" "$so")
+  if [[ $loader == "$upgradeableLoader" ]]; then
+    genesis_args+=(--upgradeable-program "$address" "$loader" "$so" none)
+  else
+    genesis_args+=(--bpf-program "$address" "$loader" "$so")
+  fi
 
   if [[ -r $so ]]; then
     return
@@ -39,7 +45,7 @@ fetch_program() {
 }
 
 fetch_program token 3.5.0 TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA BPFLoader2111111111111111111111111111111111
-fetch_program token-2022 0.5.0 TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb BPFLoader2111111111111111111111111111111111
+fetch_program token-2022 0.5.0 TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb BPFLoaderUpgradeab1e11111111111111111111111
 fetch_program memo  1.0.0 Memo1UhkJRfHyvLMcVucJwxXeuD728EqVDDwQDxFMNo BPFLoader1111111111111111111111111111111111
 fetch_program memo  3.0.0 MemoSq4gqABAXKb96qnH8TysNcWxMyWCqXgDLGmfcHr BPFLoader2111111111111111111111111111111111
 fetch_program associated-token-account 1.1.2 ATokenGPvbdGVxr1b2hvZbsiqW5xWH25efTNsLJA8knL BPFLoader2111111111111111111111111111111111

genesis/Cargo.toml

@@ -11,6 +11,7 @@ documentation = "https://docs.rs/solana-genesis"
 
 [dependencies]
 base64 = "0.13.0"
+bincode = "1.3.3"
 clap = "2.33.1"
 serde = "1.0.144"
 serde_json = "1.0.83"

genesis/src/main.rs

@@ -17,6 +17,7 @@ use {
     solana_runtime::hardened_unpack::MAX_GENESIS_ARCHIVE_UNPACKED_SIZE,
     solana_sdk::{
         account::{Account, AccountSharedData, ReadableAccount, WritableAccount},
+        bpf_loader_upgradeable::UpgradeableLoaderState,
         clock,
         epoch_schedule::EpochSchedule,
         fee_calculator::FeeRateGovernor,
@@ -27,6 +28,7 @@ use {
         pubkey::Pubkey,
         rent::Rent,
         signature::{Keypair, Signer},
+        signer::keypair::read_keypair_file,
         stake::state::StakeState,
         system_program, timing,
     },
@@ -376,6 +378,15 @@ fn main() -> Result<(), Box<dyn error::Error>> {
                 .multiple(true)
                 .help("Install a SBF program at the given address"),
         )
+        .arg(
+            Arg::with_name("upgradeable_program")
+                .long("upgradeable-program")
+                .value_name("ADDRESS UPGRADEABLE_LOADER BPF_PROGRAM.SO UPGRADE_AUTHORITY")
+                .takes_value(true)
+                .number_of_values(4)
+                .multiple(true)
+                .help("Install an upgradeable SBF program at the given address with the given upgrade authority (or \"none\")"),
+        )
         .arg(
             Arg::with_name("inflation")
                 .required(false)
@@ -585,28 +596,32 @@ fn main() -> Result<(), Box<dyn error::Error>> {
 
     add_genesis_accounts(&mut genesis_config, issued_lamports - faucet_lamports);
 
+    let parse_address = |address: &str, input_type: &str| {
+        address.parse::<Pubkey>().unwrap_or_else(|err| {
+            eprintln!("Error: invalid {input_type} {address}: {err}");
+            process::exit(1);
+        })
+    };
+
+    let parse_program_data = |program: &str| {
+        let mut program_data = vec![];
+        File::open(program)
+            .and_then(|mut file| file.read_to_end(&mut program_data))
+            .unwrap_or_else(|err| {
+                eprintln!("Error: failed to read {program}: {err}");
+                process::exit(1);
+            });
+        program_data
+    };
+
     if let Some(values) = matches.values_of("bpf_program") {
         let values: Vec<&str> = values.collect::<Vec<_>>();
         for address_loader_program in values.chunks(3) {
             match address_loader_program {
                 [address, loader, program] => {
-                    let address = address.parse::<Pubkey>().unwrap_or_else(|err| {
-                        eprintln!("Error: invalid address {address}: {err}");
-                        process::exit(1);
-                    });
-
-                    let loader = loader.parse::<Pubkey>().unwrap_or_else(|err| {
-                        eprintln!("Error: invalid loader {loader}: {err}");
-                        process::exit(1);
-                    });
-
-                    let mut program_data = vec![];
-                    File::open(program)
-                        .and_then(|mut file| file.read_to_end(&mut program_data))
-                        .unwrap_or_else(|err| {
-                            eprintln!("Error: failed to read {program}: {err}");
-                            process::exit(1);
-                        });
+                    let address = parse_address(address, "address");
+                    let loader = parse_address(loader, "loader");
+                    let program_data = parse_program_data(program);
                     genesis_config.add_account(
                         address,
                         AccountSharedData::from(Account {
@@ -623,6 +638,65 @@ fn main() -> Result<(), Box<dyn error::Error>> {
         }
     }
 
+    if let Some(values) = matches.values_of("upgradeable_program") {
+        let values: Vec<&str> = values.collect::<Vec<_>>();
+        for address_loader_program_upgrade_authority in values.chunks(4) {
+            match address_loader_program_upgrade_authority {
+                [address, loader, program, upgrade_authority] => {
+                    let address = parse_address(address, "address");
+                    let loader = parse_address(loader, "loader");
+                    let program_data_elf = parse_program_data(program);
+                    let upgrade_authority_address = if *upgrade_authority == "none" {
+                        Pubkey::default()
+                    } else {
+                        upgrade_authority.parse::<Pubkey>().unwrap_or_else(|_| {
+                          read_keypair_file(upgrade_authority).map(|keypair| keypair.pubkey()).unwrap_or_else(|err| {
+                              eprintln!("Error: invalid upgrade_authority {upgrade_authority}: {err}");
+                              process::exit(1);
+                          })
+                      })
+                    };
+
+                    let (programdata_address, _) =
+                        Pubkey::find_program_address(&[address.as_ref()], &loader);
+                    let mut program_data =
+                        bincode::serialize(&UpgradeableLoaderState::ProgramData {
+                            slot: 0,
+                            upgrade_authority_address: Some(upgrade_authority_address),
+                        })
+                        .unwrap();
+                    program_data.extend_from_slice(&program_data_elf);
+                    genesis_config.add_account(
+                        programdata_address,
+                        AccountSharedData::from(Account {
+                            lamports: genesis_config.rent.minimum_balance(program_data.len()),
+                            data: program_data,
+                            owner: loader,
+                            executable: false,
+                            rent_epoch: 0,
+                        }),
+                    );
+
+                    let program_data = bincode::serialize(&UpgradeableLoaderState::Program {
+                        programdata_address,
+                    })
+                    .unwrap();
+                    genesis_config.add_account(
+                        address,
+                        AccountSharedData::from(Account {
+                            lamports: genesis_config.rent.minimum_balance(program_data.len()),
+                            data: program_data,
+                            owner: loader,
+                            executable: true,
+                            rent_epoch: 0,
+                        }),
+                    );
+                }
+                _ => unreachable!(),
+            }
+        }
+    }
+
     solana_logger::setup();
     create_new_ledger(
         &ledger_path,