[Solana]: Add Rust FFI functions

* Fix `VersionedMessage::address_table_lookups()` method
trustwallet · Oct 30, 2024 · a659acc · a659acc
1 parent ed9bc35
commit a659acc
Show file tree

Hide file tree

Showing 5 changed files with 109 additions and 9 deletions.
diff --git a/rust/chains/tw_solana/src/modules/insert_instruction.rs b/rust/chains/tw_solana/src/modules/insert_instruction.rs
@@ -92,9 +92,7 @@ pub trait InsertInstruction {
     }
 
     /// Returns ALT (Address Lookup Tables) if supported by the message version.
-    fn address_table_lookups(&self) -> Option<&[MessageAddressTableLookup]> {
-        None
-    }
+    fn address_table_lookups(&self) -> Option<&[MessageAddressTableLookup]>;
 
     fn account_keys_mut(&mut self) -> &mut Vec<SolanaAddress>;
 

diff --git a/rust/chains/tw_solana/src/modules/utils.rs b/rust/chains/tw_solana/src/modules/utils.rs
@@ -99,7 +99,7 @@ impl SolanaTransaction {
             .find_map(try_instruction_as_set_unit_limit))
     }
 
-    pub fn set_compute_unit_price(encoded_tx: &str, price: UnitPrice) -> SigningResult<()> {
+    pub fn set_compute_unit_price(encoded_tx: &str, price: UnitPrice) -> SigningResult<String> {
         let tx_bytes = base64::decode(encoded_tx, STANDARD)?;
         let mut tx: VersionedTransaction =
             bincode::deserialize(&tx_bytes).map_err(|_| SigningErrorType::Error_input_parse)?;
@@ -114,17 +114,17 @@ impl SolanaTransaction {
         // If it presents already, it's enough to update the instruction data only.
         if let Some(pos) = ix_position {
             tx.message.instructions_mut()[pos].data = set_price_ix.data;
-            return Ok(());
+            return tx.to_base64().tw_err(|_| SigningErrorType::Error_internal);
         }
 
         // `ComputeBudgetInstruction::SetComputeUnitPrice` can be pushed to the end of the instructions list.
         tx.message
             .push_simple_instruction(set_price_ix.program_id, set_price_ix.data)?;
 
-        Ok(())
+        tx.to_base64().tw_err(|_| SigningErrorType::Error_internal)
     }
 
-    pub fn set_compute_unit_limit(encoded_tx: &str, limit: UnitLimit) -> SigningResult<()> {
+    pub fn set_compute_unit_limit(encoded_tx: &str, limit: UnitLimit) -> SigningResult<String> {
         let tx_bytes = base64::decode(encoded_tx, STANDARD)?;
         let mut tx: VersionedTransaction =
             bincode::deserialize(&tx_bytes).map_err(|_| SigningErrorType::Error_input_parse)?;
@@ -139,7 +139,7 @@ impl SolanaTransaction {
         // If it presents already, it's enough to update the instruction data only.
         if let Some(pos) = ix_position {
             tx.message.instructions_mut()[pos].data = set_limit_ix.data;
-            return Ok(());
+            return tx.to_base64().tw_err(|_| SigningErrorType::Error_internal);
         }
 
         // `ComputeBudgetInstruction::SetComputeUnitLimit` should be at the beginning of the instructions list.
@@ -156,7 +156,7 @@ impl SolanaTransaction {
             set_limit_ix.data,
         )?;
 
-        Ok(())
+        tx.to_base64().tw_err(|_| SigningErrorType::Error_internal)
     }
 }
 

diff --git a/rust/chains/tw_solana/src/transaction/legacy.rs b/rust/chains/tw_solana/src/transaction/legacy.rs
@@ -4,6 +4,7 @@
 
 use crate::address::SolanaAddress;
 use crate::modules::insert_instruction::InsertInstruction;
+use crate::transaction::v0::MessageAddressTableLookup;
 use crate::transaction::{short_vec, CompiledInstruction, MessageHeader, Signature};
 use serde::{Deserialize, Serialize};
 use tw_hash::{as_byte_sequence, H256};
@@ -30,6 +31,10 @@ pub struct Message {
 }
 
 impl InsertInstruction for Message {
+    fn address_table_lookups(&self) -> Option<&[MessageAddressTableLookup]> {
+        None
+    }
+
     fn account_keys_mut(&mut self) -> &mut Vec<SolanaAddress> {
         &mut self.account_keys
     }

diff --git a/rust/chains/tw_solana/src/transaction/versioned.rs b/rust/chains/tw_solana/src/transaction/versioned.rs
@@ -7,6 +7,7 @@
 use crate::address::SolanaAddress;
 use crate::blockhash::Blockhash;
 use crate::modules::insert_instruction::InsertInstruction;
+use crate::transaction::v0::MessageAddressTableLookup;
 use crate::transaction::{legacy, short_vec, v0, CompiledInstruction, MessageHeader, Signature};
 use serde::de::{SeqAccess, Unexpected, Visitor};
 use serde::ser::SerializeTuple;
@@ -42,6 +43,11 @@ impl VersionedTransaction {
         let tx_bytes = base64::decode(s, STANDARD)?;
         bincode::deserialize(&tx_bytes).map_err(|_| EncodingError::InvalidInput)
     }
+
+    pub fn to_base64(&self) -> EncodingResult<String> {
+        let tx_bytes = bincode::serialize(self).map_err(|_| EncodingError::InvalidInput)?;
+        Ok(base64::encode(&tx_bytes, STANDARD))
+    }
 }
 
 /// Either a legacy message or a v0 message.
@@ -142,6 +148,13 @@ impl Serialize for VersionedMessage {
 }
 
 impl InsertInstruction for VersionedMessage {
+    fn address_table_lookups(&self) -> Option<&[MessageAddressTableLookup]> {
+        match self {
+            VersionedMessage::Legacy(legacy) => legacy.address_table_lookups(),
+            VersionedMessage::V0(v0) => v0.address_table_lookups(),
+        }
+    }
+
     fn account_keys_mut(&mut self) -> &mut Vec<SolanaAddress> {
         match self {
             VersionedMessage::Legacy(legacy) => legacy.account_keys_mut(),

diff --git a/rust/wallet_core_rs/src/ffi/solana/transaction.rs b/rust/wallet_core_rs/src/ffi/solana/transaction.rs
@@ -49,3 +49,87 @@ pub unsafe extern "C" fn tw_solana_transaction_update_blockhash_and_sign(
 
     TWData::from(output_proto).into_ptr()
 }
+
+/// Try to find a `ComputeBudgetInstruction::SetComputeUnitPrice` instruction in the given transaction,
+/// and returns the specified Unit Price.
+///
+/// \param encoded_tx base64 encoded Solana transaction.
+/// \return nullable Unit Price as a decimal string. Null if no instruction found.
+#[no_mangle]
+pub unsafe extern "C" fn tw_solana_transaction_get_compute_unit_price(
+    encoded_tx: *const TWString,
+) -> *mut TWString {
+    let encoded_tx = try_or_else!(TWString::from_ptr_as_ref(encoded_tx), std::ptr::null_mut);
+    let encoded_tx = try_or_else!(encoded_tx.as_str(), std::ptr::null_mut);
+
+    match SolanaTransaction::get_compute_unit_price(encoded_tx) {
+        Ok(Some(price)) => TWString::from(price.to_string()).into_ptr(),
+        _ => std::ptr::null_mut(),
+    }
+}
+
+/// Try to find a `ComputeBudgetInstruction::SetComputeUnitLimit` instruction in the given transaction,
+/// and returns the specified Unit Limit.
+///
+/// \param encoded_tx base64 encoded Solana transaction.
+/// \return nullable Unit Limit as a decimal string. Null if no instruction found.
+#[no_mangle]
+pub unsafe extern "C" fn tw_solana_transaction_get_compute_unit_limit(
+    encoded_tx: *const TWString,
+) -> *mut TWString {
+    let encoded_tx = try_or_else!(TWString::from_ptr_as_ref(encoded_tx), std::ptr::null_mut);
+    let encoded_tx = try_or_else!(encoded_tx.as_str(), std::ptr::null_mut);
+
+    match SolanaTransaction::get_compute_unit_limit(encoded_tx) {
+        Ok(Some(limit)) => TWString::from(limit.to_string()).into_ptr(),
+        _ => std::ptr::null_mut(),
+    }
+}
+
+/// Adds or updates a `ComputeBudgetInstruction::SetComputeUnitPrice` instruction of the given transaction,
+/// and returns the updated transaction.
+///
+/// \param encoded_tx base64 encoded Solana transaction.
+/// \price Unit Price as a decimal string.
+/// \return base64 encoded Solana transaction.
+#[no_mangle]
+pub unsafe extern "C" fn tw_solana_transaction_set_compute_unit_price(
+    encoded_tx: *const TWString,
+    price: *const TWString,
+) -> *mut TWString {
+    let encoded_tx = try_or_else!(TWString::from_ptr_as_ref(encoded_tx), std::ptr::null_mut);
+    let encoded_tx = try_or_else!(encoded_tx.as_str(), std::ptr::null_mut);
+
+    let price = try_or_else!(TWString::from_ptr_as_ref(price), std::ptr::null_mut);
+    let price = try_or_else!(price.as_str(), std::ptr::null_mut);
+    let price = try_or_else!(price.parse(), std::ptr::null_mut);
+
+    match SolanaTransaction::set_compute_unit_price(encoded_tx, price) {
+        Ok(updated_tx) => TWString::from(updated_tx).into_ptr(),
+        _ => std::ptr::null_mut(),
+    }
+}
+
+/// Adds or updates a `ComputeBudgetInstruction::SetComputeUnitLimit` instruction of the given transaction,
+/// and returns the updated transaction.
+///
+/// \param encoded_tx base64 encoded Solana transaction.
+/// \limit Unit Limit as a decimal string.
+/// \return base64 encoded Solana transaction.
+#[no_mangle]
+pub unsafe extern "C" fn tw_solana_transaction_set_compute_unit_limit(
+    encoded_tx: *const TWString,
+    limit: *const TWString,
+) -> *mut TWString {
+    let encoded_tx = try_or_else!(TWString::from_ptr_as_ref(encoded_tx), std::ptr::null_mut);
+    let encoded_tx = try_or_else!(encoded_tx.as_str(), std::ptr::null_mut);
+
+    let limit = try_or_else!(TWString::from_ptr_as_ref(limit), std::ptr::null_mut);
+    let limit = try_or_else!(limit.as_str(), std::ptr::null_mut);
+    let limit = try_or_else!(limit.parse(), std::ptr::null_mut);
+
+    match SolanaTransaction::set_compute_unit_limit(encoded_tx, limit) {
+        Ok(updated_tx) => TWString::from(updated_tx).into_ptr(),
+        _ => std::ptr::null_mut(),
+    }
+}