fix: CUs discrepancy (#299)

* fix: CUs discrepancy

* sync programs mainnet loader

---------

Co-authored-by: Aursen <aursen@users.noreply.github.com>

Author: Aursen

crates/litesvm/src/error.rs

@@ -30,4 +30,6 @@ pub enum LiteSVMError {
     Instruction(#[from] InstructionError),
     #[error("{0}")]
     InvalidPath(#[from] std::io::Error),
+    #[error("{0}")]
+    InvalidLoader(String),
 }

crates/litesvm/src/lib.rs

@@ -332,6 +332,7 @@ use {
     solana_hash::Hash,
     solana_keypair::Keypair,
     solana_last_restart_slot::LastRestartSlot,
+    solana_loader_v3_interface::state::UpgradeableLoaderState,
     solana_message::{
         inner_instruction::InnerInstructionsList, Message, SanitizedMessage, VersionedMessage,
     },
@@ -342,7 +343,9 @@ use {
         loaded_programs::{LoadProgramMetrics, ProgramCacheEntry},
     },
     solana_rent::Rent,
-    solana_sdk_ids::{bpf_loader, native_loader, system_program},
+    solana_sdk_ids::{
+        bpf_loader, bpf_loader_deprecated, bpf_loader_upgradeable, native_loader, system_program,
+    },
     solana_signature::Signature,
     solana_signer::Signer,
     solana_slot_hashes::SlotHashes,
@@ -826,62 +829,129 @@ impl LiteSVM {
         Ok(())
     }
 
-    /// Adds am SBF program to the test environment.
     fn add_program_internal<const PREVERIFIED: bool>(
         &mut self,
         program_id: impl Into<Address>,
         program_bytes: &[u8],
+        loader_id: &Address,
     ) -> Result<(), LiteSVMError> {
         let program_id = program_id.into();
-        let program_len = program_bytes.len();
-        let lamports = self.minimum_balance_for_rent_exemption(program_len);
-        let mut account = AccountSharedData::new(lamports, program_len, &bpf_loader::id());
-        account.set_executable(true);
-        account.set_data_from_slice(program_bytes);
         let current_slot = self
             .accounts
             .sysvar_cache
             .get_clock()
             .unwrap_or_default()
             .slot;
+
+        let program_size = if bpf_loader_upgradeable::check_id(loader_id) {
+            let (programdata_address, _bump) =
+                Address::find_program_address(&[program_id.as_ref()], loader_id);
+
+            let programdata_metadata_len = UpgradeableLoaderState::size_of_programdata_metadata();
+            let programdata_len = programdata_metadata_len + program_bytes.len();
+            let mut programdata_data = vec![0u8; programdata_len];
+
+            bincode::serialize_into(
+                &mut programdata_data[..programdata_metadata_len],
+                &UpgradeableLoaderState::ProgramData {
+                    slot: current_slot,
+                    upgrade_authority_address: None,
+                },
+            )
+            .expect("UpgradeableLoaderState::ProgramData serialization should never fail");
+            programdata_data[programdata_metadata_len..].copy_from_slice(program_bytes);
+
+            let programdata_lamports = self.minimum_balance_for_rent_exemption(programdata_len);
+            let mut programdata_account =
+                AccountSharedData::new(programdata_lamports, programdata_len, loader_id);
+            programdata_account.set_data_from_slice(&programdata_data);
+
+            let program_account_data = bincode::serialize(&UpgradeableLoaderState::Program {
+                programdata_address,
+            })
+            .expect("UpgradeableLoaderState::Program serialization should never fail");
+
+            let program_lamports =
+                self.minimum_balance_for_rent_exemption(program_account_data.len());
+            let mut program_account =
+                AccountSharedData::new(program_lamports, program_account_data.len(), loader_id);
+            program_account.set_executable(true);
+            program_account.set_data_from_slice(&program_account_data);
+
+            self.accounts
+                .add_account_no_checks(programdata_address, programdata_account);
+            self.accounts
+                .add_account_no_checks(program_id, program_account);
+
+            programdata_len
+        } else if bpf_loader::check_id(loader_id) || bpf_loader_deprecated::check_id(loader_id) {
+            let program_len = program_bytes.len();
+            let lamports = self.minimum_balance_for_rent_exemption(program_len);
+            let mut account = AccountSharedData::new(lamports, program_len, loader_id);
+            account.set_executable(true);
+            account.set_data_from_slice(program_bytes);
+
+            self.accounts.add_account_no_checks(program_id, account);
+
+            program_len
+        } else {
+            return Err(LiteSVMError::InvalidLoader(format!(
+                "Unsupported loader: {loader_id}"
+            )));
+        };
+
         let mut loaded_program = solana_bpf_loader_program::load_program_from_bytes(
             None,
             &mut LoadProgramMetrics::default(),
-            account.data(),
-            account.owner(),
-            account.data().len(),
+            program_bytes,
+            loader_id,
+            program_size,
             current_slot,
             self.accounts.environments.program_runtime_v1.clone(),
             PREVERIFIED,
         )
         .map_err(LiteSVMError::from)?;
         loaded_program.effective_slot = current_slot;
 
-        // We already loaded and validated (or explicitly trusted) the executable above.
-        // Insert the account directly to avoid a second program load.
-        self.accounts.add_account_no_checks(program_id, account);
         self.accounts
             .programs_cache
             .replenish(program_id, Arc::new(loaded_program));
+
         Ok(())
     }
 
     /// Adds an SBF program to the test environment.
+    ///
+    /// Uses `BPFLoaderUpgradeable` by default for the loader.
     pub fn add_program(
         &mut self,
         program_id: impl Into<Address>,
         program_bytes: &[u8],
     ) -> Result<(), LiteSVMError> {
-        self.add_program_internal::<false>(program_id, program_bytes)
+        self.add_program_internal::<false>(program_id, program_bytes, &bpf_loader_upgradeable::id())
+    }
+
+    /// Adds an SBF program with a specific loader to match mainnet CU behavior.
+    ///
+    /// Use `bpf_loader::id()` for BPFLoader2, `bpf_loader_deprecated::id()` for BPFLoader1,
+    /// or `bpf_loader_upgradeable::id()` for the upgradeable loader.
+    pub fn add_program_with_loader(
+        &mut self,
+        program_id: impl Into<Address>,
+        program_bytes: &[u8],
+        loader_id: Address,
+    ) -> Result<(), LiteSVMError> {
+        self.add_program_internal::<false>(program_id, program_bytes, &loader_id)
     }
 
     /// Adds an SBF program that is known-good and already verified.
     pub(crate) fn add_program_preverified(
         &mut self,
         program_id: impl Into<Address>,
         program_bytes: &[u8],
+        loader_id: &Address,
     ) -> Result<(), LiteSVMError> {
-        self.add_program_internal::<true>(program_id, program_bytes)
+        self.add_program_internal::<true>(program_id, program_bytes, loader_id)
     }
 
     fn create_transaction_context(

crates/litesvm/src/programs/mod.rs

@@ -1,45 +1,58 @@
 use {
     crate::LiteSVM,
     solana_address::address,
-    solana_sdk_ids::{address_lookup_table, config},
+    solana_sdk_ids::{
+        address_lookup_table, bpf_loader, bpf_loader_deprecated, bpf_loader_upgradeable, config,
+    },
 };
 
 pub fn load_default_programs(svm: &mut LiteSVM) {
     svm.add_program_preverified(
         address!("TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"),
         include_bytes!("elf/spl_token-3.5.0.so"),
+        &bpf_loader::id(),
     )
     .unwrap();
     svm.add_program_preverified(
         address!("TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb"),
         include_bytes!("elf/spl_token_2022-10.0.0.so"),
+        &bpf_loader_upgradeable::id(),
     )
     .unwrap();
     svm.add_program_preverified(
         address!("Memo1UhkJRfHyvLMcVucJwxXeuD728EqVDDwQDxFMNo"),
         include_bytes!("elf/spl_memo-1.0.0.so"),
+        &bpf_loader_deprecated::id(),
     )
     .unwrap();
     svm.add_program_preverified(
         address!("MemoSq4gqABAXKb96qnH8TysNcWxMyWCqXgDLGmfcHr"),
         include_bytes!("elf/spl_memo-3.0.0.so"),
+        &bpf_loader::id(),
     )
     .unwrap();
     svm.add_program_preverified(
         address!("ATokenGPvbdGVxr1b2hvZbsiqW5xWH25efTNsLJA8knL"),
         include_bytes!("elf/spl_associated_token_account-1.1.1.so"),
+        &bpf_loader::id(),
+    )
+    .unwrap();
+    svm.add_program_preverified(
+        config::ID,
+        include_bytes!("elf/config.so"),
+        &bpf_loader_upgradeable::id(),
     )
     .unwrap();
-    svm.add_program_preverified(config::ID, include_bytes!("elf/config.so"))
-        .unwrap();
     svm.add_program_preverified(
         address_lookup_table::ID,
         include_bytes!("elf/address_lookup_table.so"),
+        &bpf_loader_upgradeable::id(),
     )
     .unwrap();
     svm.add_program_preverified(
         address!("Stake11111111111111111111111111111111111111"),
         include_bytes!("elf/core_bpf_stake-1.0.1.so"),
+        &bpf_loader_upgradeable::id(),
     )
     .unwrap()
 }

crates/loader/Cargo.toml

@@ -10,14 +10,14 @@ rust-version.workspace = true
 [dependencies]
 litesvm = { workspace = true }
 solana-address.workspace = true
-solana-instruction.workspace = true
 solana-keypair.workspace = true
 solana-loader-v3-interface = { workspace = true, features = ["bincode"] }
 solana-signer.workspace = true
 solana-transaction.workspace = true
 
 [dev-dependencies]
 agave-feature-set.workspace = true
+solana-instruction.workspace = true
 solana-message.workspace = true
 test-log.workspace = true