fix: if upgradeable_v3 program missing program data, treat as Closed

crates/litesvm/src/accounts_db.rs

@@ -20,8 +20,8 @@ use {
     solana_nonce as nonce,
     solana_program_runtime::{
         loaded_programs::{
-            LoadProgramMetrics, ProgramCacheEntry, ProgramCacheForTxBatch,
-            ProgramRuntimeEnvironments,
+            LoadProgramMetrics, ProgramCacheEntry, ProgramCacheEntryOwner, ProgramCacheEntryType,
+            ProgramCacheForTxBatch, ProgramRuntimeEnvironments,
         },
         sysvar_cache::SysvarCache,
     },
@@ -262,11 +262,13 @@ impl AccountsDb {
                 );
                 return Err(InstructionError::InvalidAccountData);
             };
-            let programdata_account =
-                self.get_account_ref(&programdata_address).ok_or_else(|| {
-                    error!("Program data account {programdata_address} not found");
-                    InstructionError::MissingAccount
-                })?;
+            let Some(programdata_account) = self.get_account_ref(&programdata_address) else {
+                return Ok(ProgramCacheEntry::new_tombstone(
+                    slot,
+                    ProgramCacheEntryOwner::LoaderV3,
+                    ProgramCacheEntryType::Closed,
+                ));
+            };
             let program_data = programdata_account.data();
             if let Some(programdata) =
                 program_data.get(UpgradeableLoaderState::size_of_programdata_metadata()..)

crates/litesvm/tests/upgradeable_loader.rs

@@ -0,0 +1,153 @@
+use {
+    bincode::{deserialize, serialize},
+    litesvm::LiteSVM,
+    solana_account::Account,
+    solana_address::{address, Address},
+    solana_clock::Clock,
+    solana_instruction::{account_meta::AccountMeta, Instruction},
+    solana_keypair::Keypair,
+    solana_loader_v3_interface::{
+        get_program_data_address, instruction::UpgradeableLoaderInstruction,
+        state::UpgradeableLoaderState,
+    },
+    solana_message::Message,
+    solana_native_token::LAMPORTS_PER_SOL,
+    solana_sdk_ids::bpf_loader_upgradeable,
+    solana_signer::Signer,
+    solana_transaction::Transaction,
+    std::path::PathBuf,
+};
+
+fn read_counter_program() -> Vec<u8> {
+    let mut so_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+    so_path.push("test_programs/target/deploy/counter.so");
+    std::fs::read(so_path).unwrap()
+}
+
+fn set_program_upgrade_authority(
+    svm: &mut LiteSVM,
+    program_id: Address,
+    authority: Address,
+) -> Address {
+    let programdata_address = get_program_data_address(&program_id);
+    let mut programdata_account = svm.get_account(&programdata_address).unwrap();
+    let metadata_len = UpgradeableLoaderState::size_of_programdata_metadata();
+    let metadata =
+        deserialize::<UpgradeableLoaderState>(&programdata_account.data[..metadata_len]).unwrap();
+    let slot = match metadata {
+        UpgradeableLoaderState::ProgramData { slot, .. } => slot,
+        other => panic!("expected ProgramData account, got {other:?}"),
+    };
+
+    let mut data = bincode::serialize(&UpgradeableLoaderState::ProgramData {
+        slot,
+        upgrade_authority_address: Some(authority),
+    })
+    .unwrap();
+    data.extend_from_slice(&programdata_account.data[metadata_len..]);
+    programdata_account.data = data;
+
+    svm.set_account(programdata_address, programdata_account)
+        .unwrap();
+    programdata_address
+}
+
+fn invoke_counter(
+    svm: &mut LiteSVM,
+    program_id: Address,
+    counter_address: Address,
+    payer: &Keypair,
+    deduper: u8,
+) {
+    let payer_address = payer.pubkey();
+    let tx = Transaction::new(
+        &[payer],
+        Message::new_with_blockhash(
+            &[Instruction {
+                program_id,
+                accounts: vec![AccountMeta::new(counter_address, false)],
+                data: vec![0, deduper],
+            }],
+            Some(&payer_address),
+            &svm.latest_blockhash(),
+        ),
+        svm.latest_blockhash(),
+    );
+    svm.send_transaction(tx).unwrap();
+}
+
+#[test_log::test]
+fn close_upgradeable_program_keeps_vm_usable() {
+    let authority_kp = Keypair::new();
+    let authority = authority_kp.pubkey();
+    let program_id = address!("GtdambwDgHWrDJdVPBkEHGhCwokqgAoch162teUjJse2");
+    let counter_address = address!("J39wvrFY2AkoAUCke5347RMNk3ditxZfVidoZ7U6Fguf");
+
+    let mut svm = LiteSVM::new();
+    svm.airdrop(&authority, LAMPORTS_PER_SOL).unwrap();
+    svm.add_program(program_id, &read_counter_program())
+        .unwrap();
+
+    let programdata_address = set_program_upgrade_authority(&mut svm, program_id, authority);
+    let original_program_account = svm.get_account(&program_id).unwrap();
+    let original_programdata_account = svm.get_account(&programdata_address).unwrap();
+
+    // confirm invoking program works at start
+    {
+        svm.set_account(
+            counter_address,
+            Account {
+                lamports: 5,
+                data: vec![0_u8; std::mem::size_of::<u32>()],
+                owner: program_id,
+                ..Default::default()
+            },
+        )
+        .unwrap();
+
+        invoke_counter(&mut svm, program_id, counter_address, &authority_kp, 0);
+        assert_eq!(
+            svm.get_account(&counter_address).unwrap().data,
+            1u32.to_le_bytes().to_vec()
+        );
+    }
+
+    let current_slot = svm.get_sysvar::<Clock>().slot;
+    svm.warp_to_slot(current_slot + 1);
+
+    // verify we can close the program
+    {
+        let close_ix = Instruction::new_with_bytes(
+            bpf_loader_upgradeable::id(),
+            &serialize(&UpgradeableLoaderInstruction::Close).unwrap(),
+            vec![
+                AccountMeta::new(programdata_address, false),
+                AccountMeta::new(authority, false),
+                AccountMeta::new_readonly(authority, true),
+                AccountMeta::new(program_id, false),
+            ],
+        );
+        let close_tx = Transaction::new(
+            &[&authority_kp],
+            Message::new_with_blockhash(&[close_ix], Some(&authority), &svm.latest_blockhash()),
+            svm.latest_blockhash(),
+        );
+        svm.send_transaction(close_tx).unwrap();
+
+        assert!(svm.get_account(&programdata_address).is_none());
+    }
+
+    // verify that if we directly write to the program data address again we can still invoke the program
+    {
+        svm.set_account(programdata_address, original_programdata_account)
+            .unwrap();
+        svm.set_account(program_id, original_program_account)
+            .unwrap();
+
+        invoke_counter(&mut svm, program_id, counter_address, &authority_kp, 1);
+        assert_eq!(
+            svm.get_account(&counter_address).unwrap().data,
+            2u32.to_le_bytes().to_vec()
+        );
+    }
+}