Support adding instruction in Solana transaction (trustwallet#4371)

Co-authored-by: Sergei Boiko <[email protected]>

Author: 10gic

rust/chains/tw_solana/src/instruction.rs

@@ -5,15 +5,19 @@
 use crate::address::SolanaAddress;
 use borsh::BorshSerialize;
 use serde::{Deserialize, Serialize};
+use tw_encoding::base58::as_base58_bitcoin;
 use tw_memory::Data;
 
+#[derive(Debug, Default, PartialEq, Eq, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
 pub struct Instruction {
     /// Pubkey of the program that executes this instruction.
     pub program_id: SolanaAddress,
     /// Metadata describing accounts that should be passed to the program.
     pub accounts: Vec<AccountMeta>,
     /// Opaque data passed to the program for its own interpretation.
-    pub data: Data,
+    #[serde(with = "as_base58_bitcoin")]
+    pub data: Data, // Rpc getTransaction uses base58 encoding, we use the same encoding for consistency
 }
 
 impl Instruction {
@@ -72,6 +76,7 @@ impl Instruction {
 }
 
 #[derive(Debug, Default, PartialEq, Eq, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
 pub struct AccountMeta {
     /// An account's public key.
     pub pubkey: SolanaAddress,

rust/chains/tw_solana/src/modules/insert_instruction.rs

@@ -3,6 +3,7 @@
 // Copyright © 2017 Trust Wallet.
 
 use crate::address::SolanaAddress;
+use crate::instruction::AccountMeta;
 use crate::modules::compiled_keys::try_into_u8;
 use crate::transaction::v0::MessageAddressTableLookup;
 use crate::transaction::{CompiledInstruction, MessageHeader};
@@ -11,6 +12,59 @@ use tw_coin_entry::error::prelude::*;
 use tw_memory::Data;
 
 pub trait InsertInstruction {
+    /// Pushes an instruction
+    fn push_instruction(
+        &mut self,
+        program_id: SolanaAddress,
+        accounts: Vec<AccountMeta>,
+        data: Data,
+    ) -> SigningResult<()> {
+        let insert_at = self.instructions_mut().len();
+        self.insert_instruction(insert_at, program_id, accounts, data)
+    }
+
+    /// Inserts an instruction at the given `insert_at` index.
+    fn insert_instruction(
+        &mut self,
+        insert_at: usize,
+        program_id: SolanaAddress,
+        accounts: Vec<AccountMeta>,
+        data: Data,
+    ) -> SigningResult<()> {
+        if insert_at > self.instructions_mut().len() {
+            return SigningError::err(SigningErrorType::Error_internal)
+            .context(format!("Unable to add '{program_id}' instruction at the '{insert_at}' index. Number of existing instructions: {}", self.instructions_mut().len()));
+        }
+
+        // Step 1 - add the `account` in the accounts list.
+        let accounts: Vec<u8> = accounts
+            .iter()
+            .map(|account_meta| self.push_account(account_meta))
+            .collect::<Result<Vec<u8>, _>>()?;
+
+        // Step 2 - find or add the `program_id` in the accounts list.
+        let program_id_index = match self
+            .account_keys_mut()
+            .iter()
+            .position(|acc| *acc == program_id)
+        {
+            Some(pos) => try_into_u8(pos)?,
+            None => self.push_readonly_unsigned_account(program_id)?,
+        };
+
+        // Step 3 - Create a `CompiledInstruction` based on the `program_id` index and instruction `accounts` and `data`.
+        let new_compiled_ix = CompiledInstruction {
+            program_id_index,
+            accounts,
+            data,
+        };
+
+        // Step 4 - Insert the created instruction at the given `insert_at` index.
+        self.instructions_mut().insert(insert_at, new_compiled_ix);
+
+        Ok(())
+    }
+
     /// Pushes a simple instruction that doesn't have accounts.
     fn push_simple_instruction(
         &mut self,
@@ -56,6 +110,107 @@ pub trait InsertInstruction {
         Ok(())
     }
 
+    /// Pushes an account to the message.
+    /// If the account already exists, it must match the `is_signer` and `is_writable` attributes
+    /// Returns the index of the account in the account list.
+    fn push_account(&mut self, account: &AccountMeta) -> SigningResult<u8> {
+        // The layout of the account keys is as follows:
+        // +-------------------------------------+
+        // | Writable and required signature     |                                     \
+        // +-------------------------------------+                                      |->  num_required_signatures
+        // | Readonly and required signature     | --> num_readonly_signed_accounts    /
+        // +-------------------------------------+
+        // | Writable and not required signature |
+        // +-------------------------------------+
+        // | Readonly and not required signature | --> num_readonly_unsigned_accounts
+        // +-------------------------------------+
+
+        // Check if the account already exists in `account_keys`,
+        // if it does, validate `is_signer` and `is_writable` match
+        if let Some(existing_index) = self
+            .account_keys_mut()
+            .iter()
+            .position(|key| *key == account.pubkey)
+        {
+            let is_signer =
+                existing_index < self.message_header_mut().num_required_signatures as usize;
+
+            let is_writable = if is_signer {
+                existing_index
+                    < (self.message_header_mut().num_required_signatures
+                        - self.message_header_mut().num_readonly_signed_accounts)
+                        as usize
+            } else {
+                existing_index
+                    < (self.account_keys_mut().len()
+                        - self.message_header_mut().num_readonly_unsigned_accounts as usize)
+            };
+
+            if account.is_signer != is_signer {
+                return SigningError::err(SigningErrorType::Error_internal).context(
+                    "Account already exists but the `is_signer` attribute does not match",
+                );
+            }
+            if account.is_writable != is_writable {
+                return SigningError::err(SigningErrorType::Error_internal).context(
+                    "Account already exists but the `is_writable` attribute does not match",
+                );
+            }
+            // Return the existing index if validation passes
+            return try_into_u8(existing_index);
+        }
+
+        // Determine the insertion position based on is_signer and is_writable
+        let insert_at = match (account.is_signer, account.is_writable) {
+            (true, true) => {
+                self.message_header_mut().num_required_signatures += 1;
+                // The account is added at the end of the writable and signer accounts
+                (self.message_header_mut().num_required_signatures
+                    - self.message_header_mut().num_readonly_signed_accounts)
+                    as usize
+                    - 1
+            },
+            (true, false) => {
+                self.message_header_mut().num_required_signatures += 1;
+                self.message_header_mut().num_readonly_signed_accounts += 1;
+                // The account is added at the end of the read-only and signer accounts
+                self.message_header_mut().num_required_signatures as usize - 1
+            },
+            (false, true) => {
+                // The account is added at the end of the writable and non-signer accounts
+                self.account_keys_mut().len()
+                    - self.message_header_mut().num_readonly_unsigned_accounts as usize
+            },
+            (false, false) => {
+                self.message_header_mut().num_readonly_unsigned_accounts += 1;
+                // The account is added at the end of the list
+                self.account_keys_mut().len()
+            },
+        };
+
+        // Insert the account at the determined position
+        self.account_keys_mut().insert(insert_at, account.pubkey);
+
+        let account_added_at = try_into_u8(insert_at)?;
+
+        // Update program ID and account indexes if the new account was added before its position
+        let instructions = self.instructions_mut();
+        instructions.iter_mut().for_each(|ix| {
+            // Update program ID index
+            if ix.program_id_index >= account_added_at {
+                ix.program_id_index += 1;
+            }
+
+            // Update account indexes
+            ix.accounts
+                .iter_mut()
+                .filter(|ix_account_id| **ix_account_id >= account_added_at)
+                .for_each(|ix_account_id| *ix_account_id += 1);
+        });
+
+        Ok(account_added_at)
+    }
+
     fn push_readonly_unsigned_account(&mut self, account: SolanaAddress) -> SigningResult<u8> {
         debug_assert!(
             !self.account_keys_mut().contains(&account),

rust/chains/tw_solana/src/modules/utils.rs

@@ -4,6 +4,7 @@
 
 use crate::address::SolanaAddress;
 use crate::defined_addresses::{COMPUTE_BUDGET_ADDRESS, SYSTEM_PROGRAM_ID_ADDRESS};
+use crate::instruction::Instruction;
 use crate::modules::insert_instruction::InsertInstruction;
 use crate::modules::instruction_builder::compute_budget_instruction::{
     ComputeBudgetInstruction, ComputeBudgetInstructionBuilder, UnitLimit, UnitPrice,
@@ -173,6 +174,27 @@ impl SolanaTransaction {
             .to_base64()
             .tw_err(SigningErrorType::Error_internal)
     }
+
+    pub fn add_instruction(encoded_tx: &str, instruction: &str) -> SigningResult<String> {
+        let tx_bytes = base64::decode(encoded_tx, STANDARD)?;
+        let mut tx: VersionedTransaction =
+            bincode::deserialize(&tx_bytes).map_err(|_| SigningErrorType::Error_input_parse)?;
+
+        let instruction: Instruction =
+            serde_json::from_str(instruction).map_err(|_| SigningErrorType::Error_input_parse)?;
+
+        tx.message.push_instruction(
+            instruction.program_id,
+            instruction.accounts,
+            instruction.data,
+        )?;
+
+        // Set the correct number of zero signatures
+        let unsigned_tx = VersionedTransaction::unsigned(tx.message);
+        unsigned_tx
+            .to_base64()
+            .tw_err(SigningErrorType::Error_internal)
+    }
 }
 
 fn try_instruction_as_compute_budget(