Skip to main content

solana_bpf_loader_program/
lib.rs

1#![cfg(feature = "agave-unstable-api")]
2#![deny(clippy::arithmetic_side_effects)]
3#![deny(clippy::indexing_slicing)]
4
5#[cfg(feature = "svm-internal")]
6use qualifier_attr::qualifiers;
7use {
8    solana_bincode::limited_deserialize,
9    solana_instruction::{AccountMeta, error::InstructionError},
10    solana_loader_v3_interface::{
11        instruction::UpgradeableLoaderInstruction, state::UpgradeableLoaderState,
12    },
13    solana_program_runtime::{
14        deploy_program,
15        invoke_context::InvokeContext,
16        loaded_programs::{ProgramCacheEntry, ProgramCacheEntryOwner, ProgramCacheEntryType},
17        sysvar_cache::get_sysvar_with_account_check,
18        vm::execute,
19    },
20    solana_pubkey::Pubkey,
21    solana_sbpf::{declare_builtin_function, memory_region::MemoryMapping},
22    solana_sdk_ids::{
23        bpf_loader, bpf_loader_deprecated, bpf_loader_upgradeable, loader_v4, native_loader,
24    },
25    solana_svm_log_collector::{LogCollector, ic_logger_msg, ic_msg},
26    solana_svm_measure::measure::Measure,
27    solana_svm_type_overrides::sync::Arc,
28    solana_system_interface::{MAX_PERMITTED_DATA_LENGTH, instruction as system_instruction},
29    solana_transaction_context::{IndexOfAccount, instruction::InstructionContext},
30    std::{cell::RefCell, rc::Rc},
31};
32
33#[cfg_attr(feature = "svm-internal", qualifiers(pub))]
34const DEFAULT_LOADER_COMPUTE_UNITS: u64 = 570;
35#[cfg_attr(feature = "svm-internal", qualifiers(pub))]
36const DEPRECATED_LOADER_COMPUTE_UNITS: u64 = 1_140;
37#[cfg_attr(feature = "svm-internal", qualifiers(pub))]
38const UPGRADEABLE_LOADER_COMPUTE_UNITS: u64 = 2_370;
39
40fn write_program_data(
41    program_data_offset: usize,
42    bytes: &[u8],
43    invoke_context: &mut InvokeContext,
44) -> Result<(), InstructionError> {
45    let transaction_context = &invoke_context.transaction_context;
46    let instruction_context = transaction_context.get_current_instruction_context()?;
47    let mut program = instruction_context.try_borrow_instruction_account(0)?;
48    let data = program.get_data_mut()?;
49    let write_offset = program_data_offset.saturating_add(bytes.len());
50    if data.len() < write_offset {
51        ic_msg!(
52            invoke_context,
53            "Write overflow: {} < {}",
54            data.len(),
55            write_offset,
56        );
57        return Err(InstructionError::AccountDataTooSmall);
58    }
59    data.get_mut(program_data_offset..write_offset)
60        .ok_or(InstructionError::AccountDataTooSmall)?
61        .copy_from_slice(bytes);
62    Ok(())
63}
64
65declare_builtin_function!(
66    Entrypoint,
67    fn rust(
68        invoke_context: &mut InvokeContext<'static, 'static>,
69        _arg0: u64,
70        _arg1: u64,
71        _arg2: u64,
72        _arg3: u64,
73        _arg4: u64,
74        _memory_mapping: &mut MemoryMapping,
75    ) -> Result<u64, Box<dyn std::error::Error>> {
76        process_instruction_inner(invoke_context)
77    }
78);
79
80mod migration_authority {
81    solana_pubkey::declare_id!("3Scf35jMNk2xXBD6areNjgMtXgp5ZspDhms8vdcbzC42");
82}
83
84#[cfg_attr(feature = "svm-internal", qualifiers(pub))]
85pub(crate) fn process_instruction_inner<'a>(
86    invoke_context: &mut InvokeContext<'a, 'a>,
87) -> Result<u64, Box<dyn std::error::Error>> {
88    let log_collector = invoke_context.get_log_collector();
89    let transaction_context = &invoke_context.transaction_context;
90    let instruction_context = transaction_context.get_current_instruction_context()?;
91    let program_id = instruction_context.get_program_key()?;
92    let owner_id = instruction_context.get_program_owner()?;
93
94    // Program Management Instruction
95    if native_loader::check_id(&owner_id) {
96        let program_id = instruction_context.get_program_key()?;
97        return if bpf_loader_upgradeable::check_id(program_id) {
98            invoke_context.consume_checked(UPGRADEABLE_LOADER_COMPUTE_UNITS)?;
99            process_loader_upgradeable_instruction(invoke_context)
100        } else if bpf_loader::check_id(program_id) {
101            invoke_context.consume_checked(DEFAULT_LOADER_COMPUTE_UNITS)?;
102            ic_logger_msg!(
103                log_collector,
104                "BPF loader management instructions are no longer supported",
105            );
106            Err(InstructionError::UnsupportedProgramId)
107        } else if bpf_loader_deprecated::check_id(program_id) {
108            invoke_context.consume_checked(DEPRECATED_LOADER_COMPUTE_UNITS)?;
109            ic_logger_msg!(log_collector, "Deprecated loader is no longer supported");
110            Err(InstructionError::UnsupportedProgramId)
111        } else {
112            ic_logger_msg!(log_collector, "Invalid BPF loader id");
113            Err(InstructionError::UnsupportedProgramId)
114        }
115        .map(|_| 0)
116        .map_err(|error| Box::new(error) as Box<dyn std::error::Error>);
117    }
118
119    // Program Invocation
120    let mut get_or_create_executor_time = Measure::start("get_or_create_executor_time");
121    let executor = invoke_context
122        .program_cache_for_tx_batch
123        .find(program_id)
124        .ok_or_else(|| {
125            ic_logger_msg!(log_collector, "Program is not cached");
126            InstructionError::UnsupportedProgramId
127        })?;
128    get_or_create_executor_time.stop();
129    invoke_context.timings.get_or_create_executor_us += get_or_create_executor_time.as_us();
130
131    match &executor.program {
132        ProgramCacheEntryType::FailedVerification(_)
133        | ProgramCacheEntryType::Closed
134        | ProgramCacheEntryType::DelayVisibility => {
135            ic_logger_msg!(log_collector, "Program is not deployed");
136            Err(Box::new(InstructionError::UnsupportedProgramId) as Box<dyn std::error::Error>)
137        }
138        ProgramCacheEntryType::Loaded(executable) => execute(executable, invoke_context),
139        _ => Err(Box::new(InstructionError::UnsupportedProgramId) as Box<dyn std::error::Error>),
140    }
141    .map(|_| 0)
142}
143
144fn process_loader_upgradeable_instruction(
145    invoke_context: &mut InvokeContext,
146) -> Result<(), InstructionError> {
147    let log_collector = invoke_context.get_log_collector();
148    let transaction_context = &invoke_context.transaction_context;
149    let instruction_context = transaction_context.get_current_instruction_context()?;
150    let instruction_data = instruction_context.get_instruction_data();
151    let program_id = instruction_context.get_program_key()?;
152
153    match limited_deserialize(instruction_data, solana_packet::PACKET_DATA_SIZE as u64)? {
154        UpgradeableLoaderInstruction::InitializeBuffer => {
155            instruction_context.check_number_of_instruction_accounts(2)?;
156            let mut buffer = instruction_context.try_borrow_instruction_account(0)?;
157
158            if UpgradeableLoaderState::Uninitialized != buffer.get_state()? {
159                ic_logger_msg!(log_collector, "Buffer account already initialized");
160                return Err(InstructionError::AccountAlreadyInitialized);
161            }
162
163            let authority_key = Some(*instruction_context.get_key_of_instruction_account(1)?);
164
165            buffer.set_state(&UpgradeableLoaderState::Buffer {
166                authority_address: authority_key,
167            })?;
168        }
169        UpgradeableLoaderInstruction::Write { offset, bytes } => {
170            instruction_context.check_number_of_instruction_accounts(2)?;
171            let buffer = instruction_context.try_borrow_instruction_account(0)?;
172
173            if let UpgradeableLoaderState::Buffer { authority_address } = buffer.get_state()? {
174                if authority_address.is_none() {
175                    ic_logger_msg!(log_collector, "Buffer is immutable");
176                    return Err(InstructionError::Immutable); // TODO better error code
177                }
178                let authority_key = Some(*instruction_context.get_key_of_instruction_account(1)?);
179                if authority_address != authority_key {
180                    ic_logger_msg!(log_collector, "Incorrect buffer authority provided");
181                    return Err(InstructionError::IncorrectAuthority);
182                }
183                if !instruction_context.is_instruction_account_signer(1)? {
184                    ic_logger_msg!(log_collector, "Buffer authority did not sign");
185                    return Err(InstructionError::MissingRequiredSignature);
186                }
187            } else {
188                ic_logger_msg!(log_collector, "Invalid Buffer account");
189                return Err(InstructionError::InvalidAccountData);
190            }
191            drop(buffer);
192            write_program_data(
193                UpgradeableLoaderState::size_of_buffer_metadata().saturating_add(offset as usize),
194                &bytes,
195                invoke_context,
196            )?;
197        }
198        UpgradeableLoaderInstruction::DeployWithMaxDataLen { max_data_len } => {
199            instruction_context.check_number_of_instruction_accounts(4)?;
200            let payer_key = *instruction_context.get_key_of_instruction_account(0)?;
201            let programdata_key = *instruction_context.get_key_of_instruction_account(1)?;
202            let rent =
203                get_sysvar_with_account_check::rent(invoke_context, &instruction_context, 4)?;
204            let clock =
205                get_sysvar_with_account_check::clock(invoke_context, &instruction_context, 5)?;
206            instruction_context.check_number_of_instruction_accounts(8)?;
207            let authority_key = Some(*instruction_context.get_key_of_instruction_account(7)?);
208
209            // Verify Program account
210
211            let program = instruction_context.try_borrow_instruction_account(2)?;
212            if UpgradeableLoaderState::Uninitialized != program.get_state()? {
213                ic_logger_msg!(log_collector, "Program account already initialized");
214                return Err(InstructionError::AccountAlreadyInitialized);
215            }
216            if program.get_data().len() < UpgradeableLoaderState::size_of_program() {
217                ic_logger_msg!(log_collector, "Program account too small");
218                return Err(InstructionError::AccountDataTooSmall);
219            }
220            if program.get_lamports() < rent.minimum_balance(program.get_data().len()) {
221                ic_logger_msg!(log_collector, "Program account not rent-exempt");
222                return Err(InstructionError::ExecutableAccountNotRentExempt);
223            }
224            let new_program_id = *program.get_key();
225            drop(program);
226
227            // Verify Buffer account
228
229            let buffer = instruction_context.try_borrow_instruction_account(3)?;
230            if let UpgradeableLoaderState::Buffer { authority_address } = buffer.get_state()? {
231                if authority_address != authority_key {
232                    ic_logger_msg!(log_collector, "Buffer and upgrade authority don't match");
233                    return Err(InstructionError::IncorrectAuthority);
234                }
235                if !instruction_context.is_instruction_account_signer(7)? {
236                    ic_logger_msg!(log_collector, "Upgrade authority did not sign");
237                    return Err(InstructionError::MissingRequiredSignature);
238                }
239            } else {
240                ic_logger_msg!(log_collector, "Invalid Buffer account");
241                return Err(InstructionError::InvalidArgument);
242            }
243            let buffer_key = *buffer.get_key();
244            let buffer_data_offset = UpgradeableLoaderState::size_of_buffer_metadata();
245            let buffer_data_len = buffer.get_data().len().saturating_sub(buffer_data_offset);
246            let programdata_data_offset = UpgradeableLoaderState::size_of_programdata_metadata();
247            let programdata_len = UpgradeableLoaderState::size_of_programdata(max_data_len);
248            if buffer.get_data().len() < UpgradeableLoaderState::size_of_buffer_metadata()
249                || buffer_data_len == 0
250            {
251                ic_logger_msg!(log_collector, "Buffer account too small");
252                return Err(InstructionError::InvalidAccountData);
253            }
254            drop(buffer);
255            if max_data_len < buffer_data_len {
256                ic_logger_msg!(
257                    log_collector,
258                    "Max data length is too small to hold Buffer data"
259                );
260                return Err(InstructionError::AccountDataTooSmall);
261            }
262            if programdata_len > MAX_PERMITTED_DATA_LENGTH as usize {
263                ic_logger_msg!(log_collector, "Max data length is too large");
264                return Err(InstructionError::InvalidArgument);
265            }
266
267            // Create ProgramData account
268            let (derived_address, bump_seed) =
269                Pubkey::find_program_address(&[new_program_id.as_ref()], program_id);
270            if derived_address != programdata_key {
271                ic_logger_msg!(log_collector, "ProgramData address is not derived");
272                return Err(InstructionError::InvalidArgument);
273            }
274
275            // Drain the Buffer account to payer before paying for programdata account
276            {
277                let mut buffer = instruction_context.try_borrow_instruction_account(3)?;
278                let mut payer = instruction_context.try_borrow_instruction_account(0)?;
279                payer.checked_add_lamports(buffer.get_lamports())?;
280                buffer.set_lamports(0)?;
281            }
282
283            let owner_id = *program_id;
284            let mut instruction = system_instruction::create_account(
285                &payer_key,
286                &programdata_key,
287                1.max(rent.minimum_balance(programdata_len)),
288                programdata_len as u64,
289                program_id,
290            );
291
292            // pass an extra account to avoid the overly strict UnbalancedInstruction error
293            instruction
294                .accounts
295                .push(AccountMeta::new(buffer_key, false));
296
297            invoke_context
298                .native_invoke_signed(instruction, &[&[new_program_id.as_ref(), &[bump_seed]]])?;
299
300            // Load and verify the program bits
301            let transaction_context = &invoke_context.transaction_context;
302            let instruction_context = transaction_context.get_current_instruction_context()?;
303            let buffer = instruction_context.try_borrow_instruction_account(3)?;
304            deploy_program!(
305                invoke_context,
306                &new_program_id,
307                &owner_id,
308                UpgradeableLoaderState::size_of_program().saturating_add(programdata_len),
309                buffer
310                    .get_data()
311                    .get(buffer_data_offset..)
312                    .ok_or(InstructionError::AccountDataTooSmall)?,
313                clock.slot,
314            );
315            drop(buffer);
316
317            let transaction_context = &invoke_context.transaction_context;
318            let instruction_context = transaction_context.get_current_instruction_context()?;
319
320            // Update the ProgramData account and record the program bits
321            {
322                let mut programdata = instruction_context.try_borrow_instruction_account(1)?;
323                programdata.set_state(&UpgradeableLoaderState::ProgramData {
324                    slot: clock.slot,
325                    upgrade_authority_address: authority_key,
326                })?;
327                let dst_slice = programdata
328                    .get_data_mut()?
329                    .get_mut(
330                        programdata_data_offset
331                            ..programdata_data_offset.saturating_add(buffer_data_len),
332                    )
333                    .ok_or(InstructionError::AccountDataTooSmall)?;
334                let mut buffer = instruction_context.try_borrow_instruction_account(3)?;
335                let src_slice = buffer
336                    .get_data()
337                    .get(buffer_data_offset..)
338                    .ok_or(InstructionError::AccountDataTooSmall)?;
339                dst_slice.copy_from_slice(src_slice);
340                buffer.set_data_length(UpgradeableLoaderState::size_of_buffer(0))?;
341            }
342
343            // Update the Program account
344            let mut program = instruction_context.try_borrow_instruction_account(2)?;
345            program.set_state(&UpgradeableLoaderState::Program {
346                programdata_address: programdata_key,
347            })?;
348            program.set_executable(true)?;
349            drop(program);
350
351            ic_logger_msg!(log_collector, "Deployed program {:?}", new_program_id);
352        }
353        UpgradeableLoaderInstruction::Upgrade => {
354            instruction_context.check_number_of_instruction_accounts(3)?;
355            let programdata_key = *instruction_context.get_key_of_instruction_account(0)?;
356            let rent =
357                get_sysvar_with_account_check::rent(invoke_context, &instruction_context, 4)?;
358            let clock =
359                get_sysvar_with_account_check::clock(invoke_context, &instruction_context, 5)?;
360            instruction_context.check_number_of_instruction_accounts(7)?;
361            let authority_key = Some(*instruction_context.get_key_of_instruction_account(6)?);
362
363            // Verify Program account
364
365            let program = instruction_context.try_borrow_instruction_account(1)?;
366            if !program.is_writable() {
367                ic_logger_msg!(log_collector, "Program account not writeable");
368                return Err(InstructionError::InvalidArgument);
369            }
370            if program.get_owner() != program_id {
371                ic_logger_msg!(log_collector, "Program account not owned by loader");
372                return Err(InstructionError::IncorrectProgramId);
373            }
374            if let UpgradeableLoaderState::Program {
375                programdata_address,
376            } = program.get_state()?
377            {
378                if programdata_address != programdata_key {
379                    ic_logger_msg!(log_collector, "Program and ProgramData account mismatch");
380                    return Err(InstructionError::InvalidArgument);
381                }
382            } else {
383                ic_logger_msg!(log_collector, "Invalid Program account");
384                return Err(InstructionError::InvalidAccountData);
385            }
386            let new_program_id = *program.get_key();
387            drop(program);
388
389            // Verify Buffer account
390
391            let buffer = instruction_context.try_borrow_instruction_account(2)?;
392            if let UpgradeableLoaderState::Buffer { authority_address } = buffer.get_state()? {
393                if authority_address != authority_key {
394                    ic_logger_msg!(log_collector, "Buffer and upgrade authority don't match");
395                    return Err(InstructionError::IncorrectAuthority);
396                }
397                if !instruction_context.is_instruction_account_signer(6)? {
398                    ic_logger_msg!(log_collector, "Upgrade authority did not sign");
399                    return Err(InstructionError::MissingRequiredSignature);
400                }
401            } else {
402                ic_logger_msg!(log_collector, "Invalid Buffer account");
403                return Err(InstructionError::InvalidArgument);
404            }
405            let buffer_lamports = buffer.get_lamports();
406            let buffer_data_offset = UpgradeableLoaderState::size_of_buffer_metadata();
407            let buffer_data_len = buffer.get_data().len().saturating_sub(buffer_data_offset);
408            if buffer.get_data().len() < UpgradeableLoaderState::size_of_buffer_metadata()
409                || buffer_data_len == 0
410            {
411                ic_logger_msg!(log_collector, "Buffer account too small");
412                return Err(InstructionError::InvalidAccountData);
413            }
414            drop(buffer);
415
416            // Verify ProgramData account
417
418            let programdata = instruction_context.try_borrow_instruction_account(0)?;
419            let programdata_data_offset = UpgradeableLoaderState::size_of_programdata_metadata();
420            let programdata_balance_required =
421                1.max(rent.minimum_balance(programdata.get_data().len()));
422            if programdata.get_data().len()
423                < UpgradeableLoaderState::size_of_programdata(buffer_data_len)
424            {
425                ic_logger_msg!(log_collector, "ProgramData account not large enough");
426                return Err(InstructionError::AccountDataTooSmall);
427            }
428            if programdata.get_lamports().saturating_add(buffer_lamports)
429                < programdata_balance_required
430            {
431                ic_logger_msg!(
432                    log_collector,
433                    "Buffer account balance too low to fund upgrade"
434                );
435                return Err(InstructionError::InsufficientFunds);
436            }
437            if let UpgradeableLoaderState::ProgramData {
438                slot,
439                upgrade_authority_address,
440            } = programdata.get_state()?
441            {
442                if clock.slot == slot {
443                    ic_logger_msg!(log_collector, "Program was deployed in this block already");
444                    return Err(InstructionError::InvalidArgument);
445                }
446                if upgrade_authority_address.is_none() {
447                    ic_logger_msg!(log_collector, "Program not upgradeable");
448                    return Err(InstructionError::Immutable);
449                }
450                if upgrade_authority_address != authority_key {
451                    ic_logger_msg!(log_collector, "Incorrect upgrade authority provided");
452                    return Err(InstructionError::IncorrectAuthority);
453                }
454                if !instruction_context.is_instruction_account_signer(6)? {
455                    ic_logger_msg!(log_collector, "Upgrade authority did not sign");
456                    return Err(InstructionError::MissingRequiredSignature);
457                }
458            } else {
459                ic_logger_msg!(log_collector, "Invalid ProgramData account");
460                return Err(InstructionError::InvalidAccountData);
461            };
462            let programdata_len = programdata.get_data().len();
463            drop(programdata);
464
465            // Load and verify the program bits
466            let buffer = instruction_context.try_borrow_instruction_account(2)?;
467            deploy_program!(
468                invoke_context,
469                &new_program_id,
470                program_id,
471                UpgradeableLoaderState::size_of_program().saturating_add(programdata_len),
472                buffer
473                    .get_data()
474                    .get(buffer_data_offset..)
475                    .ok_or(InstructionError::AccountDataTooSmall)?,
476                clock.slot,
477            );
478            drop(buffer);
479
480            let transaction_context = &invoke_context.transaction_context;
481            let instruction_context = transaction_context.get_current_instruction_context()?;
482
483            // Update the ProgramData account, record the upgraded data, and zero
484            // the rest
485            let mut programdata = instruction_context.try_borrow_instruction_account(0)?;
486            {
487                programdata.set_state(&UpgradeableLoaderState::ProgramData {
488                    slot: clock.slot,
489                    upgrade_authority_address: authority_key,
490                })?;
491                let dst_slice = programdata
492                    .get_data_mut()?
493                    .get_mut(
494                        programdata_data_offset
495                            ..programdata_data_offset.saturating_add(buffer_data_len),
496                    )
497                    .ok_or(InstructionError::AccountDataTooSmall)?;
498                let buffer = instruction_context.try_borrow_instruction_account(2)?;
499                let src_slice = buffer
500                    .get_data()
501                    .get(buffer_data_offset..)
502                    .ok_or(InstructionError::AccountDataTooSmall)?;
503                dst_slice.copy_from_slice(src_slice);
504            }
505            programdata
506                .get_data_mut()?
507                .get_mut(programdata_data_offset.saturating_add(buffer_data_len)..)
508                .ok_or(InstructionError::AccountDataTooSmall)?
509                .fill(0);
510
511            // Fund ProgramData to rent-exemption, spill the rest
512            let mut buffer = instruction_context.try_borrow_instruction_account(2)?;
513            let mut spill = instruction_context.try_borrow_instruction_account(3)?;
514            spill.checked_add_lamports(
515                programdata
516                    .get_lamports()
517                    .saturating_add(buffer_lamports)
518                    .saturating_sub(programdata_balance_required),
519            )?;
520            buffer.set_lamports(0)?;
521            programdata.set_lamports(programdata_balance_required)?;
522            buffer.set_data_length(UpgradeableLoaderState::size_of_buffer(0))?;
523
524            ic_logger_msg!(log_collector, "Upgraded program {:?}", new_program_id);
525        }
526        UpgradeableLoaderInstruction::SetAuthority => {
527            instruction_context.check_number_of_instruction_accounts(2)?;
528            let mut account = instruction_context.try_borrow_instruction_account(0)?;
529            let present_authority_key = instruction_context.get_key_of_instruction_account(1)?;
530            let new_authority = instruction_context.get_key_of_instruction_account(2).ok();
531
532            match account.get_state()? {
533                UpgradeableLoaderState::Buffer { authority_address } => {
534                    if new_authority.is_none() {
535                        ic_logger_msg!(log_collector, "Buffer authority is not optional");
536                        return Err(InstructionError::IncorrectAuthority);
537                    }
538                    if authority_address.is_none() {
539                        ic_logger_msg!(log_collector, "Buffer is immutable");
540                        return Err(InstructionError::Immutable);
541                    }
542                    if authority_address != Some(*present_authority_key) {
543                        ic_logger_msg!(log_collector, "Incorrect buffer authority provided");
544                        return Err(InstructionError::IncorrectAuthority);
545                    }
546                    if !instruction_context.is_instruction_account_signer(1)? {
547                        ic_logger_msg!(log_collector, "Buffer authority did not sign");
548                        return Err(InstructionError::MissingRequiredSignature);
549                    }
550                    account.set_state(&UpgradeableLoaderState::Buffer {
551                        authority_address: new_authority.cloned(),
552                    })?;
553                }
554                UpgradeableLoaderState::ProgramData {
555                    slot,
556                    upgrade_authority_address,
557                } => {
558                    if upgrade_authority_address.is_none() {
559                        ic_logger_msg!(log_collector, "Program not upgradeable");
560                        return Err(InstructionError::Immutable);
561                    }
562                    if upgrade_authority_address != Some(*present_authority_key) {
563                        ic_logger_msg!(log_collector, "Incorrect upgrade authority provided");
564                        return Err(InstructionError::IncorrectAuthority);
565                    }
566                    if !instruction_context.is_instruction_account_signer(1)? {
567                        ic_logger_msg!(log_collector, "Upgrade authority did not sign");
568                        return Err(InstructionError::MissingRequiredSignature);
569                    }
570                    account.set_state(&UpgradeableLoaderState::ProgramData {
571                        slot,
572                        upgrade_authority_address: new_authority.cloned(),
573                    })?;
574                }
575                _ => {
576                    ic_logger_msg!(log_collector, "Account does not support authorities");
577                    return Err(InstructionError::InvalidArgument);
578                }
579            }
580
581            ic_logger_msg!(log_collector, "New authority {:?}", new_authority);
582        }
583        UpgradeableLoaderInstruction::SetAuthorityChecked => {
584            if !invoke_context
585                .get_feature_set()
586                .enable_bpf_loader_set_authority_checked_ix
587            {
588                return Err(InstructionError::InvalidInstructionData);
589            }
590
591            instruction_context.check_number_of_instruction_accounts(3)?;
592            let mut account = instruction_context.try_borrow_instruction_account(0)?;
593            let present_authority_key = instruction_context.get_key_of_instruction_account(1)?;
594            let new_authority_key = instruction_context.get_key_of_instruction_account(2)?;
595
596            match account.get_state()? {
597                UpgradeableLoaderState::Buffer { authority_address } => {
598                    if authority_address.is_none() {
599                        ic_logger_msg!(log_collector, "Buffer is immutable");
600                        return Err(InstructionError::Immutable);
601                    }
602                    if authority_address != Some(*present_authority_key) {
603                        ic_logger_msg!(log_collector, "Incorrect buffer authority provided");
604                        return Err(InstructionError::IncorrectAuthority);
605                    }
606                    if !instruction_context.is_instruction_account_signer(1)? {
607                        ic_logger_msg!(log_collector, "Buffer authority did not sign");
608                        return Err(InstructionError::MissingRequiredSignature);
609                    }
610                    if !instruction_context.is_instruction_account_signer(2)? {
611                        ic_logger_msg!(log_collector, "New authority did not sign");
612                        return Err(InstructionError::MissingRequiredSignature);
613                    }
614                    account.set_state(&UpgradeableLoaderState::Buffer {
615                        authority_address: Some(*new_authority_key),
616                    })?;
617                }
618                UpgradeableLoaderState::ProgramData {
619                    slot,
620                    upgrade_authority_address,
621                } => {
622                    if upgrade_authority_address.is_none() {
623                        ic_logger_msg!(log_collector, "Program not upgradeable");
624                        return Err(InstructionError::Immutable);
625                    }
626                    if upgrade_authority_address != Some(*present_authority_key) {
627                        ic_logger_msg!(log_collector, "Incorrect upgrade authority provided");
628                        return Err(InstructionError::IncorrectAuthority);
629                    }
630                    if !instruction_context.is_instruction_account_signer(1)? {
631                        ic_logger_msg!(log_collector, "Upgrade authority did not sign");
632                        return Err(InstructionError::MissingRequiredSignature);
633                    }
634                    if !instruction_context.is_instruction_account_signer(2)? {
635                        ic_logger_msg!(log_collector, "New authority did not sign");
636                        return Err(InstructionError::MissingRequiredSignature);
637                    }
638                    account.set_state(&UpgradeableLoaderState::ProgramData {
639                        slot,
640                        upgrade_authority_address: Some(*new_authority_key),
641                    })?;
642                }
643                _ => {
644                    ic_logger_msg!(log_collector, "Account does not support authorities");
645                    return Err(InstructionError::InvalidArgument);
646                }
647            }
648
649            ic_logger_msg!(log_collector, "New authority {:?}", new_authority_key);
650        }
651        UpgradeableLoaderInstruction::Close => {
652            instruction_context.check_number_of_instruction_accounts(2)?;
653            if instruction_context.get_index_of_instruction_account_in_transaction(0)?
654                == instruction_context.get_index_of_instruction_account_in_transaction(1)?
655            {
656                ic_logger_msg!(
657                    log_collector,
658                    "Recipient is the same as the account being closed"
659                );
660                return Err(InstructionError::InvalidArgument);
661            }
662            let mut close_account = instruction_context.try_borrow_instruction_account(0)?;
663            let close_key = *close_account.get_key();
664            let close_account_state = close_account.get_state()?;
665            close_account.set_data_length(UpgradeableLoaderState::size_of_uninitialized())?;
666            match close_account_state {
667                UpgradeableLoaderState::Uninitialized => {
668                    let mut recipient_account =
669                        instruction_context.try_borrow_instruction_account(1)?;
670                    recipient_account.checked_add_lamports(close_account.get_lamports())?;
671                    close_account.set_lamports(0)?;
672
673                    ic_logger_msg!(log_collector, "Closed Uninitialized {}", close_key);
674                }
675                UpgradeableLoaderState::Buffer { authority_address } => {
676                    instruction_context.check_number_of_instruction_accounts(3)?;
677                    drop(close_account);
678                    common_close_account(&authority_address, &instruction_context, &log_collector)?;
679
680                    ic_logger_msg!(log_collector, "Closed Buffer {}", close_key);
681                }
682                UpgradeableLoaderState::ProgramData {
683                    slot,
684                    upgrade_authority_address: authority_address,
685                } => {
686                    instruction_context.check_number_of_instruction_accounts(4)?;
687                    drop(close_account);
688                    let program_account = instruction_context.try_borrow_instruction_account(3)?;
689                    let program_key = *program_account.get_key();
690
691                    if !program_account.is_writable() {
692                        ic_logger_msg!(log_collector, "Program account is not writable");
693                        return Err(InstructionError::InvalidArgument);
694                    }
695                    if program_account.get_owner() != program_id {
696                        ic_logger_msg!(log_collector, "Program account not owned by loader");
697                        return Err(InstructionError::IncorrectProgramId);
698                    }
699                    let clock = invoke_context.get_sysvar_cache().get_clock()?;
700                    if clock.slot == slot {
701                        ic_logger_msg!(log_collector, "Program was deployed in this block already");
702                        return Err(InstructionError::InvalidArgument);
703                    }
704
705                    match program_account.get_state()? {
706                        UpgradeableLoaderState::Program {
707                            programdata_address,
708                        } => {
709                            if programdata_address != close_key {
710                                ic_logger_msg!(
711                                    log_collector,
712                                    "ProgramData account does not match ProgramData account"
713                                );
714                                return Err(InstructionError::InvalidArgument);
715                            }
716
717                            drop(program_account);
718                            common_close_account(
719                                &authority_address,
720                                &instruction_context,
721                                &log_collector,
722                            )?;
723                            let clock = invoke_context.get_sysvar_cache().get_clock()?;
724                            invoke_context
725                                .program_cache_for_tx_batch
726                                .store_modified_entry(
727                                    program_key,
728                                    Arc::new(ProgramCacheEntry::new_tombstone(
729                                        clock.slot,
730                                        ProgramCacheEntryOwner::LoaderV3,
731                                        ProgramCacheEntryType::Closed,
732                                    )),
733                                );
734                        }
735                        _ => {
736                            ic_logger_msg!(log_collector, "Invalid Program account");
737                            return Err(InstructionError::InvalidArgument);
738                        }
739                    }
740
741                    ic_logger_msg!(log_collector, "Closed Program {}", program_key);
742                }
743                _ => {
744                    ic_logger_msg!(log_collector, "Account does not support closing");
745                    return Err(InstructionError::InvalidArgument);
746                }
747            }
748        }
749        UpgradeableLoaderInstruction::ExtendProgram { additional_bytes } => {
750            if invoke_context
751                .get_feature_set()
752                .enable_extend_program_checked
753            {
754                ic_logger_msg!(
755                    log_collector,
756                    "ExtendProgram was superseded by ExtendProgramChecked"
757                );
758                return Err(InstructionError::InvalidInstructionData);
759            }
760            common_extend_program(invoke_context, additional_bytes, false)?;
761        }
762        UpgradeableLoaderInstruction::ExtendProgramChecked { additional_bytes } => {
763            if !invoke_context
764                .get_feature_set()
765                .enable_extend_program_checked
766            {
767                return Err(InstructionError::InvalidInstructionData);
768            }
769            common_extend_program(invoke_context, additional_bytes, true)?;
770        }
771        UpgradeableLoaderInstruction::Migrate => {
772            if !invoke_context.get_feature_set().enable_loader_v4 {
773                return Err(InstructionError::InvalidInstructionData);
774            }
775
776            instruction_context.check_number_of_instruction_accounts(3)?;
777            let programdata_address = *instruction_context.get_key_of_instruction_account(0)?;
778            let program_address = *instruction_context.get_key_of_instruction_account(1)?;
779            let provided_authority_address =
780                *instruction_context.get_key_of_instruction_account(2)?;
781            let clock_slot = invoke_context
782                .get_sysvar_cache()
783                .get_clock()
784                .map(|clock| clock.slot)?;
785
786            // Verify ProgramData account
787            let programdata = instruction_context.try_borrow_instruction_account(0)?;
788            if !programdata.is_writable() {
789                ic_logger_msg!(log_collector, "ProgramData account not writeable");
790                return Err(InstructionError::InvalidArgument);
791            }
792            let (program_len, upgrade_authority_address) =
793                if let Ok(UpgradeableLoaderState::ProgramData {
794                    slot,
795                    upgrade_authority_address,
796                }) = programdata.get_state()
797                {
798                    if clock_slot == slot {
799                        ic_logger_msg!(log_collector, "Program was deployed in this block already");
800                        return Err(InstructionError::InvalidArgument);
801                    }
802                    (
803                        programdata
804                            .get_data()
805                            .len()
806                            .saturating_sub(UpgradeableLoaderState::size_of_programdata_metadata()),
807                        upgrade_authority_address,
808                    )
809                } else {
810                    (0, None)
811                };
812            let programdata_funds = programdata.get_lamports();
813            drop(programdata);
814
815            // Verify authority signature
816            if !migration_authority::check_id(&provided_authority_address)
817                && provided_authority_address
818                    != upgrade_authority_address.unwrap_or(program_address)
819            {
820                ic_logger_msg!(log_collector, "Incorrect migration authority provided");
821                return Err(InstructionError::IncorrectAuthority);
822            }
823            if !instruction_context.is_instruction_account_signer(2)? {
824                ic_logger_msg!(log_collector, "Migration authority did not sign");
825                return Err(InstructionError::MissingRequiredSignature);
826            }
827
828            // Verify Program account
829            let mut program = instruction_context.try_borrow_instruction_account(1)?;
830            if !program.is_writable() {
831                ic_logger_msg!(log_collector, "Program account not writeable");
832                return Err(InstructionError::InvalidArgument);
833            }
834            if program.get_owner() != program_id {
835                ic_logger_msg!(log_collector, "Program account not owned by loader");
836                return Err(InstructionError::IncorrectProgramId);
837            }
838            if let UpgradeableLoaderState::Program {
839                programdata_address: stored_programdata_address,
840            } = program.get_state()?
841            {
842                if programdata_address != stored_programdata_address {
843                    ic_logger_msg!(log_collector, "Program and ProgramData account mismatch");
844                    return Err(InstructionError::InvalidArgument);
845                }
846            } else {
847                ic_logger_msg!(log_collector, "Invalid Program account");
848                return Err(InstructionError::InvalidAccountData);
849            }
850            program.set_data_from_slice(&[])?;
851            program.checked_add_lamports(programdata_funds)?;
852            program.set_owner(&loader_v4::id().to_bytes())?;
853            drop(program);
854
855            let mut programdata = instruction_context.try_borrow_instruction_account(0)?;
856            programdata.set_lamports(0)?;
857            drop(programdata);
858
859            if program_len == 0 {
860                invoke_context
861                    .program_cache_for_tx_batch
862                    .store_modified_entry(
863                        program_address,
864                        Arc::new(ProgramCacheEntry::new_tombstone(
865                            clock_slot,
866                            ProgramCacheEntryOwner::LoaderV4,
867                            ProgramCacheEntryType::Closed,
868                        )),
869                    );
870            } else {
871                invoke_context.native_invoke_signed(
872                    solana_loader_v4_interface::instruction::set_program_length(
873                        &program_address,
874                        &provided_authority_address,
875                        program_len as u32,
876                        &program_address,
877                    ),
878                    &[],
879                )?;
880
881                invoke_context.native_invoke_signed(
882                    solana_loader_v4_interface::instruction::copy(
883                        &program_address,
884                        &provided_authority_address,
885                        &programdata_address,
886                        0,
887                        0,
888                        program_len as u32,
889                    ),
890                    &[],
891                )?;
892
893                invoke_context.native_invoke_signed(
894                    solana_loader_v4_interface::instruction::deploy(
895                        &program_address,
896                        &provided_authority_address,
897                    ),
898                    &[],
899                )?;
900
901                if let Some(upgrade_authority_address) = upgrade_authority_address {
902                    if migration_authority::check_id(&provided_authority_address) {
903                        invoke_context.native_invoke_signed(
904                            solana_loader_v4_interface::instruction::transfer_authority(
905                                &program_address,
906                                &provided_authority_address,
907                                &upgrade_authority_address,
908                            ),
909                            &[],
910                        )?;
911                    }
912                } else {
913                    invoke_context.native_invoke_signed(
914                        solana_loader_v4_interface::instruction::finalize(
915                            &program_address,
916                            &provided_authority_address,
917                            &program_address,
918                        ),
919                        &[],
920                    )?;
921                }
922            }
923
924            let transaction_context = &invoke_context.transaction_context;
925            let instruction_context = transaction_context.get_current_instruction_context()?;
926            let mut programdata = instruction_context.try_borrow_instruction_account(0)?;
927            programdata.set_data_from_slice(&[])?;
928            drop(programdata);
929
930            ic_logger_msg!(log_collector, "Migrated program {:?}", &program_address);
931        }
932    }
933
934    Ok(())
935}
936
937fn common_extend_program(
938    invoke_context: &mut InvokeContext,
939    additional_bytes: u32,
940    check_authority: bool,
941) -> Result<(), InstructionError> {
942    let log_collector = invoke_context.get_log_collector();
943    let transaction_context = &invoke_context.transaction_context;
944    let instruction_context = transaction_context.get_current_instruction_context()?;
945    let program_id = instruction_context.get_program_key()?;
946
947    const PROGRAM_DATA_ACCOUNT_INDEX: IndexOfAccount = 0;
948    const PROGRAM_ACCOUNT_INDEX: IndexOfAccount = 1;
949    const AUTHORITY_ACCOUNT_INDEX: IndexOfAccount = 2;
950    // The unused `system_program_account_index` is 3 if `check_authority` and 2 otherwise.
951    let optional_payer_account_index = if check_authority { 4 } else { 3 };
952
953    if additional_bytes == 0 {
954        ic_logger_msg!(log_collector, "Additional bytes must be greater than 0");
955        return Err(InstructionError::InvalidInstructionData);
956    }
957
958    let programdata_account =
959        instruction_context.try_borrow_instruction_account(PROGRAM_DATA_ACCOUNT_INDEX)?;
960    let programdata_key = *programdata_account.get_key();
961
962    if program_id != programdata_account.get_owner() {
963        ic_logger_msg!(log_collector, "ProgramData owner is invalid");
964        return Err(InstructionError::InvalidAccountOwner);
965    }
966    if !programdata_account.is_writable() {
967        ic_logger_msg!(log_collector, "ProgramData is not writable");
968        return Err(InstructionError::InvalidArgument);
969    }
970
971    let program_account =
972        instruction_context.try_borrow_instruction_account(PROGRAM_ACCOUNT_INDEX)?;
973    if !program_account.is_writable() {
974        ic_logger_msg!(log_collector, "Program account is not writable");
975        return Err(InstructionError::InvalidArgument);
976    }
977    if program_account.get_owner() != program_id {
978        ic_logger_msg!(log_collector, "Program account not owned by loader");
979        return Err(InstructionError::InvalidAccountOwner);
980    }
981    let program_key = *program_account.get_key();
982    match program_account.get_state()? {
983        UpgradeableLoaderState::Program {
984            programdata_address,
985        } => {
986            if programdata_address != programdata_key {
987                ic_logger_msg!(
988                    log_collector,
989                    "Program account does not match ProgramData account"
990                );
991                return Err(InstructionError::InvalidArgument);
992            }
993        }
994        _ => {
995            ic_logger_msg!(log_collector, "Invalid Program account");
996            return Err(InstructionError::InvalidAccountData);
997        }
998    }
999    drop(program_account);
1000
1001    let old_len = programdata_account.get_data().len();
1002    let new_len = old_len.saturating_add(additional_bytes as usize);
1003    if new_len > MAX_PERMITTED_DATA_LENGTH as usize {
1004        ic_logger_msg!(
1005            log_collector,
1006            "Extended ProgramData length of {} bytes exceeds max account data length of {} bytes",
1007            new_len,
1008            MAX_PERMITTED_DATA_LENGTH
1009        );
1010        return Err(InstructionError::InvalidRealloc);
1011    }
1012
1013    let clock_slot = invoke_context
1014        .get_sysvar_cache()
1015        .get_clock()
1016        .map(|clock| clock.slot)?;
1017
1018    let upgrade_authority_address = if let UpgradeableLoaderState::ProgramData {
1019        slot,
1020        upgrade_authority_address,
1021    } = programdata_account.get_state()?
1022    {
1023        if clock_slot == slot {
1024            ic_logger_msg!(log_collector, "Program was extended in this block already");
1025            return Err(InstructionError::InvalidArgument);
1026        }
1027
1028        if upgrade_authority_address.is_none() {
1029            ic_logger_msg!(
1030                log_collector,
1031                "Cannot extend ProgramData accounts that are not upgradeable"
1032            );
1033            return Err(InstructionError::Immutable);
1034        }
1035
1036        if check_authority {
1037            let authority_key =
1038                Some(*instruction_context.get_key_of_instruction_account(AUTHORITY_ACCOUNT_INDEX)?);
1039            if upgrade_authority_address != authority_key {
1040                ic_logger_msg!(log_collector, "Incorrect upgrade authority provided");
1041                return Err(InstructionError::IncorrectAuthority);
1042            }
1043            if !instruction_context.is_instruction_account_signer(AUTHORITY_ACCOUNT_INDEX)? {
1044                ic_logger_msg!(log_collector, "Upgrade authority did not sign");
1045                return Err(InstructionError::MissingRequiredSignature);
1046            }
1047        }
1048
1049        upgrade_authority_address
1050    } else {
1051        ic_logger_msg!(log_collector, "ProgramData state is invalid");
1052        return Err(InstructionError::InvalidAccountData);
1053    };
1054
1055    let required_payment = {
1056        let balance = programdata_account.get_lamports();
1057        let rent = invoke_context.get_sysvar_cache().get_rent()?;
1058        let min_balance = rent.minimum_balance(new_len).max(1);
1059        min_balance.saturating_sub(balance)
1060    };
1061
1062    // Borrowed accounts need to be dropped before native_invoke_signed
1063    drop(programdata_account);
1064
1065    // Dereference the program ID to prevent overlapping mutable/immutable borrow of invoke context
1066    let program_id = *program_id;
1067    if required_payment > 0 {
1068        let payer_key =
1069            *instruction_context.get_key_of_instruction_account(optional_payer_account_index)?;
1070
1071        invoke_context.native_invoke_signed(
1072            system_instruction::transfer(&payer_key, &programdata_key, required_payment),
1073            &[],
1074        )?;
1075    }
1076
1077    let transaction_context = &invoke_context.transaction_context;
1078    let instruction_context = transaction_context.get_current_instruction_context()?;
1079    let mut programdata_account =
1080        instruction_context.try_borrow_instruction_account(PROGRAM_DATA_ACCOUNT_INDEX)?;
1081    programdata_account.set_data_length(new_len)?;
1082
1083    let programdata_data_offset = UpgradeableLoaderState::size_of_programdata_metadata();
1084
1085    deploy_program!(
1086        invoke_context,
1087        &program_key,
1088        &program_id,
1089        UpgradeableLoaderState::size_of_program().saturating_add(new_len),
1090        programdata_account
1091            .get_data()
1092            .get(programdata_data_offset..)
1093            .ok_or(InstructionError::AccountDataTooSmall)?,
1094        clock_slot,
1095    );
1096    drop(programdata_account);
1097
1098    let mut programdata_account =
1099        instruction_context.try_borrow_instruction_account(PROGRAM_DATA_ACCOUNT_INDEX)?;
1100    programdata_account.set_state(&UpgradeableLoaderState::ProgramData {
1101        slot: clock_slot,
1102        upgrade_authority_address,
1103    })?;
1104
1105    ic_logger_msg!(
1106        log_collector,
1107        "Extended ProgramData account by {} bytes",
1108        additional_bytes
1109    );
1110
1111    Ok(())
1112}
1113
1114fn common_close_account(
1115    authority_address: &Option<Pubkey>,
1116    instruction_context: &InstructionContext,
1117    log_collector: &Option<Rc<RefCell<LogCollector>>>,
1118) -> Result<(), InstructionError> {
1119    if authority_address.is_none() {
1120        ic_logger_msg!(log_collector, "Account is immutable");
1121        return Err(InstructionError::Immutable);
1122    }
1123    if *authority_address != Some(*instruction_context.get_key_of_instruction_account(2)?) {
1124        ic_logger_msg!(log_collector, "Incorrect authority provided");
1125        return Err(InstructionError::IncorrectAuthority);
1126    }
1127    if !instruction_context.is_instruction_account_signer(2)? {
1128        ic_logger_msg!(log_collector, "Authority did not sign");
1129        return Err(InstructionError::MissingRequiredSignature);
1130    }
1131
1132    let mut close_account = instruction_context.try_borrow_instruction_account(0)?;
1133    let mut recipient_account = instruction_context.try_borrow_instruction_account(1)?;
1134
1135    recipient_account.checked_add_lamports(close_account.get_lamports())?;
1136    close_account.set_lamports(0)?;
1137    close_account.set_state(&UpgradeableLoaderState::Uninitialized)?;
1138    Ok(())
1139}
1140
1141#[cfg_attr(feature = "svm-internal", qualifiers(pub))]
1142mod test_utils {
1143    #[cfg(all(feature = "svm-internal", feature = "metrics"))]
1144    use solana_program_runtime::loaded_programs::LoadProgramMetrics;
1145    #[cfg(feature = "svm-internal")]
1146    use {
1147        super::*, agave_syscalls::create_program_runtime_environment_v1,
1148        solana_account::ReadableAccount, solana_loader_v4_interface::state::LoaderV4State,
1149        solana_program_runtime::loaded_programs::DELAY_VISIBILITY_SLOT_OFFSET,
1150        solana_sdk_ids::loader_v4,
1151    };
1152
1153    #[cfg(feature = "svm-internal")]
1154    fn check_loader_id(id: &Pubkey) -> bool {
1155        bpf_loader::check_id(id)
1156            || bpf_loader_deprecated::check_id(id)
1157            || bpf_loader_upgradeable::check_id(id)
1158            || loader_v4::check_id(id)
1159    }
1160
1161    #[cfg(feature = "svm-internal")]
1162    #[cfg_attr(feature = "svm-internal", qualifiers(pub))]
1163    fn load_all_invoked_programs(invoke_context: &mut InvokeContext) {
1164        let program_runtime_environment = create_program_runtime_environment_v1(
1165            invoke_context.get_feature_set(),
1166            invoke_context.get_compute_budget(),
1167            false, /* deployment */
1168            false, /* debugging_features */
1169        );
1170        let program_runtime_environment = Arc::new(program_runtime_environment.unwrap());
1171        let num_accounts = invoke_context.transaction_context.get_number_of_accounts();
1172        for index in 0..num_accounts {
1173            let account = invoke_context
1174                .transaction_context
1175                .accounts()
1176                .try_borrow(index)
1177                .expect("Failed to get the account");
1178
1179            let owner = account.owner();
1180            if check_loader_id(owner) {
1181                let programdata_data_offset = if loader_v4::check_id(owner) {
1182                    LoaderV4State::program_data_offset()
1183                } else {
1184                    0
1185                };
1186                let pubkey = invoke_context
1187                    .transaction_context
1188                    .get_key_of_account_at_index(index)
1189                    .expect("Failed to get account key");
1190
1191                let programdata = account
1192                    .data()
1193                    .get(programdata_data_offset.min(account.data().len())..)
1194                    .unwrap();
1195                let program_runtime_environment = program_runtime_environment.clone();
1196                let effective_slot = DELAY_VISIBILITY_SLOT_OFFSET;
1197                let loaded_program = ProgramCacheEntry::new(
1198                    owner,
1199                    program_runtime_environment,
1200                    0,
1201                    effective_slot,
1202                    programdata,
1203                    account.data().len(),
1204                    #[cfg(feature = "metrics")]
1205                    &mut LoadProgramMetrics::default(),
1206                )
1207                .map_err(|_| InstructionError::InvalidAccountData);
1208                if let Ok(loaded_program) = loaded_program {
1209                    invoke_context
1210                        .program_cache_for_tx_batch
1211                        .store_modified_entry(*pubkey, Arc::new(loaded_program));
1212                }
1213            }
1214        }
1215    }
1216}
1217
1218#[cfg(test)]
1219mod tests {
1220    use {
1221        super::*,
1222        assert_matches::assert_matches,
1223        rand::Rng,
1224        solana_account::{
1225            AccountSharedData, ReadableAccount, WritableAccount,
1226            create_account_shared_data_for_test as create_account_for_test, state_traits::StateMut,
1227        },
1228        solana_clock::Clock,
1229        solana_epoch_schedule::EpochSchedule,
1230        solana_instruction::{AccountMeta, error::InstructionError},
1231        solana_program_runtime::{
1232            invoke_context::mock_process_instruction, vm::calculate_heap_cost,
1233            with_mock_invoke_context,
1234        },
1235        solana_pubkey::Pubkey,
1236        solana_rent::Rent,
1237        solana_sbpf::program::BuiltinProgram,
1238        solana_sdk_ids::{system_program, sysvar},
1239        solana_svm_type_overrides::sync::atomic::Ordering,
1240        std::{fs::File, io::Read, ops::Range, sync::atomic::AtomicU64},
1241    };
1242
1243    fn process_instruction(
1244        loader_id: &Pubkey,
1245        program_index: Option<IndexOfAccount>,
1246        instruction_data: &[u8],
1247        transaction_accounts: Vec<(Pubkey, AccountSharedData)>,
1248        instruction_accounts: Vec<AccountMeta>,
1249        expected_result: Result<(), InstructionError>,
1250    ) -> Vec<AccountSharedData> {
1251        mock_process_instruction(
1252            loader_id,
1253            program_index,
1254            instruction_data,
1255            transaction_accounts,
1256            instruction_accounts,
1257            expected_result,
1258            Entrypoint::vm,
1259            |invoke_context| {
1260                test_utils::load_all_invoked_programs(invoke_context);
1261            },
1262            |_invoke_context| {},
1263        )
1264    }
1265
1266    fn load_program_account_from_elf(loader_id: &Pubkey, path: &str) -> AccountSharedData {
1267        let mut file = File::open(path).expect("file open failed");
1268        let mut elf = Vec::new();
1269        file.read_to_end(&mut elf).unwrap();
1270        let rent = Rent::default();
1271        let mut program_account =
1272            AccountSharedData::new(rent.minimum_balance(elf.len()), 0, loader_id);
1273        program_account.set_data(elf);
1274        program_account.set_executable(true);
1275        program_account
1276    }
1277
1278    #[test]
1279    fn test_bpf_loader_invoke_main() {
1280        let loader_id = bpf_loader::id();
1281        let program_id = Pubkey::new_unique();
1282        let program_account =
1283            load_program_account_from_elf(&loader_id, "test_elfs/out/sbpfv3_return_ok.so");
1284        let parameter_id = Pubkey::new_unique();
1285        let parameter_account = AccountSharedData::new(1, 0, &loader_id);
1286        let parameter_meta = AccountMeta {
1287            pubkey: parameter_id,
1288            is_signer: false,
1289            is_writable: false,
1290        };
1291
1292        // Case: No program account
1293        process_instruction(
1294            &loader_id,
1295            None,
1296            &[],
1297            Vec::new(),
1298            Vec::new(),
1299            Err(InstructionError::UnsupportedProgramId),
1300        );
1301
1302        // Case: Only a program account
1303        process_instruction(
1304            &loader_id,
1305            Some(0),
1306            &[],
1307            vec![(program_id, program_account.clone())],
1308            Vec::new(),
1309            Ok(()),
1310        );
1311
1312        // Case: With program and parameter account
1313        process_instruction(
1314            &loader_id,
1315            Some(0),
1316            &[],
1317            vec![
1318                (program_id, program_account.clone()),
1319                (parameter_id, parameter_account.clone()),
1320            ],
1321            vec![parameter_meta.clone()],
1322            Ok(()),
1323        );
1324
1325        // Case: With duplicate accounts
1326        process_instruction(
1327            &loader_id,
1328            Some(0),
1329            &[],
1330            vec![
1331                (program_id, program_account.clone()),
1332                (parameter_id, parameter_account.clone()),
1333            ],
1334            vec![parameter_meta.clone(), parameter_meta],
1335            Ok(()),
1336        );
1337
1338        // Case: limited budget
1339        mock_process_instruction(
1340            &loader_id,
1341            Some(0),
1342            &[],
1343            vec![(program_id, program_account)],
1344            Vec::new(),
1345            Err(InstructionError::ProgramFailedToComplete),
1346            Entrypoint::vm,
1347            |invoke_context| {
1348                invoke_context.mock_set_remaining(0);
1349                test_utils::load_all_invoked_programs(invoke_context);
1350            },
1351            |_invoke_context| {},
1352        );
1353
1354        // Case: Account not a program
1355        mock_process_instruction(
1356            &loader_id,
1357            Some(0),
1358            &[],
1359            vec![(program_id, parameter_account.clone())],
1360            Vec::new(),
1361            Err(InstructionError::UnsupportedProgramId),
1362            Entrypoint::vm,
1363            |invoke_context| {
1364                test_utils::load_all_invoked_programs(invoke_context);
1365            },
1366            |_invoke_context| {},
1367        );
1368        process_instruction(
1369            &loader_id,
1370            Some(0),
1371            &[],
1372            vec![(program_id, parameter_account)],
1373            Vec::new(),
1374            Err(InstructionError::UnsupportedProgramId),
1375        );
1376    }
1377
1378    #[test]
1379    fn test_bpf_loader_serialize_unaligned() {
1380        let loader_id = bpf_loader_deprecated::id();
1381        let program_id = Pubkey::new_unique();
1382        let program_account =
1383            load_program_account_from_elf(&loader_id, "test_elfs/out/noop_unaligned.so");
1384        let parameter_id = Pubkey::new_unique();
1385        let parameter_account = AccountSharedData::new(1, 0, &loader_id);
1386        let parameter_meta = AccountMeta {
1387            pubkey: parameter_id,
1388            is_signer: false,
1389            is_writable: false,
1390        };
1391
1392        // Case: With program and parameter account
1393        process_instruction(
1394            &loader_id,
1395            Some(0),
1396            &[],
1397            vec![
1398                (program_id, program_account.clone()),
1399                (parameter_id, parameter_account.clone()),
1400            ],
1401            vec![parameter_meta.clone()],
1402            Ok(()),
1403        );
1404
1405        // Case: With duplicate accounts
1406        process_instruction(
1407            &loader_id,
1408            Some(0),
1409            &[],
1410            vec![
1411                (program_id, program_account),
1412                (parameter_id, parameter_account),
1413            ],
1414            vec![parameter_meta.clone(), parameter_meta],
1415            Ok(()),
1416        );
1417    }
1418
1419    #[test]
1420    fn test_bpf_loader_serialize_aligned() {
1421        let loader_id = bpf_loader::id();
1422        let program_id = Pubkey::new_unique();
1423        let program_account =
1424            load_program_account_from_elf(&loader_id, "test_elfs/out/noop_aligned.so");
1425        let parameter_id = Pubkey::new_unique();
1426        let parameter_account = AccountSharedData::new(1, 0, &loader_id);
1427        let parameter_meta = AccountMeta {
1428            pubkey: parameter_id,
1429            is_signer: false,
1430            is_writable: false,
1431        };
1432
1433        // Case: With program and parameter account
1434        process_instruction(
1435            &loader_id,
1436            Some(0),
1437            &[],
1438            vec![
1439                (program_id, program_account.clone()),
1440                (parameter_id, parameter_account.clone()),
1441            ],
1442            vec![parameter_meta.clone()],
1443            Ok(()),
1444        );
1445
1446        // Case: With duplicate accounts
1447        process_instruction(
1448            &loader_id,
1449            Some(0),
1450            &[],
1451            vec![
1452                (program_id, program_account),
1453                (parameter_id, parameter_account),
1454            ],
1455            vec![parameter_meta.clone(), parameter_meta],
1456            Ok(()),
1457        );
1458    }
1459
1460    #[test]
1461    fn test_bpf_loader_upgradeable_initialize_buffer() {
1462        let loader_id = bpf_loader_upgradeable::id();
1463        let buffer_address = Pubkey::new_unique();
1464        let buffer_account =
1465            AccountSharedData::new(1, UpgradeableLoaderState::size_of_buffer(9), &loader_id);
1466        let authority_address = Pubkey::new_unique();
1467        let authority_account =
1468            AccountSharedData::new(1, UpgradeableLoaderState::size_of_buffer(9), &loader_id);
1469        let instruction_data =
1470            bincode::serialize(&UpgradeableLoaderInstruction::InitializeBuffer).unwrap();
1471        let instruction_accounts = vec![
1472            AccountMeta {
1473                pubkey: buffer_address,
1474                is_signer: false,
1475                is_writable: true,
1476            },
1477            AccountMeta {
1478                pubkey: authority_address,
1479                is_signer: false,
1480                is_writable: false,
1481            },
1482        ];
1483
1484        // Case: Success
1485        let accounts = process_instruction(
1486            &loader_id,
1487            None,
1488            &instruction_data,
1489            vec![
1490                (buffer_address, buffer_account),
1491                (authority_address, authority_account),
1492            ],
1493            instruction_accounts.clone(),
1494            Ok(()),
1495        );
1496        let state: UpgradeableLoaderState = accounts.first().unwrap().state().unwrap();
1497        assert_eq!(
1498            state,
1499            UpgradeableLoaderState::Buffer {
1500                authority_address: Some(authority_address)
1501            }
1502        );
1503
1504        // Case: Already initialized
1505        let accounts = process_instruction(
1506            &loader_id,
1507            None,
1508            &instruction_data,
1509            vec![
1510                (buffer_address, accounts.first().unwrap().clone()),
1511                (authority_address, accounts.get(1).unwrap().clone()),
1512            ],
1513            instruction_accounts,
1514            Err(InstructionError::AccountAlreadyInitialized),
1515        );
1516        let state: UpgradeableLoaderState = accounts.first().unwrap().state().unwrap();
1517        assert_eq!(
1518            state,
1519            UpgradeableLoaderState::Buffer {
1520                authority_address: Some(authority_address)
1521            }
1522        );
1523    }
1524
1525    #[test]
1526    fn test_bpf_loader_upgradeable_write() {
1527        let loader_id = bpf_loader_upgradeable::id();
1528        let buffer_address = Pubkey::new_unique();
1529        let mut buffer_account =
1530            AccountSharedData::new(1, UpgradeableLoaderState::size_of_buffer(9), &loader_id);
1531        let instruction_accounts = vec![
1532            AccountMeta {
1533                pubkey: buffer_address,
1534                is_signer: false,
1535                is_writable: true,
1536            },
1537            AccountMeta {
1538                pubkey: buffer_address,
1539                is_signer: true,
1540                is_writable: false,
1541            },
1542        ];
1543
1544        // Case: Not initialized
1545        let instruction = bincode::serialize(&UpgradeableLoaderInstruction::Write {
1546            offset: 0,
1547            bytes: vec![42; 9],
1548        })
1549        .unwrap();
1550        process_instruction(
1551            &loader_id,
1552            None,
1553            &instruction,
1554            vec![(buffer_address, buffer_account.clone())],
1555            instruction_accounts.clone(),
1556            Err(InstructionError::InvalidAccountData),
1557        );
1558
1559        // Case: Write entire buffer
1560        let instruction = bincode::serialize(&UpgradeableLoaderInstruction::Write {
1561            offset: 0,
1562            bytes: vec![42; 9],
1563        })
1564        .unwrap();
1565        buffer_account
1566            .set_state(&UpgradeableLoaderState::Buffer {
1567                authority_address: Some(buffer_address),
1568            })
1569            .unwrap();
1570        let accounts = process_instruction(
1571            &loader_id,
1572            None,
1573            &instruction,
1574            vec![(buffer_address, buffer_account.clone())],
1575            instruction_accounts.clone(),
1576            Ok(()),
1577        );
1578        let state: UpgradeableLoaderState = accounts.first().unwrap().state().unwrap();
1579        assert_eq!(
1580            state,
1581            UpgradeableLoaderState::Buffer {
1582                authority_address: Some(buffer_address)
1583            }
1584        );
1585        assert_eq!(
1586            &accounts
1587                .first()
1588                .unwrap()
1589                .data()
1590                .get(UpgradeableLoaderState::size_of_buffer_metadata()..)
1591                .unwrap(),
1592            &[42; 9]
1593        );
1594
1595        // Case: Write portion of the buffer
1596        let instruction = bincode::serialize(&UpgradeableLoaderInstruction::Write {
1597            offset: 3,
1598            bytes: vec![42; 6],
1599        })
1600        .unwrap();
1601        let mut buffer_account =
1602            AccountSharedData::new(1, UpgradeableLoaderState::size_of_buffer(9), &loader_id);
1603        buffer_account
1604            .set_state(&UpgradeableLoaderState::Buffer {
1605                authority_address: Some(buffer_address),
1606            })
1607            .unwrap();
1608        let accounts = process_instruction(
1609            &loader_id,
1610            None,
1611            &instruction,
1612            vec![(buffer_address, buffer_account.clone())],
1613            instruction_accounts.clone(),
1614            Ok(()),
1615        );
1616        let state: UpgradeableLoaderState = accounts.first().unwrap().state().unwrap();
1617        assert_eq!(
1618            state,
1619            UpgradeableLoaderState::Buffer {
1620                authority_address: Some(buffer_address)
1621            }
1622        );
1623        assert_eq!(
1624            &accounts
1625                .first()
1626                .unwrap()
1627                .data()
1628                .get(UpgradeableLoaderState::size_of_buffer_metadata()..)
1629                .unwrap(),
1630            &[0, 0, 0, 42, 42, 42, 42, 42, 42]
1631        );
1632
1633        // Case: overflow size
1634        let instruction = bincode::serialize(&UpgradeableLoaderInstruction::Write {
1635            offset: 0,
1636            bytes: vec![42; 10],
1637        })
1638        .unwrap();
1639        buffer_account
1640            .set_state(&UpgradeableLoaderState::Buffer {
1641                authority_address: Some(buffer_address),
1642            })
1643            .unwrap();
1644        process_instruction(
1645            &loader_id,
1646            None,
1647            &instruction,
1648            vec![(buffer_address, buffer_account.clone())],
1649            instruction_accounts.clone(),
1650            Err(InstructionError::AccountDataTooSmall),
1651        );
1652
1653        // Case: overflow offset
1654        let instruction = bincode::serialize(&UpgradeableLoaderInstruction::Write {
1655            offset: 1,
1656            bytes: vec![42; 9],
1657        })
1658        .unwrap();
1659        buffer_account
1660            .set_state(&UpgradeableLoaderState::Buffer {
1661                authority_address: Some(buffer_address),
1662            })
1663            .unwrap();
1664        process_instruction(
1665            &loader_id,
1666            None,
1667            &instruction,
1668            vec![(buffer_address, buffer_account.clone())],
1669            instruction_accounts.clone(),
1670            Err(InstructionError::AccountDataTooSmall),
1671        );
1672
1673        // Case: Not signed
1674        let instruction = bincode::serialize(&UpgradeableLoaderInstruction::Write {
1675            offset: 0,
1676            bytes: vec![42; 9],
1677        })
1678        .unwrap();
1679        buffer_account
1680            .set_state(&UpgradeableLoaderState::Buffer {
1681                authority_address: Some(buffer_address),
1682            })
1683            .unwrap();
1684        process_instruction(
1685            &loader_id,
1686            None,
1687            &instruction,
1688            vec![(buffer_address, buffer_account.clone())],
1689            vec![
1690                AccountMeta {
1691                    pubkey: buffer_address,
1692                    is_signer: false,
1693                    is_writable: false,
1694                },
1695                AccountMeta {
1696                    pubkey: buffer_address,
1697                    is_signer: false,
1698                    is_writable: false,
1699                },
1700            ],
1701            Err(InstructionError::MissingRequiredSignature),
1702        );
1703
1704        // Case: wrong authority
1705        let authority_address = Pubkey::new_unique();
1706        let instruction = bincode::serialize(&UpgradeableLoaderInstruction::Write {
1707            offset: 1,
1708            bytes: vec![42; 9],
1709        })
1710        .unwrap();
1711        buffer_account
1712            .set_state(&UpgradeableLoaderState::Buffer {
1713                authority_address: Some(buffer_address),
1714            })
1715            .unwrap();
1716        process_instruction(
1717            &loader_id,
1718            None,
1719            &instruction,
1720            vec![
1721                (buffer_address, buffer_account.clone()),
1722                (authority_address, buffer_account.clone()),
1723            ],
1724            vec![
1725                AccountMeta {
1726                    pubkey: buffer_address,
1727                    is_signer: false,
1728                    is_writable: false,
1729                },
1730                AccountMeta {
1731                    pubkey: authority_address,
1732                    is_signer: false,
1733                    is_writable: false,
1734                },
1735            ],
1736            Err(InstructionError::IncorrectAuthority),
1737        );
1738
1739        // Case: None authority
1740        let instruction = bincode::serialize(&UpgradeableLoaderInstruction::Write {
1741            offset: 1,
1742            bytes: vec![42; 9],
1743        })
1744        .unwrap();
1745        buffer_account
1746            .set_state(&UpgradeableLoaderState::Buffer {
1747                authority_address: None,
1748            })
1749            .unwrap();
1750        process_instruction(
1751            &loader_id,
1752            None,
1753            &instruction,
1754            vec![(buffer_address, buffer_account.clone())],
1755            instruction_accounts,
1756            Err(InstructionError::Immutable),
1757        );
1758    }
1759
1760    fn truncate_data(account: &mut AccountSharedData, len: usize) {
1761        let mut data = account.data().to_vec();
1762        data.truncate(len);
1763        account.set_data(data);
1764    }
1765
1766    #[test]
1767    fn test_bpf_loader_upgradeable_upgrade() {
1768        let mut file = File::open("test_elfs/out/sbpfv3_return_ok.so").expect("file open failed");
1769        let mut elf_orig = Vec::new();
1770        file.read_to_end(&mut elf_orig).unwrap();
1771        let mut file = File::open("test_elfs/out/sbpfv3_return_err.so").expect("file open failed");
1772        let mut elf_new = Vec::new();
1773        file.read_to_end(&mut elf_new).unwrap();
1774        assert_ne!(elf_orig.len(), elf_new.len());
1775        const SLOT: u64 = 42;
1776        let buffer_address = Pubkey::new_unique();
1777        let upgrade_authority_address = Pubkey::new_unique();
1778
1779        fn get_accounts(
1780            buffer_address: &Pubkey,
1781            buffer_authority: &Pubkey,
1782            upgrade_authority_address: &Pubkey,
1783            elf_orig: &[u8],
1784            elf_new: &[u8],
1785        ) -> (Vec<(Pubkey, AccountSharedData)>, Vec<AccountMeta>) {
1786            let loader_id = bpf_loader_upgradeable::id();
1787            let program_address = Pubkey::new_unique();
1788            let spill_address = Pubkey::new_unique();
1789            let rent = Rent::default();
1790            let min_program_balance =
1791                1.max(rent.minimum_balance(UpgradeableLoaderState::size_of_program()));
1792            let min_programdata_balance = 1.max(rent.minimum_balance(
1793                UpgradeableLoaderState::size_of_programdata(elf_orig.len().max(elf_new.len())),
1794            ));
1795            let (programdata_address, _) =
1796                Pubkey::find_program_address(&[program_address.as_ref()], &loader_id);
1797            let mut buffer_account = AccountSharedData::new(
1798                1,
1799                UpgradeableLoaderState::size_of_buffer(elf_new.len()),
1800                &bpf_loader_upgradeable::id(),
1801            );
1802            buffer_account
1803                .set_state(&UpgradeableLoaderState::Buffer {
1804                    authority_address: Some(*buffer_authority),
1805                })
1806                .unwrap();
1807            buffer_account
1808                .data_as_mut_slice()
1809                .get_mut(UpgradeableLoaderState::size_of_buffer_metadata()..)
1810                .unwrap()
1811                .copy_from_slice(elf_new);
1812            let mut programdata_account = AccountSharedData::new(
1813                min_programdata_balance,
1814                UpgradeableLoaderState::size_of_programdata(elf_orig.len().max(elf_new.len())),
1815                &bpf_loader_upgradeable::id(),
1816            );
1817            programdata_account
1818                .set_state(&UpgradeableLoaderState::ProgramData {
1819                    slot: SLOT,
1820                    upgrade_authority_address: Some(*upgrade_authority_address),
1821                })
1822                .unwrap();
1823            let mut program_account = AccountSharedData::new(
1824                min_program_balance,
1825                UpgradeableLoaderState::size_of_program(),
1826                &bpf_loader_upgradeable::id(),
1827            );
1828            program_account.set_executable(true);
1829            program_account
1830                .set_state(&UpgradeableLoaderState::Program {
1831                    programdata_address,
1832                })
1833                .unwrap();
1834            let spill_account = AccountSharedData::new(0, 0, &Pubkey::new_unique());
1835            let rent_account = create_account_for_test(&rent);
1836            let clock_account = create_account_for_test(&Clock {
1837                slot: SLOT.saturating_add(1),
1838                ..Clock::default()
1839            });
1840            let upgrade_authority_account = AccountSharedData::new(1, 0, &Pubkey::new_unique());
1841            let transaction_accounts = vec![
1842                (programdata_address, programdata_account),
1843                (program_address, program_account),
1844                (*buffer_address, buffer_account),
1845                (spill_address, spill_account),
1846                (sysvar::rent::id(), rent_account),
1847                (sysvar::clock::id(), clock_account),
1848                (*upgrade_authority_address, upgrade_authority_account),
1849            ];
1850            let instruction_accounts = vec![
1851                AccountMeta {
1852                    pubkey: programdata_address,
1853                    is_signer: false,
1854                    is_writable: true,
1855                },
1856                AccountMeta {
1857                    pubkey: program_address,
1858                    is_signer: false,
1859                    is_writable: true,
1860                },
1861                AccountMeta {
1862                    pubkey: *buffer_address,
1863                    is_signer: false,
1864                    is_writable: true,
1865                },
1866                AccountMeta {
1867                    pubkey: spill_address,
1868                    is_signer: false,
1869                    is_writable: true,
1870                },
1871                AccountMeta {
1872                    pubkey: sysvar::rent::id(),
1873                    is_signer: false,
1874                    is_writable: false,
1875                },
1876                AccountMeta {
1877                    pubkey: sysvar::clock::id(),
1878                    is_signer: false,
1879                    is_writable: false,
1880                },
1881                AccountMeta {
1882                    pubkey: *upgrade_authority_address,
1883                    is_signer: true,
1884                    is_writable: false,
1885                },
1886            ];
1887            (transaction_accounts, instruction_accounts)
1888        }
1889
1890        fn process_instruction(
1891            transaction_accounts: Vec<(Pubkey, AccountSharedData)>,
1892            instruction_accounts: Vec<AccountMeta>,
1893            expected_result: Result<(), InstructionError>,
1894        ) -> Vec<AccountSharedData> {
1895            let instruction_data =
1896                bincode::serialize(&UpgradeableLoaderInstruction::Upgrade).unwrap();
1897            mock_process_instruction(
1898                &bpf_loader_upgradeable::id(),
1899                None,
1900                &instruction_data,
1901                transaction_accounts,
1902                instruction_accounts,
1903                expected_result,
1904                Entrypoint::vm,
1905                |_invoke_context| {},
1906                |_invoke_context| {},
1907            )
1908        }
1909
1910        // Case: Success
1911        let (transaction_accounts, instruction_accounts) = get_accounts(
1912            &buffer_address,
1913            &upgrade_authority_address,
1914            &upgrade_authority_address,
1915            &elf_orig,
1916            &elf_new,
1917        );
1918        let accounts = process_instruction(transaction_accounts, instruction_accounts, Ok(()));
1919        let min_programdata_balance = Rent::default().minimum_balance(
1920            UpgradeableLoaderState::size_of_programdata(elf_orig.len().max(elf_new.len())),
1921        );
1922        assert_eq!(
1923            min_programdata_balance,
1924            accounts.first().unwrap().lamports()
1925        );
1926        assert_eq!(0, accounts.get(2).unwrap().lamports());
1927        assert_eq!(1, accounts.get(3).unwrap().lamports());
1928        let state: UpgradeableLoaderState = accounts.first().unwrap().state().unwrap();
1929        assert_eq!(
1930            state,
1931            UpgradeableLoaderState::ProgramData {
1932                slot: SLOT.saturating_add(1),
1933                upgrade_authority_address: Some(upgrade_authority_address)
1934            }
1935        );
1936        for (i, byte) in accounts
1937            .first()
1938            .unwrap()
1939            .data()
1940            .get(
1941                UpgradeableLoaderState::size_of_programdata_metadata()
1942                    ..UpgradeableLoaderState::size_of_programdata(elf_new.len()),
1943            )
1944            .unwrap()
1945            .iter()
1946            .enumerate()
1947        {
1948            assert_eq!(*elf_new.get(i).unwrap(), *byte);
1949        }
1950
1951        // Case: not upgradable
1952        let (mut transaction_accounts, instruction_accounts) = get_accounts(
1953            &buffer_address,
1954            &upgrade_authority_address,
1955            &upgrade_authority_address,
1956            &elf_orig,
1957            &elf_new,
1958        );
1959        transaction_accounts
1960            .get_mut(0)
1961            .unwrap()
1962            .1
1963            .set_state(&UpgradeableLoaderState::ProgramData {
1964                slot: SLOT,
1965                upgrade_authority_address: None,
1966            })
1967            .unwrap();
1968        process_instruction(
1969            transaction_accounts,
1970            instruction_accounts,
1971            Err(InstructionError::Immutable),
1972        );
1973
1974        // Case: wrong authority
1975        let (mut transaction_accounts, mut instruction_accounts) = get_accounts(
1976            &buffer_address,
1977            &upgrade_authority_address,
1978            &upgrade_authority_address,
1979            &elf_orig,
1980            &elf_new,
1981        );
1982        let invalid_upgrade_authority_address = Pubkey::new_unique();
1983        transaction_accounts.get_mut(6).unwrap().0 = invalid_upgrade_authority_address;
1984        instruction_accounts.get_mut(6).unwrap().pubkey = invalid_upgrade_authority_address;
1985        process_instruction(
1986            transaction_accounts,
1987            instruction_accounts,
1988            Err(InstructionError::IncorrectAuthority),
1989        );
1990
1991        // Case: authority did not sign
1992        let (transaction_accounts, mut instruction_accounts) = get_accounts(
1993            &buffer_address,
1994            &upgrade_authority_address,
1995            &upgrade_authority_address,
1996            &elf_orig,
1997            &elf_new,
1998        );
1999        instruction_accounts.get_mut(6).unwrap().is_signer = false;
2000        process_instruction(
2001            transaction_accounts,
2002            instruction_accounts,
2003            Err(InstructionError::MissingRequiredSignature),
2004        );
2005
2006        // Case: Buffer account and spill account alias
2007        let (transaction_accounts, mut instruction_accounts) = get_accounts(
2008            &buffer_address,
2009            &upgrade_authority_address,
2010            &upgrade_authority_address,
2011            &elf_orig,
2012            &elf_new,
2013        );
2014        *instruction_accounts.get_mut(3).unwrap() = instruction_accounts.get(2).unwrap().clone();
2015        process_instruction(
2016            transaction_accounts,
2017            instruction_accounts,
2018            Err(InstructionError::AccountBorrowFailed),
2019        );
2020
2021        // Case: Programdata account and spill account alias
2022        let (transaction_accounts, mut instruction_accounts) = get_accounts(
2023            &buffer_address,
2024            &upgrade_authority_address,
2025            &upgrade_authority_address,
2026            &elf_orig,
2027            &elf_new,
2028        );
2029        *instruction_accounts.get_mut(3).unwrap() = instruction_accounts.first().unwrap().clone();
2030        process_instruction(
2031            transaction_accounts,
2032            instruction_accounts,
2033            Err(InstructionError::AccountBorrowFailed),
2034        );
2035
2036        // Case: Program account not a program
2037        let (transaction_accounts, mut instruction_accounts) = get_accounts(
2038            &buffer_address,
2039            &upgrade_authority_address,
2040            &upgrade_authority_address,
2041            &elf_orig,
2042            &elf_new,
2043        );
2044        *instruction_accounts.get_mut(1).unwrap() = instruction_accounts.get(2).unwrap().clone();
2045        let instruction_data = bincode::serialize(&UpgradeableLoaderInstruction::Upgrade).unwrap();
2046
2047        mock_process_instruction(
2048            &bpf_loader_upgradeable::id(),
2049            None,
2050            &instruction_data,
2051            transaction_accounts.clone(),
2052            instruction_accounts.clone(),
2053            Err(InstructionError::InvalidAccountData),
2054            Entrypoint::vm,
2055            |invoke_context| {
2056                test_utils::load_all_invoked_programs(invoke_context);
2057            },
2058            |_invoke_context| {},
2059        );
2060        process_instruction(
2061            transaction_accounts.clone(),
2062            instruction_accounts.clone(),
2063            Err(InstructionError::InvalidAccountData),
2064        );
2065
2066        // Case: Program account now owned by loader
2067        let (mut transaction_accounts, instruction_accounts) = get_accounts(
2068            &buffer_address,
2069            &upgrade_authority_address,
2070            &upgrade_authority_address,
2071            &elf_orig,
2072            &elf_new,
2073        );
2074        transaction_accounts
2075            .get_mut(1)
2076            .unwrap()
2077            .1
2078            .set_owner(Pubkey::new_unique());
2079        process_instruction(
2080            transaction_accounts,
2081            instruction_accounts,
2082            Err(InstructionError::IncorrectProgramId),
2083        );
2084
2085        // Case: Program account not writable
2086        let (transaction_accounts, mut instruction_accounts) = get_accounts(
2087            &buffer_address,
2088            &upgrade_authority_address,
2089            &upgrade_authority_address,
2090            &elf_orig,
2091            &elf_new,
2092        );
2093        instruction_accounts.get_mut(1).unwrap().is_writable = false;
2094        process_instruction(
2095            transaction_accounts,
2096            instruction_accounts,
2097            Err(InstructionError::InvalidArgument),
2098        );
2099
2100        // Case: Program account not initialized
2101        let (mut transaction_accounts, instruction_accounts) = get_accounts(
2102            &buffer_address,
2103            &upgrade_authority_address,
2104            &upgrade_authority_address,
2105            &elf_orig,
2106            &elf_new,
2107        );
2108        transaction_accounts
2109            .get_mut(1)
2110            .unwrap()
2111            .1
2112            .set_state(&UpgradeableLoaderState::Uninitialized)
2113            .unwrap();
2114        process_instruction(
2115            transaction_accounts,
2116            instruction_accounts,
2117            Err(InstructionError::InvalidAccountData),
2118        );
2119
2120        // Case: Program ProgramData account mismatch
2121        let (mut transaction_accounts, mut instruction_accounts) = get_accounts(
2122            &buffer_address,
2123            &upgrade_authority_address,
2124            &upgrade_authority_address,
2125            &elf_orig,
2126            &elf_new,
2127        );
2128        let invalid_programdata_address = Pubkey::new_unique();
2129        transaction_accounts.get_mut(0).unwrap().0 = invalid_programdata_address;
2130        instruction_accounts.get_mut(0).unwrap().pubkey = invalid_programdata_address;
2131        process_instruction(
2132            transaction_accounts,
2133            instruction_accounts,
2134            Err(InstructionError::InvalidArgument),
2135        );
2136
2137        // Case: Buffer account not initialized
2138        let (mut transaction_accounts, instruction_accounts) = get_accounts(
2139            &buffer_address,
2140            &upgrade_authority_address,
2141            &upgrade_authority_address,
2142            &elf_orig,
2143            &elf_new,
2144        );
2145        transaction_accounts
2146            .get_mut(2)
2147            .unwrap()
2148            .1
2149            .set_state(&UpgradeableLoaderState::Uninitialized)
2150            .unwrap();
2151        process_instruction(
2152            transaction_accounts,
2153            instruction_accounts,
2154            Err(InstructionError::InvalidArgument),
2155        );
2156
2157        // Case: Buffer account too big
2158        let (mut transaction_accounts, instruction_accounts) = get_accounts(
2159            &buffer_address,
2160            &upgrade_authority_address,
2161            &upgrade_authority_address,
2162            &elf_orig,
2163            &elf_new,
2164        );
2165        transaction_accounts.get_mut(2).unwrap().1 = AccountSharedData::new(
2166            1,
2167            UpgradeableLoaderState::size_of_buffer(
2168                elf_orig.len().max(elf_new.len()).saturating_add(1),
2169            ),
2170            &bpf_loader_upgradeable::id(),
2171        );
2172        transaction_accounts
2173            .get_mut(2)
2174            .unwrap()
2175            .1
2176            .set_state(&UpgradeableLoaderState::Buffer {
2177                authority_address: Some(upgrade_authority_address),
2178            })
2179            .unwrap();
2180        process_instruction(
2181            transaction_accounts,
2182            instruction_accounts,
2183            Err(InstructionError::AccountDataTooSmall),
2184        );
2185
2186        // Case: Buffer account too small
2187        let (mut transaction_accounts, instruction_accounts) = get_accounts(
2188            &buffer_address,
2189            &upgrade_authority_address,
2190            &upgrade_authority_address,
2191            &elf_orig,
2192            &elf_new,
2193        );
2194        transaction_accounts
2195            .get_mut(2)
2196            .unwrap()
2197            .1
2198            .set_state(&UpgradeableLoaderState::Buffer {
2199                authority_address: Some(upgrade_authority_address),
2200            })
2201            .unwrap();
2202        truncate_data(&mut transaction_accounts.get_mut(2).unwrap().1, 5);
2203        process_instruction(
2204            transaction_accounts,
2205            instruction_accounts,
2206            Err(InstructionError::InvalidAccountData),
2207        );
2208
2209        // Case: Mismatched buffer and program authority
2210        let (transaction_accounts, instruction_accounts) = get_accounts(
2211            &buffer_address,
2212            &buffer_address,
2213            &upgrade_authority_address,
2214            &elf_orig,
2215            &elf_new,
2216        );
2217        process_instruction(
2218            transaction_accounts,
2219            instruction_accounts,
2220            Err(InstructionError::IncorrectAuthority),
2221        );
2222
2223        // Case: No buffer authority
2224        let (mut transaction_accounts, instruction_accounts) = get_accounts(
2225            &buffer_address,
2226            &buffer_address,
2227            &upgrade_authority_address,
2228            &elf_orig,
2229            &elf_new,
2230        );
2231        transaction_accounts
2232            .get_mut(2)
2233            .unwrap()
2234            .1
2235            .set_state(&UpgradeableLoaderState::Buffer {
2236                authority_address: None,
2237            })
2238            .unwrap();
2239        process_instruction(
2240            transaction_accounts,
2241            instruction_accounts,
2242            Err(InstructionError::IncorrectAuthority),
2243        );
2244
2245        // Case: No buffer and program authority
2246        let (mut transaction_accounts, instruction_accounts) = get_accounts(
2247            &buffer_address,
2248            &buffer_address,
2249            &upgrade_authority_address,
2250            &elf_orig,
2251            &elf_new,
2252        );
2253        transaction_accounts
2254            .get_mut(0)
2255            .unwrap()
2256            .1
2257            .set_state(&UpgradeableLoaderState::ProgramData {
2258                slot: SLOT,
2259                upgrade_authority_address: None,
2260            })
2261            .unwrap();
2262        transaction_accounts
2263            .get_mut(2)
2264            .unwrap()
2265            .1
2266            .set_state(&UpgradeableLoaderState::Buffer {
2267                authority_address: None,
2268            })
2269            .unwrap();
2270        process_instruction(
2271            transaction_accounts,
2272            instruction_accounts,
2273            Err(InstructionError::IncorrectAuthority),
2274        );
2275    }
2276
2277    #[test]
2278    fn test_bpf_loader_upgradeable_set_upgrade_authority() {
2279        let instruction = bincode::serialize(&UpgradeableLoaderInstruction::SetAuthority).unwrap();
2280        let loader_id = bpf_loader_upgradeable::id();
2281        let slot = 0;
2282        let upgrade_authority_address = Pubkey::new_unique();
2283        let upgrade_authority_account = AccountSharedData::new(1, 0, &Pubkey::new_unique());
2284        let new_upgrade_authority_address = Pubkey::new_unique();
2285        let new_upgrade_authority_account = AccountSharedData::new(1, 0, &Pubkey::new_unique());
2286        let program_address = Pubkey::new_unique();
2287        let (programdata_address, _) = Pubkey::find_program_address(
2288            &[program_address.as_ref()],
2289            &bpf_loader_upgradeable::id(),
2290        );
2291        let mut programdata_account = AccountSharedData::new(
2292            1,
2293            UpgradeableLoaderState::size_of_programdata(0),
2294            &bpf_loader_upgradeable::id(),
2295        );
2296        programdata_account
2297            .set_state(&UpgradeableLoaderState::ProgramData {
2298                slot,
2299                upgrade_authority_address: Some(upgrade_authority_address),
2300            })
2301            .unwrap();
2302        let programdata_meta = AccountMeta {
2303            pubkey: programdata_address,
2304            is_signer: false,
2305            is_writable: true,
2306        };
2307        let upgrade_authority_meta = AccountMeta {
2308            pubkey: upgrade_authority_address,
2309            is_signer: true,
2310            is_writable: false,
2311        };
2312        let new_upgrade_authority_meta = AccountMeta {
2313            pubkey: new_upgrade_authority_address,
2314            is_signer: false,
2315            is_writable: false,
2316        };
2317
2318        // Case: Set to new authority
2319        let accounts = process_instruction(
2320            &loader_id,
2321            None,
2322            &instruction,
2323            vec![
2324                (programdata_address, programdata_account.clone()),
2325                (upgrade_authority_address, upgrade_authority_account.clone()),
2326                (
2327                    new_upgrade_authority_address,
2328                    new_upgrade_authority_account.clone(),
2329                ),
2330            ],
2331            vec![
2332                programdata_meta.clone(),
2333                upgrade_authority_meta.clone(),
2334                new_upgrade_authority_meta.clone(),
2335            ],
2336            Ok(()),
2337        );
2338        let state: UpgradeableLoaderState = accounts.first().unwrap().state().unwrap();
2339        assert_eq!(
2340            state,
2341            UpgradeableLoaderState::ProgramData {
2342                slot,
2343                upgrade_authority_address: Some(new_upgrade_authority_address),
2344            }
2345        );
2346
2347        // Case: Not upgradeable
2348        let accounts = process_instruction(
2349            &loader_id,
2350            None,
2351            &instruction,
2352            vec![
2353                (programdata_address, programdata_account.clone()),
2354                (upgrade_authority_address, upgrade_authority_account.clone()),
2355            ],
2356            vec![programdata_meta.clone(), upgrade_authority_meta.clone()],
2357            Ok(()),
2358        );
2359        let state: UpgradeableLoaderState = accounts.first().unwrap().state().unwrap();
2360        assert_eq!(
2361            state,
2362            UpgradeableLoaderState::ProgramData {
2363                slot,
2364                upgrade_authority_address: None,
2365            }
2366        );
2367
2368        // Case: Authority did not sign
2369        process_instruction(
2370            &loader_id,
2371            None,
2372            &instruction,
2373            vec![
2374                (programdata_address, programdata_account.clone()),
2375                (upgrade_authority_address, upgrade_authority_account.clone()),
2376            ],
2377            vec![
2378                programdata_meta.clone(),
2379                AccountMeta {
2380                    pubkey: upgrade_authority_address,
2381                    is_signer: false,
2382                    is_writable: false,
2383                },
2384            ],
2385            Err(InstructionError::MissingRequiredSignature),
2386        );
2387
2388        // Case: wrong authority
2389        let invalid_upgrade_authority_address = Pubkey::new_unique();
2390        process_instruction(
2391            &loader_id,
2392            None,
2393            &instruction,
2394            vec![
2395                (programdata_address, programdata_account.clone()),
2396                (
2397                    invalid_upgrade_authority_address,
2398                    upgrade_authority_account.clone(),
2399                ),
2400                (new_upgrade_authority_address, new_upgrade_authority_account),
2401            ],
2402            vec![
2403                programdata_meta.clone(),
2404                AccountMeta {
2405                    pubkey: invalid_upgrade_authority_address,
2406                    is_signer: true,
2407                    is_writable: false,
2408                },
2409                new_upgrade_authority_meta,
2410            ],
2411            Err(InstructionError::IncorrectAuthority),
2412        );
2413
2414        // Case: No authority
2415        programdata_account
2416            .set_state(&UpgradeableLoaderState::ProgramData {
2417                slot,
2418                upgrade_authority_address: None,
2419            })
2420            .unwrap();
2421        process_instruction(
2422            &loader_id,
2423            None,
2424            &instruction,
2425            vec![
2426                (programdata_address, programdata_account.clone()),
2427                (upgrade_authority_address, upgrade_authority_account.clone()),
2428            ],
2429            vec![programdata_meta.clone(), upgrade_authority_meta.clone()],
2430            Err(InstructionError::Immutable),
2431        );
2432
2433        // Case: Not a ProgramData account
2434        programdata_account
2435            .set_state(&UpgradeableLoaderState::Program {
2436                programdata_address: Pubkey::new_unique(),
2437            })
2438            .unwrap();
2439        process_instruction(
2440            &loader_id,
2441            None,
2442            &instruction,
2443            vec![
2444                (programdata_address, programdata_account.clone()),
2445                (upgrade_authority_address, upgrade_authority_account),
2446            ],
2447            vec![programdata_meta, upgrade_authority_meta],
2448            Err(InstructionError::InvalidArgument),
2449        );
2450    }
2451
2452    #[test]
2453    fn test_bpf_loader_upgradeable_set_upgrade_authority_checked() {
2454        let instruction =
2455            bincode::serialize(&UpgradeableLoaderInstruction::SetAuthorityChecked).unwrap();
2456        let loader_id = bpf_loader_upgradeable::id();
2457        let slot = 0;
2458        let upgrade_authority_address = Pubkey::new_unique();
2459        let upgrade_authority_account = AccountSharedData::new(1, 0, &Pubkey::new_unique());
2460        let new_upgrade_authority_address = Pubkey::new_unique();
2461        let new_upgrade_authority_account = AccountSharedData::new(1, 0, &Pubkey::new_unique());
2462        let program_address = Pubkey::new_unique();
2463        let (programdata_address, _) = Pubkey::find_program_address(
2464            &[program_address.as_ref()],
2465            &bpf_loader_upgradeable::id(),
2466        );
2467        let mut programdata_account = AccountSharedData::new(
2468            1,
2469            UpgradeableLoaderState::size_of_programdata(0),
2470            &bpf_loader_upgradeable::id(),
2471        );
2472        programdata_account
2473            .set_state(&UpgradeableLoaderState::ProgramData {
2474                slot,
2475                upgrade_authority_address: Some(upgrade_authority_address),
2476            })
2477            .unwrap();
2478        let programdata_meta = AccountMeta {
2479            pubkey: programdata_address,
2480            is_signer: false,
2481            is_writable: true,
2482        };
2483        let upgrade_authority_meta = AccountMeta {
2484            pubkey: upgrade_authority_address,
2485            is_signer: true,
2486            is_writable: false,
2487        };
2488        let new_upgrade_authority_meta = AccountMeta {
2489            pubkey: new_upgrade_authority_address,
2490            is_signer: true,
2491            is_writable: false,
2492        };
2493
2494        // Case: Set to new authority
2495        let accounts = process_instruction(
2496            &loader_id,
2497            None,
2498            &instruction,
2499            vec![
2500                (programdata_address, programdata_account.clone()),
2501                (upgrade_authority_address, upgrade_authority_account.clone()),
2502                (
2503                    new_upgrade_authority_address,
2504                    new_upgrade_authority_account.clone(),
2505                ),
2506            ],
2507            vec![
2508                programdata_meta.clone(),
2509                upgrade_authority_meta.clone(),
2510                new_upgrade_authority_meta.clone(),
2511            ],
2512            Ok(()),
2513        );
2514
2515        let state: UpgradeableLoaderState = accounts.first().unwrap().state().unwrap();
2516        assert_eq!(
2517            state,
2518            UpgradeableLoaderState::ProgramData {
2519                slot,
2520                upgrade_authority_address: Some(new_upgrade_authority_address),
2521            }
2522        );
2523
2524        // Case: set to same authority
2525        process_instruction(
2526            &loader_id,
2527            None,
2528            &instruction,
2529            vec![
2530                (programdata_address, programdata_account.clone()),
2531                (upgrade_authority_address, upgrade_authority_account.clone()),
2532            ],
2533            vec![
2534                programdata_meta.clone(),
2535                upgrade_authority_meta.clone(),
2536                upgrade_authority_meta.clone(),
2537            ],
2538            Ok(()),
2539        );
2540
2541        // Case: present authority not in instruction
2542        process_instruction(
2543            &loader_id,
2544            None,
2545            &instruction,
2546            vec![
2547                (programdata_address, programdata_account.clone()),
2548                (upgrade_authority_address, upgrade_authority_account.clone()),
2549                (
2550                    new_upgrade_authority_address,
2551                    new_upgrade_authority_account.clone(),
2552                ),
2553            ],
2554            vec![programdata_meta.clone(), new_upgrade_authority_meta.clone()],
2555            Err(InstructionError::MissingAccount),
2556        );
2557
2558        // Case: new authority not in instruction
2559        process_instruction(
2560            &loader_id,
2561            None,
2562            &instruction,
2563            vec![
2564                (programdata_address, programdata_account.clone()),
2565                (upgrade_authority_address, upgrade_authority_account.clone()),
2566                (
2567                    new_upgrade_authority_address,
2568                    new_upgrade_authority_account.clone(),
2569                ),
2570            ],
2571            vec![programdata_meta.clone(), upgrade_authority_meta.clone()],
2572            Err(InstructionError::MissingAccount),
2573        );
2574
2575        // Case: present authority did not sign
2576        process_instruction(
2577            &loader_id,
2578            None,
2579            &instruction,
2580            vec![
2581                (programdata_address, programdata_account.clone()),
2582                (upgrade_authority_address, upgrade_authority_account.clone()),
2583                (
2584                    new_upgrade_authority_address,
2585                    new_upgrade_authority_account.clone(),
2586                ),
2587            ],
2588            vec![
2589                programdata_meta.clone(),
2590                AccountMeta {
2591                    pubkey: upgrade_authority_address,
2592                    is_signer: false,
2593                    is_writable: false,
2594                },
2595                new_upgrade_authority_meta.clone(),
2596            ],
2597            Err(InstructionError::MissingRequiredSignature),
2598        );
2599
2600        // Case: New authority did not sign
2601        process_instruction(
2602            &loader_id,
2603            None,
2604            &instruction,
2605            vec![
2606                (programdata_address, programdata_account.clone()),
2607                (upgrade_authority_address, upgrade_authority_account.clone()),
2608                (
2609                    new_upgrade_authority_address,
2610                    new_upgrade_authority_account.clone(),
2611                ),
2612            ],
2613            vec![
2614                programdata_meta.clone(),
2615                upgrade_authority_meta.clone(),
2616                AccountMeta {
2617                    pubkey: new_upgrade_authority_address,
2618                    is_signer: false,
2619                    is_writable: false,
2620                },
2621            ],
2622            Err(InstructionError::MissingRequiredSignature),
2623        );
2624
2625        // Case: wrong present authority
2626        let invalid_upgrade_authority_address = Pubkey::new_unique();
2627        process_instruction(
2628            &loader_id,
2629            None,
2630            &instruction,
2631            vec![
2632                (programdata_address, programdata_account.clone()),
2633                (
2634                    invalid_upgrade_authority_address,
2635                    upgrade_authority_account.clone(),
2636                ),
2637                (new_upgrade_authority_address, new_upgrade_authority_account),
2638            ],
2639            vec![
2640                programdata_meta.clone(),
2641                AccountMeta {
2642                    pubkey: invalid_upgrade_authority_address,
2643                    is_signer: true,
2644                    is_writable: false,
2645                },
2646                new_upgrade_authority_meta.clone(),
2647            ],
2648            Err(InstructionError::IncorrectAuthority),
2649        );
2650
2651        // Case: programdata is immutable
2652        programdata_account
2653            .set_state(&UpgradeableLoaderState::ProgramData {
2654                slot,
2655                upgrade_authority_address: None,
2656            })
2657            .unwrap();
2658        process_instruction(
2659            &loader_id,
2660            None,
2661            &instruction,
2662            vec![
2663                (programdata_address, programdata_account.clone()),
2664                (upgrade_authority_address, upgrade_authority_account.clone()),
2665            ],
2666            vec![
2667                programdata_meta.clone(),
2668                upgrade_authority_meta.clone(),
2669                new_upgrade_authority_meta.clone(),
2670            ],
2671            Err(InstructionError::Immutable),
2672        );
2673
2674        // Case: Not a ProgramData account
2675        programdata_account
2676            .set_state(&UpgradeableLoaderState::Program {
2677                programdata_address: Pubkey::new_unique(),
2678            })
2679            .unwrap();
2680        process_instruction(
2681            &loader_id,
2682            None,
2683            &instruction,
2684            vec![
2685                (programdata_address, programdata_account.clone()),
2686                (upgrade_authority_address, upgrade_authority_account),
2687            ],
2688            vec![
2689                programdata_meta,
2690                upgrade_authority_meta,
2691                new_upgrade_authority_meta,
2692            ],
2693            Err(InstructionError::InvalidArgument),
2694        );
2695    }
2696
2697    #[test]
2698    fn test_bpf_loader_upgradeable_set_buffer_authority() {
2699        let instruction = bincode::serialize(&UpgradeableLoaderInstruction::SetAuthority).unwrap();
2700        let loader_id = bpf_loader_upgradeable::id();
2701        let invalid_authority_address = Pubkey::new_unique();
2702        let authority_address = Pubkey::new_unique();
2703        let authority_account = AccountSharedData::new(1, 0, &Pubkey::new_unique());
2704        let new_authority_address = Pubkey::new_unique();
2705        let new_authority_account = AccountSharedData::new(1, 0, &Pubkey::new_unique());
2706        let buffer_address = Pubkey::new_unique();
2707        let mut buffer_account =
2708            AccountSharedData::new(1, UpgradeableLoaderState::size_of_buffer(0), &loader_id);
2709        buffer_account
2710            .set_state(&UpgradeableLoaderState::Buffer {
2711                authority_address: Some(authority_address),
2712            })
2713            .unwrap();
2714        let mut transaction_accounts = vec![
2715            (buffer_address, buffer_account.clone()),
2716            (authority_address, authority_account.clone()),
2717            (new_authority_address, new_authority_account.clone()),
2718        ];
2719        let buffer_meta = AccountMeta {
2720            pubkey: buffer_address,
2721            is_signer: false,
2722            is_writable: true,
2723        };
2724        let authority_meta = AccountMeta {
2725            pubkey: authority_address,
2726            is_signer: true,
2727            is_writable: false,
2728        };
2729        let new_authority_meta = AccountMeta {
2730            pubkey: new_authority_address,
2731            is_signer: false,
2732            is_writable: false,
2733        };
2734
2735        // Case: New authority required
2736        let accounts = process_instruction(
2737            &loader_id,
2738            None,
2739            &instruction,
2740            transaction_accounts.clone(),
2741            vec![buffer_meta.clone(), authority_meta.clone()],
2742            Err(InstructionError::IncorrectAuthority),
2743        );
2744        let state: UpgradeableLoaderState = accounts.first().unwrap().state().unwrap();
2745        assert_eq!(
2746            state,
2747            UpgradeableLoaderState::Buffer {
2748                authority_address: Some(authority_address),
2749            }
2750        );
2751
2752        // Case: Set to new authority
2753        buffer_account
2754            .set_state(&UpgradeableLoaderState::Buffer {
2755                authority_address: Some(authority_address),
2756            })
2757            .unwrap();
2758        let accounts = process_instruction(
2759            &loader_id,
2760            None,
2761            &instruction,
2762            transaction_accounts.clone(),
2763            vec![
2764                buffer_meta.clone(),
2765                authority_meta.clone(),
2766                new_authority_meta.clone(),
2767            ],
2768            Ok(()),
2769        );
2770        let state: UpgradeableLoaderState = accounts.first().unwrap().state().unwrap();
2771        assert_eq!(
2772            state,
2773            UpgradeableLoaderState::Buffer {
2774                authority_address: Some(new_authority_address),
2775            }
2776        );
2777
2778        // Case: Authority did not sign
2779        process_instruction(
2780            &loader_id,
2781            None,
2782            &instruction,
2783            transaction_accounts.clone(),
2784            vec![
2785                buffer_meta.clone(),
2786                AccountMeta {
2787                    pubkey: authority_address,
2788                    is_signer: false,
2789                    is_writable: false,
2790                },
2791                new_authority_meta.clone(),
2792            ],
2793            Err(InstructionError::MissingRequiredSignature),
2794        );
2795
2796        // Case: wrong authority
2797        process_instruction(
2798            &loader_id,
2799            None,
2800            &instruction,
2801            vec![
2802                (buffer_address, buffer_account.clone()),
2803                (invalid_authority_address, authority_account),
2804                (new_authority_address, new_authority_account),
2805            ],
2806            vec![
2807                buffer_meta.clone(),
2808                AccountMeta {
2809                    pubkey: invalid_authority_address,
2810                    is_signer: true,
2811                    is_writable: false,
2812                },
2813                new_authority_meta.clone(),
2814            ],
2815            Err(InstructionError::IncorrectAuthority),
2816        );
2817
2818        // Case: No authority
2819        process_instruction(
2820            &loader_id,
2821            None,
2822            &instruction,
2823            transaction_accounts.clone(),
2824            vec![buffer_meta.clone(), authority_meta.clone()],
2825            Err(InstructionError::IncorrectAuthority),
2826        );
2827
2828        // Case: Set to no authority
2829        transaction_accounts
2830            .get_mut(0)
2831            .unwrap()
2832            .1
2833            .set_state(&UpgradeableLoaderState::Buffer {
2834                authority_address: None,
2835            })
2836            .unwrap();
2837        process_instruction(
2838            &loader_id,
2839            None,
2840            &instruction,
2841            transaction_accounts.clone(),
2842            vec![
2843                buffer_meta.clone(),
2844                authority_meta.clone(),
2845                new_authority_meta.clone(),
2846            ],
2847            Err(InstructionError::Immutable),
2848        );
2849
2850        // Case: Not a Buffer account
2851        transaction_accounts
2852            .get_mut(0)
2853            .unwrap()
2854            .1
2855            .set_state(&UpgradeableLoaderState::Program {
2856                programdata_address: Pubkey::new_unique(),
2857            })
2858            .unwrap();
2859        process_instruction(
2860            &loader_id,
2861            None,
2862            &instruction,
2863            transaction_accounts.clone(),
2864            vec![buffer_meta, authority_meta, new_authority_meta],
2865            Err(InstructionError::InvalidArgument),
2866        );
2867    }
2868
2869    #[test]
2870    fn test_bpf_loader_upgradeable_set_buffer_authority_checked() {
2871        let instruction =
2872            bincode::serialize(&UpgradeableLoaderInstruction::SetAuthorityChecked).unwrap();
2873        let loader_id = bpf_loader_upgradeable::id();
2874        let invalid_authority_address = Pubkey::new_unique();
2875        let authority_address = Pubkey::new_unique();
2876        let authority_account = AccountSharedData::new(1, 0, &Pubkey::new_unique());
2877        let new_authority_address = Pubkey::new_unique();
2878        let new_authority_account = AccountSharedData::new(1, 0, &Pubkey::new_unique());
2879        let buffer_address = Pubkey::new_unique();
2880        let mut buffer_account =
2881            AccountSharedData::new(1, UpgradeableLoaderState::size_of_buffer(0), &loader_id);
2882        buffer_account
2883            .set_state(&UpgradeableLoaderState::Buffer {
2884                authority_address: Some(authority_address),
2885            })
2886            .unwrap();
2887        let mut transaction_accounts = vec![
2888            (buffer_address, buffer_account.clone()),
2889            (authority_address, authority_account.clone()),
2890            (new_authority_address, new_authority_account.clone()),
2891        ];
2892        let buffer_meta = AccountMeta {
2893            pubkey: buffer_address,
2894            is_signer: false,
2895            is_writable: true,
2896        };
2897        let authority_meta = AccountMeta {
2898            pubkey: authority_address,
2899            is_signer: true,
2900            is_writable: false,
2901        };
2902        let new_authority_meta = AccountMeta {
2903            pubkey: new_authority_address,
2904            is_signer: true,
2905            is_writable: false,
2906        };
2907
2908        // Case: Set to new authority
2909        buffer_account
2910            .set_state(&UpgradeableLoaderState::Buffer {
2911                authority_address: Some(authority_address),
2912            })
2913            .unwrap();
2914        let accounts = process_instruction(
2915            &loader_id,
2916            None,
2917            &instruction,
2918            transaction_accounts.clone(),
2919            vec![
2920                buffer_meta.clone(),
2921                authority_meta.clone(),
2922                new_authority_meta.clone(),
2923            ],
2924            Ok(()),
2925        );
2926        let state: UpgradeableLoaderState = accounts.first().unwrap().state().unwrap();
2927        assert_eq!(
2928            state,
2929            UpgradeableLoaderState::Buffer {
2930                authority_address: Some(new_authority_address),
2931            }
2932        );
2933
2934        // Case: set to same authority
2935        process_instruction(
2936            &loader_id,
2937            None,
2938            &instruction,
2939            transaction_accounts.clone(),
2940            vec![
2941                buffer_meta.clone(),
2942                authority_meta.clone(),
2943                authority_meta.clone(),
2944            ],
2945            Ok(()),
2946        );
2947
2948        // Case: Missing current authority
2949        process_instruction(
2950            &loader_id,
2951            None,
2952            &instruction,
2953            transaction_accounts.clone(),
2954            vec![buffer_meta.clone(), new_authority_meta.clone()],
2955            Err(InstructionError::MissingAccount),
2956        );
2957
2958        // Case: Missing new authority
2959        process_instruction(
2960            &loader_id,
2961            None,
2962            &instruction,
2963            transaction_accounts.clone(),
2964            vec![buffer_meta.clone(), authority_meta.clone()],
2965            Err(InstructionError::MissingAccount),
2966        );
2967
2968        // Case: wrong present authority
2969        process_instruction(
2970            &loader_id,
2971            None,
2972            &instruction,
2973            vec![
2974                (buffer_address, buffer_account.clone()),
2975                (invalid_authority_address, authority_account),
2976                (new_authority_address, new_authority_account),
2977            ],
2978            vec![
2979                buffer_meta.clone(),
2980                AccountMeta {
2981                    pubkey: invalid_authority_address,
2982                    is_signer: true,
2983                    is_writable: false,
2984                },
2985                new_authority_meta.clone(),
2986            ],
2987            Err(InstructionError::IncorrectAuthority),
2988        );
2989
2990        // Case: present authority did not sign
2991        process_instruction(
2992            &loader_id,
2993            None,
2994            &instruction,
2995            transaction_accounts.clone(),
2996            vec![
2997                buffer_meta.clone(),
2998                AccountMeta {
2999                    pubkey: authority_address,
3000                    is_signer: false,
3001                    is_writable: false,
3002                },
3003                new_authority_meta.clone(),
3004            ],
3005            Err(InstructionError::MissingRequiredSignature),
3006        );
3007
3008        // Case: new authority did not sign
3009        process_instruction(
3010            &loader_id,
3011            None,
3012            &instruction,
3013            transaction_accounts.clone(),
3014            vec![
3015                buffer_meta.clone(),
3016                authority_meta.clone(),
3017                AccountMeta {
3018                    pubkey: new_authority_address,
3019                    is_signer: false,
3020                    is_writable: false,
3021                },
3022            ],
3023            Err(InstructionError::MissingRequiredSignature),
3024        );
3025
3026        // Case: Not a Buffer account
3027        transaction_accounts
3028            .get_mut(0)
3029            .unwrap()
3030            .1
3031            .set_state(&UpgradeableLoaderState::Program {
3032                programdata_address: Pubkey::new_unique(),
3033            })
3034            .unwrap();
3035        process_instruction(
3036            &loader_id,
3037            None,
3038            &instruction,
3039            transaction_accounts.clone(),
3040            vec![
3041                buffer_meta.clone(),
3042                authority_meta.clone(),
3043                new_authority_meta.clone(),
3044            ],
3045            Err(InstructionError::InvalidArgument),
3046        );
3047
3048        // Case: Buffer is immutable
3049        transaction_accounts
3050            .get_mut(0)
3051            .unwrap()
3052            .1
3053            .set_state(&UpgradeableLoaderState::Buffer {
3054                authority_address: None,
3055            })
3056            .unwrap();
3057        process_instruction(
3058            &loader_id,
3059            None,
3060            &instruction,
3061            transaction_accounts.clone(),
3062            vec![buffer_meta, authority_meta, new_authority_meta],
3063            Err(InstructionError::Immutable),
3064        );
3065    }
3066
3067    #[test]
3068    fn test_bpf_loader_upgradeable_close() {
3069        let instruction = bincode::serialize(&UpgradeableLoaderInstruction::Close).unwrap();
3070        let loader_id = bpf_loader_upgradeable::id();
3071        let invalid_authority_address = Pubkey::new_unique();
3072        let authority_address = Pubkey::new_unique();
3073        let authority_account = AccountSharedData::new(1, 0, &Pubkey::new_unique());
3074        let recipient_address = Pubkey::new_unique();
3075        let recipient_account = AccountSharedData::new(1, 0, &Pubkey::new_unique());
3076        let buffer_address = Pubkey::new_unique();
3077        let mut buffer_account =
3078            AccountSharedData::new(1, UpgradeableLoaderState::size_of_buffer(128), &loader_id);
3079        buffer_account
3080            .set_state(&UpgradeableLoaderState::Buffer {
3081                authority_address: Some(authority_address),
3082            })
3083            .unwrap();
3084        let uninitialized_address = Pubkey::new_unique();
3085        let mut uninitialized_account = AccountSharedData::new(
3086            1,
3087            UpgradeableLoaderState::size_of_programdata(0),
3088            &loader_id,
3089        );
3090        uninitialized_account
3091            .set_state(&UpgradeableLoaderState::Uninitialized)
3092            .unwrap();
3093        let programdata_address = Pubkey::new_unique();
3094        let mut programdata_account = AccountSharedData::new(
3095            1,
3096            UpgradeableLoaderState::size_of_programdata(128),
3097            &loader_id,
3098        );
3099        programdata_account
3100            .set_state(&UpgradeableLoaderState::ProgramData {
3101                slot: 0,
3102                upgrade_authority_address: Some(authority_address),
3103            })
3104            .unwrap();
3105        let program_address = Pubkey::new_unique();
3106        let mut program_account =
3107            AccountSharedData::new(1, UpgradeableLoaderState::size_of_program(), &loader_id);
3108        program_account.set_executable(true);
3109        program_account
3110            .set_state(&UpgradeableLoaderState::Program {
3111                programdata_address,
3112            })
3113            .unwrap();
3114        let clock_account = create_account_for_test(&Clock {
3115            slot: 1,
3116            ..Clock::default()
3117        });
3118        let transaction_accounts = vec![
3119            (buffer_address, buffer_account.clone()),
3120            (recipient_address, recipient_account.clone()),
3121            (authority_address, authority_account.clone()),
3122        ];
3123        let buffer_meta = AccountMeta {
3124            pubkey: buffer_address,
3125            is_signer: false,
3126            is_writable: true,
3127        };
3128        let recipient_meta = AccountMeta {
3129            pubkey: recipient_address,
3130            is_signer: false,
3131            is_writable: true,
3132        };
3133        let authority_meta = AccountMeta {
3134            pubkey: authority_address,
3135            is_signer: true,
3136            is_writable: false,
3137        };
3138
3139        // Case: close a buffer account
3140        let accounts = process_instruction(
3141            &loader_id,
3142            None,
3143            &instruction,
3144            transaction_accounts,
3145            vec![
3146                buffer_meta.clone(),
3147                recipient_meta.clone(),
3148                authority_meta.clone(),
3149            ],
3150            Ok(()),
3151        );
3152        assert_eq!(0, accounts.first().unwrap().lamports());
3153        assert_eq!(2, accounts.get(1).unwrap().lamports());
3154        let state: UpgradeableLoaderState = accounts.first().unwrap().state().unwrap();
3155        assert_eq!(state, UpgradeableLoaderState::Uninitialized);
3156        assert_eq!(
3157            UpgradeableLoaderState::size_of_uninitialized(),
3158            accounts.first().unwrap().data().len()
3159        );
3160
3161        // Case: close with wrong authority
3162        process_instruction(
3163            &loader_id,
3164            None,
3165            &instruction,
3166            vec![
3167                (buffer_address, buffer_account.clone()),
3168                (recipient_address, recipient_account.clone()),
3169                (invalid_authority_address, authority_account.clone()),
3170            ],
3171            vec![
3172                buffer_meta,
3173                recipient_meta.clone(),
3174                AccountMeta {
3175                    pubkey: invalid_authority_address,
3176                    is_signer: true,
3177                    is_writable: false,
3178                },
3179            ],
3180            Err(InstructionError::IncorrectAuthority),
3181        );
3182
3183        // Case: close an uninitialized account
3184        let accounts = process_instruction(
3185            &loader_id,
3186            None,
3187            &instruction,
3188            vec![
3189                (uninitialized_address, uninitialized_account.clone()),
3190                (recipient_address, recipient_account.clone()),
3191                (invalid_authority_address, authority_account.clone()),
3192            ],
3193            vec![
3194                AccountMeta {
3195                    pubkey: uninitialized_address,
3196                    is_signer: false,
3197                    is_writable: true,
3198                },
3199                recipient_meta.clone(),
3200                authority_meta.clone(),
3201            ],
3202            Ok(()),
3203        );
3204        assert_eq!(0, accounts.first().unwrap().lamports());
3205        assert_eq!(2, accounts.get(1).unwrap().lamports());
3206        let state: UpgradeableLoaderState = accounts.first().unwrap().state().unwrap();
3207        assert_eq!(state, UpgradeableLoaderState::Uninitialized);
3208        assert_eq!(
3209            UpgradeableLoaderState::size_of_uninitialized(),
3210            accounts.first().unwrap().data().len()
3211        );
3212
3213        // Case: close a program account
3214        let accounts = process_instruction(
3215            &loader_id,
3216            None,
3217            &instruction,
3218            vec![
3219                (programdata_address, programdata_account.clone()),
3220                (recipient_address, recipient_account.clone()),
3221                (authority_address, authority_account.clone()),
3222                (program_address, program_account.clone()),
3223                (sysvar::clock::id(), clock_account.clone()),
3224            ],
3225            vec![
3226                AccountMeta {
3227                    pubkey: programdata_address,
3228                    is_signer: false,
3229                    is_writable: true,
3230                },
3231                recipient_meta,
3232                authority_meta,
3233                AccountMeta {
3234                    pubkey: program_address,
3235                    is_signer: false,
3236                    is_writable: true,
3237                },
3238            ],
3239            Ok(()),
3240        );
3241        assert_eq!(0, accounts.first().unwrap().lamports());
3242        assert_eq!(2, accounts.get(1).unwrap().lamports());
3243        let state: UpgradeableLoaderState = accounts.first().unwrap().state().unwrap();
3244        assert_eq!(state, UpgradeableLoaderState::Uninitialized);
3245        assert_eq!(
3246            UpgradeableLoaderState::size_of_uninitialized(),
3247            accounts.first().unwrap().data().len()
3248        );
3249
3250        // Try to invoke closed account
3251        programdata_account = accounts.first().unwrap().clone();
3252        program_account = accounts.get(3).unwrap().clone();
3253        process_instruction(
3254            &loader_id,
3255            Some(1),
3256            &[],
3257            vec![
3258                (programdata_address, programdata_account.clone()),
3259                (program_address, program_account.clone()),
3260            ],
3261            Vec::new(),
3262            Err(InstructionError::UnsupportedProgramId),
3263        );
3264
3265        // Case: Reopen should fail
3266        process_instruction(
3267            &loader_id,
3268            None,
3269            &bincode::serialize(&UpgradeableLoaderInstruction::DeployWithMaxDataLen {
3270                max_data_len: 0,
3271            })
3272            .unwrap(),
3273            vec![
3274                (recipient_address, recipient_account),
3275                (programdata_address, programdata_account),
3276                (program_address, program_account),
3277                (buffer_address, buffer_account),
3278                (
3279                    sysvar::rent::id(),
3280                    create_account_for_test(&Rent::default()),
3281                ),
3282                (sysvar::clock::id(), clock_account),
3283                (
3284                    system_program::id(),
3285                    AccountSharedData::new(0, 0, &system_program::id()),
3286                ),
3287                (authority_address, authority_account),
3288            ],
3289            vec![
3290                AccountMeta {
3291                    pubkey: recipient_address,
3292                    is_signer: true,
3293                    is_writable: true,
3294                },
3295                AccountMeta {
3296                    pubkey: programdata_address,
3297                    is_signer: false,
3298                    is_writable: true,
3299                },
3300                AccountMeta {
3301                    pubkey: program_address,
3302                    is_signer: false,
3303                    is_writable: true,
3304                },
3305                AccountMeta {
3306                    pubkey: buffer_address,
3307                    is_signer: false,
3308                    is_writable: false,
3309                },
3310                AccountMeta {
3311                    pubkey: sysvar::rent::id(),
3312                    is_signer: false,
3313                    is_writable: false,
3314                },
3315                AccountMeta {
3316                    pubkey: sysvar::clock::id(),
3317                    is_signer: false,
3318                    is_writable: false,
3319                },
3320                AccountMeta {
3321                    pubkey: system_program::id(),
3322                    is_signer: false,
3323                    is_writable: false,
3324                },
3325                AccountMeta {
3326                    pubkey: authority_address,
3327                    is_signer: false,
3328                    is_writable: false,
3329                },
3330            ],
3331            Err(InstructionError::AccountAlreadyInitialized),
3332        );
3333    }
3334
3335    /// fuzzing utility function
3336    fn fuzz<F>(
3337        bytes: &[u8],
3338        outer_iters: usize,
3339        inner_iters: usize,
3340        offset: Range<usize>,
3341        value: Range<u8>,
3342        work: F,
3343    ) where
3344        F: Fn(&mut [u8]),
3345    {
3346        let mut rng = rand::rng();
3347        for _ in 0..outer_iters {
3348            let mut mangled_bytes = bytes.to_vec();
3349            for _ in 0..inner_iters {
3350                let offset = rng.random_range(offset.start..offset.end);
3351                let value = rng.random_range(value.start..value.end);
3352                *mangled_bytes.get_mut(offset).unwrap() = value;
3353                work(&mut mangled_bytes);
3354            }
3355        }
3356    }
3357
3358    #[test]
3359    #[ignore]
3360    fn test_fuzz() {
3361        let loader_id = bpf_loader::id();
3362        let program_id = Pubkey::new_unique();
3363
3364        // Create program account
3365        let mut file = File::open("test_elfs/out/sbpfv3_return_ok.so").expect("file open failed");
3366        let mut elf = Vec::new();
3367        file.read_to_end(&mut elf).unwrap();
3368
3369        // Mangle the whole file
3370        fuzz(
3371            &elf,
3372            1_000_000_000,
3373            100,
3374            0..elf.len(),
3375            0..255,
3376            |bytes: &mut [u8]| {
3377                let mut program_account = AccountSharedData::new(1, 0, &loader_id);
3378                program_account.set_data(bytes.to_vec());
3379                program_account.set_executable(true);
3380                process_instruction(
3381                    &loader_id,
3382                    None,
3383                    &[],
3384                    vec![(program_id, program_account)],
3385                    Vec::new(),
3386                    Ok(()),
3387                );
3388            },
3389        );
3390    }
3391
3392    #[test]
3393    fn test_calculate_heap_cost() {
3394        let heap_cost = 8_u64;
3395
3396        // heap allocations are in 32K block, `heap_cost` of CU is consumed per additional 32k
3397
3398        // assert less than 32K heap should cost zero unit
3399        assert_eq!(0, calculate_heap_cost(31 * 1024, heap_cost));
3400
3401        // assert exact 32K heap should be cost zero unit
3402        assert_eq!(0, calculate_heap_cost(32 * 1024, heap_cost));
3403
3404        // assert slightly more than 32K heap should cost 1 * heap_cost
3405        assert_eq!(heap_cost, calculate_heap_cost(33 * 1024, heap_cost));
3406
3407        // assert exact 64K heap should cost 1 * heap_cost
3408        assert_eq!(heap_cost, calculate_heap_cost(64 * 1024, heap_cost));
3409    }
3410
3411    fn deploy_test_program(
3412        invoke_context: &mut InvokeContext,
3413        program_id: Pubkey,
3414    ) -> Result<(), InstructionError> {
3415        let mut file = File::open("test_elfs/out/sbpfv3_return_ok.so").expect("file open failed");
3416        let mut elf = Vec::new();
3417        file.read_to_end(&mut elf).unwrap();
3418        deploy_program!(
3419            invoke_context,
3420            &program_id,
3421            &bpf_loader_upgradeable::id(),
3422            elf.len(),
3423            &elf,
3424            2_u64,
3425        );
3426        Ok(())
3427    }
3428
3429    #[test]
3430    fn test_program_usage_count_on_upgrade() {
3431        let transaction_accounts = vec![(
3432            sysvar::epoch_schedule::id(),
3433            create_account_for_test(&EpochSchedule::default()),
3434        )];
3435        with_mock_invoke_context!(invoke_context, transaction_context, transaction_accounts);
3436        let program_id = Pubkey::new_unique();
3437        let env = Arc::new(BuiltinProgram::new_mock());
3438        let program = ProgramCacheEntry {
3439            program: ProgramCacheEntryType::Unloaded(env),
3440            account_owner: ProgramCacheEntryOwner::LoaderV2,
3441            account_size: 0,
3442            deployment_slot: 0,
3443            effective_slot: 0,
3444            tx_usage_counter: Arc::new(AtomicU64::new(100)),
3445            latest_access_slot: AtomicU64::new(0),
3446        };
3447        invoke_context
3448            .program_cache_for_tx_batch
3449            .replenish(program_id, Arc::new(program));
3450        invoke_context
3451            .program_cache_for_tx_batch
3452            .set_slot_for_tests(2);
3453
3454        assert_matches!(
3455            deploy_test_program(&mut invoke_context, program_id,),
3456            Ok(())
3457        );
3458
3459        let updated_program = invoke_context
3460            .program_cache_for_tx_batch
3461            .find(&program_id)
3462            .expect("Didn't find upgraded program in the cache");
3463
3464        assert_eq!(updated_program.deployment_slot, 2);
3465        assert_eq!(
3466            updated_program.tx_usage_counter.load(Ordering::Relaxed),
3467            100
3468        );
3469    }
3470
3471    #[test]
3472    fn test_program_usage_count_on_non_upgrade() {
3473        let transaction_accounts = vec![(
3474            sysvar::epoch_schedule::id(),
3475            create_account_for_test(&EpochSchedule::default()),
3476        )];
3477        with_mock_invoke_context!(invoke_context, transaction_context, transaction_accounts);
3478        let program_id = Pubkey::new_unique();
3479        let env = Arc::new(BuiltinProgram::new_mock());
3480        let program = ProgramCacheEntry {
3481            program: ProgramCacheEntryType::Unloaded(env),
3482            account_owner: ProgramCacheEntryOwner::LoaderV2,
3483            account_size: 0,
3484            deployment_slot: 0,
3485            effective_slot: 0,
3486            tx_usage_counter: Arc::new(AtomicU64::new(100)),
3487            latest_access_slot: AtomicU64::new(0),
3488        };
3489        invoke_context
3490            .program_cache_for_tx_batch
3491            .replenish(program_id, Arc::new(program));
3492        invoke_context
3493            .program_cache_for_tx_batch
3494            .set_slot_for_tests(2);
3495
3496        let program_id2 = Pubkey::new_unique();
3497        assert_matches!(
3498            deploy_test_program(&mut invoke_context, program_id2),
3499            Ok(())
3500        );
3501
3502        let program2 = invoke_context
3503            .program_cache_for_tx_batch
3504            .find(&program_id2)
3505            .expect("Didn't find upgraded program in the cache");
3506
3507        assert_eq!(program2.deployment_slot, 2);
3508        assert_eq!(program2.tx_usage_counter.load(Ordering::Relaxed), 0);
3509    }
3510}