Adjusts CLI manual extend and auto extend.

Author: Lichtso

cli/src/program.rs

@@ -169,7 +169,7 @@ pub enum ProgramCliCommand {
         use_lamports_unit: bool,
         bypass_warning: bool,
     },
-    ExtendProgram {
+    ExtendProgramChecked {
         program_pubkey: Pubkey,
         additional_bytes: u32,
     },
@@ -1018,7 +1018,7 @@ pub fn parse_program_subcommand(
             )?;
 
             CliCommandInfo {
-                command: CliCommand::Program(ProgramCliCommand::ExtendProgram {
+                command: CliCommand::Program(ProgramCliCommand::ExtendProgramChecked {
                     program_pubkey,
                     additional_bytes,
                 }),
@@ -1231,7 +1231,7 @@ pub fn process_program_subcommand(
             *use_lamports_unit,
             *bypass_warning,
         ),
-        ProgramCliCommand::ExtendProgram {
+        ProgramCliCommand::ExtendProgramChecked {
             program_pubkey,
             additional_bytes,
         } => process_extend_program(&rpc_client, config, *program_pubkey, *additional_bytes),
@@ -2411,21 +2411,28 @@ fn process_extend_program(
         _ => Err(format!("Program {program_pubkey} is closed")),
     }?;
 
-    match upgrade_authority_address {
-        None => Err(format!("Program {program_pubkey} is not upgradeable")),
-        _ => Ok(()),
-    }?;
+    let upgrade_authority_address = upgrade_authority_address
+        .ok_or_else(|| format!("Program {program_pubkey} is not upgradeable"))?;
 
     let blockhash = rpc_client.get_latest_blockhash()?;
-
-    let mut tx = Transaction::new_unsigned(Message::new(
-        &[loader_v3_instruction::extend_program(
-            &program_pubkey,
-            Some(&payer_pubkey),
-            additional_bytes,
-        )],
-        Some(&payer_pubkey),
-    ));
+    let feature_set = fetch_feature_set(rpc_client)?;
+
+    let instruction =
+        if feature_set.is_active(&agave_feature_set::enable_extend_program_checked::id()) {
+            loader_v3_instruction::extend_program_checked(
+                &program_pubkey,
+                &upgrade_authority_address,
+                Some(&payer_pubkey),
+                additional_bytes,
+            )
+        } else {
+            loader_v3_instruction::extend_program(
+                &program_pubkey,
+                Some(&payer_pubkey),
+                additional_bytes,
+            )
+        };
+    let mut tx = Transaction::new_unsigned(Message::new(&[instruction], Some(&payer_pubkey)));
 
     tx.try_sign(&[config.signers[0]], blockhash)?;
     let result = rpc_client.send_and_confirm_transaction_with_spinner_and_config(
@@ -2972,6 +2979,17 @@ fn extend_program_data_if_needed(
         return Ok(());
     };
 
+    let upgrade_authority_address = match program_data_account.state() {
+        Ok(UpgradeableLoaderState::ProgramData {
+            slot: _,
+            upgrade_authority_address,
+        }) => Ok(upgrade_authority_address),
+        _ => Err(format!("Program {program_id} is closed")),
+    }?;
+
+    let upgrade_authority_address = upgrade_authority_address
+        .ok_or_else(|| format!("Program {program_id} is not upgradeable"))?;
+
     let required_len = UpgradeableLoaderState::size_of_programdata(program_len);
     let max_permitted_data_length = usize::try_from(MAX_PERMITTED_DATA_LENGTH).unwrap();
     if required_len > max_permitted_data_length {
@@ -2993,11 +3011,20 @@ fn extend_program_data_if_needed(
 
     let additional_bytes =
         u32::try_from(additional_bytes).expect("`u32` is big enough to hold an account size");
-    initial_instructions.push(loader_v3_instruction::extend_program(
-        program_id,
-        Some(fee_payer),
-        additional_bytes,
-    ));
+
+    let feature_set = fetch_feature_set(rpc_client)?;
+    let instruction =
+        if feature_set.is_active(&agave_feature_set::enable_extend_program_checked::id()) {
+            loader_v3_instruction::extend_program_checked(
+                program_id,
+                &upgrade_authority_address,
+                Some(fee_payer),
+                additional_bytes,
+            )
+        } else {
+            loader_v3_instruction::extend_program(program_id, Some(fee_payer), additional_bytes)
+        };
+    initial_instructions.push(instruction);
 
     Ok(())
 }
@@ -3096,7 +3123,12 @@ fn send_deploy_messages(
             // account to sign the transaction. One (transfer) only requires the fee-payer signature.
             // This check is to ensure signing does not fail on a KeypairPubkeyMismatch error from an
             // extraneous signature.
-            if message.header.num_required_signatures == 2 {
+            if message.header.num_required_signatures == 3 {
+                initial_transaction.try_sign(
+                    &[fee_payer_signer, initial_signer, write_signer.unwrap()],
+                    blockhash,
+                )?;
+            } else if message.header.num_required_signatures == 2 {
                 initial_transaction.try_sign(&[fee_payer_signer, initial_signer], blockhash)?;
             } else {
                 initial_transaction.try_sign(&[fee_payer_signer], blockhash)?;
@@ -4406,7 +4438,7 @@ mod tests {
         assert_eq!(
             parse_command(&test_command, &default_signer, &mut None).unwrap(),
             CliCommandInfo {
-                command: CliCommand::Program(ProgramCliCommand::ExtendProgram {
+                command: CliCommand::Program(ProgramCliCommand::ExtendProgramChecked {
                     program_pubkey,
                     additional_bytes
                 }),

cli/tests/program.rs

@@ -1444,7 +1444,7 @@ fn test_cli_program_extend_program() {
     let new_max_len = new_program_data.len();
     let additional_bytes = (new_max_len - max_len) as u32;
     config.signers = vec![&keypair];
-    config.command = CliCommand::Program(ProgramCliCommand::ExtendProgram {
+    config.command = CliCommand::Program(ProgramCliCommand::ExtendProgramChecked {
         program_pubkey: program_keypair.pubkey(),
         additional_bytes: additional_bytes - 1,
     });
@@ -1492,7 +1492,7 @@ fn test_cli_program_extend_program() {
 
     // Extend 1 last byte
     config.signers = vec![&keypair];
-    config.command = CliCommand::Program(ProgramCliCommand::ExtendProgram {
+    config.command = CliCommand::Program(ProgramCliCommand::ExtendProgramChecked {
         program_pubkey: program_keypair.pubkey(),
         additional_bytes: 1,
     });