Reverify programs that are extended using ExtendProgram (#31886)

programs/bpf-loader-tests/tests/extend_program_ix.rs

@@ -10,7 +10,7 @@ use {
         signature::{Keypair, Signer},
         system_instruction::{self, SystemError, MAX_PERMITTED_DATA_LENGTH},
         system_program,
-        transaction::Transaction,
+        transaction::{Transaction, TransactionError},
     },
 };
 
@@ -19,10 +19,11 @@ mod common;
 #[tokio::test]
 async fn test_extend_program() {
     let mut context = setup_test_context().await;
+    let program_file = find_file("noop.so").expect("Failed to find the file");
+    let data = read_file(program_file);
 
     let program_address = Pubkey::new_unique();
     let (programdata_address, _) = Pubkey::find_program_address(&[program_address.as_ref()], &id());
-    let program_data_len = 100;
     add_upgradeable_loader_account(
         &mut context,
         &program_address,
@@ -33,6 +34,8 @@ async fn test_extend_program() {
         |_| {},
     )
     .await;
+    let programdata_data_offset = UpgradeableLoaderState::size_of_programdata_metadata();
+    let program_data_len = data.len() + programdata_data_offset;
     add_upgradeable_loader_account(
         &mut context,
         &programdata_address,
@@ -41,9 +44,68 @@ async fn test_extend_program() {
             upgrade_authority_address: Some(Pubkey::new_unique()),
         },
         program_data_len,
+        |account| account.data_as_mut_slice()[programdata_data_offset..].copy_from_slice(&data),
+    )
+    .await;
+
+    let client = &mut context.banks_client;
+    let payer = &context.payer;
+    let recent_blockhash = context.last_blockhash;
+    const ADDITIONAL_BYTES: u32 = 42;
+    let transaction = Transaction::new_signed_with_payer(
+        &[extend_program(
+            &program_address,
+            Some(&payer.pubkey()),
+            ADDITIONAL_BYTES,
+        )],
+        Some(&payer.pubkey()),
+        &[payer],
+        recent_blockhash,
+    );
+
+    assert_matches!(client.process_transaction(transaction).await, Ok(()));
+    let updated_program_data_account = client
+        .get_account(programdata_address)
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(
+        updated_program_data_account.data().len(),
+        program_data_len + ADDITIONAL_BYTES as usize
+    );
+}
+
+#[tokio::test]
+async fn test_failed_extend_twice_in_same_slot() {
+    let mut context = setup_test_context().await;
+    let program_file = find_file("noop.so").expect("Failed to find the file");
+    let data = read_file(program_file);
+
+    let program_address = Pubkey::new_unique();
+    let (programdata_address, _) = Pubkey::find_program_address(&[program_address.as_ref()], &id());
+    add_upgradeable_loader_account(
+        &mut context,
+        &program_address,
+        &UpgradeableLoaderState::Program {
+            programdata_address,
+        },
+        UpgradeableLoaderState::size_of_program(),
         |_| {},
     )
     .await;
+    let programdata_data_offset = UpgradeableLoaderState::size_of_programdata_metadata();
+    let program_data_len = data.len() + programdata_data_offset;
+    add_upgradeable_loader_account(
+        &mut context,
+        &programdata_address,
+        &UpgradeableLoaderState::ProgramData {
+            slot: 0,
+            upgrade_authority_address: Some(Pubkey::new_unique()),
+        },
+        program_data_len,
+        |account| account.data_as_mut_slice()[programdata_data_offset..].copy_from_slice(&data),
+    )
+    .await;
 
     let client = &mut context.banks_client;
     let payer = &context.payer;
@@ -70,6 +132,31 @@ async fn test_extend_program() {
         updated_program_data_account.data().len(),
         program_data_len + ADDITIONAL_BYTES as usize
     );
+
+    let recent_blockhash = client
+        .get_new_latest_blockhash(&recent_blockhash)
+        .await
+        .unwrap();
+    // Extending the program in the same slot should fail
+    let transaction = Transaction::new_signed_with_payer(
+        &[extend_program(
+            &program_address,
+            Some(&payer.pubkey()),
+            ADDITIONAL_BYTES,
+        )],
+        Some(&payer.pubkey()),
+        &[payer],
+        recent_blockhash,
+    );
+
+    assert_matches!(
+        client
+            .process_transaction(transaction)
+            .await
+            .unwrap_err()
+            .unwrap(),
+        TransactionError::InstructionError(0, InstructionError::InvalidArgument)
+    );
 }
 
 #[tokio::test]
@@ -293,6 +380,9 @@ async fn test_extend_program_without_payer() {
     let mut context = setup_test_context().await;
     let rent = context.banks_client.get_rent().await.unwrap();
 
+    let program_file = find_file("noop.so").expect("Failed to find the file");
+    let data = read_file(program_file);
+
     let program_address = Pubkey::new_unique();
     let (programdata_address, _) = Pubkey::find_program_address(&[program_address.as_ref()], &id());
     add_upgradeable_loader_account(
@@ -305,7 +395,8 @@ async fn test_extend_program_without_payer() {
         |_| {},
     )
     .await;
-    let program_data_len = 100;
+    let programdata_data_offset = UpgradeableLoaderState::size_of_programdata_metadata();
+    let program_data_len = data.len() + programdata_data_offset;
     add_upgradeable_loader_account(
         &mut context,
         &programdata_address,
@@ -314,7 +405,7 @@ async fn test_extend_program_without_payer() {
             upgrade_authority_address: Some(Pubkey::new_unique()),
         },
         program_data_len,
-        |_| {},
+        |account| account.data_as_mut_slice()[programdata_data_offset..].copy_from_slice(&data),
     )
     .await;
 

programs/bpf_loader/src/lib.rs

@@ -1325,6 +1325,7 @@ fn process_loader_upgradeable_instruction(
                 ic_logger_msg!(log_collector, "Program account not owned by loader");
                 return Err(InstructionError::InvalidAccountOwner);
             }
+            let program_key = *program_account.get_key();
             match program_account.get_state()? {
                 UpgradeableLoaderState::Program {
                     programdata_address,
@@ -1356,11 +1357,21 @@ fn process_loader_upgradeable_instruction(
                 return Err(InstructionError::InvalidRealloc);
             }
 
-            if let UpgradeableLoaderState::ProgramData {
-                slot: _,
+            let clock_slot = invoke_context
+                .get_sysvar_cache()
+                .get_clock()
+                .map(|clock| clock.slot)?;
+
+            let upgrade_authority_address = if let UpgradeableLoaderState::ProgramData {
+                slot,
                 upgrade_authority_address,
             } = programdata_account.get_state()?
             {
+                if clock_slot == slot {
+                    ic_logger_msg!(log_collector, "Program was extended in this block already");
+                    return Err(InstructionError::InvalidArgument);
+                }
+
                 if upgrade_authority_address.is_none() {
                     ic_logger_msg!(
                         log_collector,
@@ -1368,10 +1379,11 @@ fn process_loader_upgradeable_instruction(
                     );
                     return Err(InstructionError::Immutable);
                 }
+                upgrade_authority_address
             } else {
                 ic_logger_msg!(log_collector, "ProgramData state is invalid");
                 return Err(InstructionError::InvalidAccountData);
-            }
+            };
 
             let required_payment = {
                 let balance = programdata_account.get_lamports();
@@ -1383,6 +1395,8 @@ fn process_loader_upgradeable_instruction(
             // Borrowed accounts need to be dropped before native_invoke
             drop(programdata_account);
 
+            // Dereference the program ID to prevent overlapping mutable/immutable borrow of invoke context
+            let program_id = *program_id;
             if required_payment > 0 {
                 let payer_key = *transaction_context.get_key_of_account_at_index(
                     instruction_context.get_index_of_instruction_account_in_transaction(
@@ -1403,6 +1417,30 @@ fn process_loader_upgradeable_instruction(
                 .try_borrow_instruction_account(transaction_context, PROGRAM_DATA_ACCOUNT_INDEX)?;
             programdata_account.set_data_length(new_len)?;
 
+            let programdata_data_offset = UpgradeableLoaderState::size_of_programdata_metadata();
+
+            deploy_program!(
+                invoke_context,
+                program_key,
+                &program_id,
+                UpgradeableLoaderState::size_of_program().saturating_add(new_len),
+                clock_slot,
+                {
+                    drop(programdata_account);
+                },
+                programdata_account
+                    .get_data()
+                    .get(programdata_data_offset..)
+                    .ok_or(InstructionError::AccountDataTooSmall)?,
+            );
+
+            let mut programdata_account = instruction_context
+                .try_borrow_instruction_account(transaction_context, PROGRAM_DATA_ACCOUNT_INDEX)?;
+            programdata_account.set_state(&UpgradeableLoaderState::ProgramData {
+                slot: clock_slot,
+                upgrade_authority_address,
+            })?;
+
             ic_logger_msg!(
                 log_collector,
                 "Extended ProgramData account by {} bytes",