Avoid unwrap in get_token_account for solana (#1138)

Author: 0xh3rman

crates/gem_solana/src/address.rs

@@ -11,11 +11,11 @@ mod tests {
     #[test]
     fn test_solana_address() {
         let address = "GvhwZwtV32kYUXUw965CUM3KGPdtBsDwPVpi92brY5R2";
-        let parsed = Pubkey::from_base58(address).unwrap();
+        let pubkey = Pubkey::from_base58(address).unwrap();
 
         assert!(validate_address(address));
-        assert_eq!(parsed.as_bytes().len(), 32);
-        assert_eq!(parsed.to_base58(), address);
+        assert_eq!(pubkey.as_bytes().len(), 32);
+        assert_eq!(pubkey.to_base58(), address);
         assert!(!validate_address("invalid"));
     }
 }

crates/gem_solana/src/lib.rs

@@ -19,7 +19,7 @@ pub mod signer;
 
 pub use address::validate_address;
 pub use jsonrpc::SolanaRpc;
-pub use solana_primitives::{Pubkey, find_program_address};
+pub use solana_primitives::{Pubkey, SolanaError, find_program_address};
 pub use transaction::{decode_transaction, try_decode_transaction};
 
 #[cfg(all(feature = "reqwest", not(feature = "rpc")))]

crates/gem_solana/src/signer/instructions/stake_account.rs

@@ -94,7 +94,10 @@ pub(super) fn from_blockhash(sender: &Pubkey, input: &SignerInput) -> Result<Pub
 
 pub(super) fn seed_from_blockhash(input: &SignerInput) -> Result<String, SignerError> {
     let block_hash = input.metadata.get_block_hash()?;
-    Ok(block_hash[..block_hash.len().min(32)].to_string())
+    block_hash
+        .get(..block_hash.len().min(32))
+        .map(String::from)
+        .ok_or_else(|| SignerError::invalid_input("invalid Solana block hash"))
 }
 
 fn create_with_seed_instruction(sender: Pubkey, stake_account: Pubkey, seed: String, lamports: u64) -> Result<Instruction, SignerError> {
@@ -201,3 +204,21 @@ fn delegate_instruction(stake_account: Pubkey, validator: Pubkey, authority: Pub
 fn program() -> Result<Pubkey, SignerError> {
     Pubkey::from_base58(STAKE_PROGRAM_ID).map_err(SignerError::from_display)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use primitives::SignerInput;
+
+    #[test]
+    fn test_seed_from_blockhash() {
+        let valid_block_hash = "1".repeat(44);
+        assert_eq!(seed_from_blockhash(&SignerInput::mock_solana(&valid_block_hash)).unwrap(), "1".repeat(32));
+
+        let invalid_block_hash = format!("{}é", "1".repeat(31));
+        assert_eq!(
+            seed_from_blockhash(&SignerInput::mock_solana(&invalid_block_hash)).unwrap_err().to_string(),
+            "Invalid input: invalid Solana block hash"
+        );
+    }
+}

crates/gem_solana/src/signer/testkit.rs

@@ -3,7 +3,6 @@ use primitives::testkit::signer_mock::TEST_PRIVATE_KEY;
 use primitives::{SolanaTokenProgramId, TransactionLoadMetadata};
 use solana_primitives::{Pubkey, VersionedTransaction, get_address};
 
-pub const TEST_BLOCK_HASH: &str = "11111111111111111111111111111111";
 pub const TEST_RECIPIENT: &str = "EN2sCsJ1WDV8UFqsiTXHcUPUxQ4juE71eCknHYYMifkd";
 pub const TEST_SENDER_TOKEN_ADDRESS: &str = "HEeranxp3y7kVQKVSLdZW1rUmnbs7bAtUTMu8o88Jash";
 
@@ -16,12 +15,7 @@ pub fn sender_address_for_key(private_key: &[u8]) -> String {
 }
 
 pub fn solana_metadata(sender_token_address: Option<&str>, recipient_token_address: Option<&str>, token_program: Option<SolanaTokenProgramId>) -> TransactionLoadMetadata {
-    TransactionLoadMetadata::Solana {
-        sender_token_address: sender_token_address.map(String::from),
-        recipient_token_address: recipient_token_address.map(String::from),
-        token_program,
-        block_hash: TEST_BLOCK_HASH.to_string(),
-    }
+    TransactionLoadMetadata::mock_solana_token(sender_token_address, recipient_token_address, token_program)
 }
 
 pub fn private_key_base58(value: &str) -> Vec<u8> {

crates/gem_solana/src/token_account.rs

@@ -1,13 +1,16 @@
-use crate::{ASSOCIATED_TOKEN_ACCOUNT_PROGRAM, Pubkey, find_program_address};
+use std::sync::LazyLock;
 
-pub fn get_token_account(wallet: &str, token_mint: &str, token_program: &str) -> String {
-    let owner = Pubkey::from_base58(wallet).unwrap();
-    let token_program = Pubkey::from_base58(token_program).unwrap();
-    let mint = Pubkey::from_base58(token_mint).unwrap();
-    let associated_token_program = Pubkey::from_base58(ASSOCIATED_TOKEN_ACCOUNT_PROGRAM).unwrap();
+use crate::{ASSOCIATED_TOKEN_ACCOUNT_PROGRAM, Pubkey, SolanaError, find_program_address};
+
+static ASSOCIATED_TOKEN_PROGRAM: LazyLock<Pubkey> = LazyLock::new(|| Pubkey::from_base58(ASSOCIATED_TOKEN_ACCOUNT_PROGRAM).unwrap());
+
+pub fn get_token_account(wallet: &str, token_mint: &str, token_program: &str) -> Result<String, SolanaError> {
+    let owner = Pubkey::from_base58(wallet)?;
+    let token_program = Pubkey::from_base58(token_program)?;
+    let mint = Pubkey::from_base58(token_mint)?;
     let seeds = [owner.as_bytes().as_ref(), token_program.as_bytes().as_ref(), mint.as_bytes().as_ref()];
 
-    find_program_address(&associated_token_program, &seeds).unwrap().0.to_string()
+    Ok(find_program_address(&ASSOCIATED_TOKEN_PROGRAM, &seeds)?.0.to_string())
 }
 
 #[cfg(test)]
@@ -16,7 +19,7 @@ mod tests {
     use crate::{PYUSD_TOKEN_MINT, TOKEN_PROGRAM, TOKEN_PROGRAM_2022, USDC_TOKEN_MINT, USDS_TOKEN_MINT, USDT_TOKEN_MINT, WSOL_TOKEN_ADDRESS};
 
     #[test]
-    fn test_get_token_account() {
+    fn test_get_token_account() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
         let test_cases = [
             (
                 "CzVqG98YbFNiMREwgTswSML59CNrfobsNX4N9j6K8fbC",
@@ -51,8 +54,12 @@ mod tests {
         ];
 
         for (wallet, token_mint, token_program, expected_token_account) in test_cases.iter() {
-            let fee_token_account = get_token_account(wallet, token_mint, token_program);
+            let fee_token_account = get_token_account(wallet, token_mint, token_program)?;
             assert_eq!(fee_token_account, *expected_token_account);
         }
+
+        assert!(get_token_account("invalid", USDC_TOKEN_MINT, TOKEN_PROGRAM).is_err());
+
+        Ok(())
     }
 }

crates/primitives/src/testkit/transaction_load_input_mock.rs

@@ -131,6 +131,10 @@ impl SignerInput {
             TransactionFee::default(),
         )
     }
+
+    pub fn mock_solana(block_hash: &str) -> Self {
+        SignerInput::new(TransactionLoadInput::mock_solana(block_hash), TransactionFee::default())
+    }
 }
 
 impl TransactionLoadInput {
@@ -150,6 +154,19 @@ impl TransactionLoadInput {
         }
     }
 
+    pub fn mock_solana(block_hash: &str) -> Self {
+        TransactionLoadInput {
+            input_type: TransactionInputType::Transfer(Asset::mock_sol()),
+            sender_address: String::new(),
+            destination_address: String::new(),
+            value: "0".to_string(),
+            gas_price: GasPriceType::regular(0),
+            memo: None,
+            is_max_value: false,
+            metadata: TransactionLoadMetadata::mock_solana(block_hash),
+        }
+    }
+
     pub fn mock_transfer(asset: Asset, sender: &str, destination: &str, value: &str, fee: u64, memo: Option<&str>, metadata: TransactionLoadMetadata) -> Self {
         TransactionLoadInput {
             input_type: TransactionInputType::Transfer(asset),

crates/primitives/src/testkit/transaction_load_metadata_mock.rs

@@ -1,4 +1,4 @@
-use crate::{TransactionLoadMetadata, stake_type::TronStakeData};
+use crate::{SolanaTokenProgramId, TransactionLoadMetadata, stake_type::TronStakeData};
 
 impl TransactionLoadMetadata {
     pub fn mock_aptos() -> Self {
@@ -48,4 +48,22 @@ impl TransactionLoadMetadata {
             sequence,
         }
     }
+
+    pub fn mock_solana(block_hash: &str) -> Self {
+        TransactionLoadMetadata::Solana {
+            sender_token_address: None,
+            recipient_token_address: None,
+            token_program: None,
+            block_hash: block_hash.to_string(),
+        }
+    }
+
+    pub fn mock_solana_token(sender_token_address: Option<&str>, recipient_token_address: Option<&str>, token_program: Option<SolanaTokenProgramId>) -> Self {
+        TransactionLoadMetadata::Solana {
+            sender_token_address: sender_token_address.map(String::from),
+            recipient_token_address: recipient_token_address.map(String::from),
+            token_program,
+            block_hash: "11111111111111111111111111111111".to_string(),
+        }
+    }
 }

crates/swapper/src/error.rs

@@ -110,6 +110,12 @@ impl From<gem_ton::tvm::TvmError> for SwapperError {
     }
 }
 
+impl From<gem_solana::SolanaError> for SwapperError {
+    fn from(err: gem_solana::SolanaError) -> Self {
+        Self::ComputeQuoteError(format!("Solana error: {err}"))
+    }
+}
+
 impl From<primitives::AddressError> for SwapperError {
     fn from(err: primitives::AddressError) -> Self {
         Self::ComputeQuoteError(format!("{INVALID_ADDRESS}: {err}"))
@@ -145,3 +151,16 @@ impl From<number_formatter::NumberFormatterError> for SwapperError {
         Self::ComputeQuoteError(format!("{}: {err}", INVALID_AMOUNT))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_solana_error_mapping() {
+        assert_eq!(
+            SwapperError::from(gem_solana::SolanaError::InvalidTransaction),
+            SwapperError::ComputeQuoteError("Solana error: Invalid transaction".to_string())
+        );
+    }
+}

crates/swapper/src/jupiter/provider.rs

@@ -60,12 +60,12 @@ where
         input.to_string()
     }
 
-    fn get_fee_token_account(&self, options: &Options, mint: &str, token_program: &str) -> Option<String> {
+    fn get_fee_token_account(&self, options: &Options, mint: &str, token_program: &str) -> Result<Option<String>, SwapperError> {
         if let Some(fee) = &options.fee {
-            let fee_account = get_token_account(&fee.solana.address, mint, token_program);
-            return Some(fee_account);
+            let fee_account = get_token_account(&fee.solana.address, mint, token_program)?;
+            return Ok(Some(fee_account));
         }
-        None
+        Ok(None)
     }
 
     async fn fetch_token_program(&self, mint: &str) -> Result<String, SwapperError> {
@@ -83,12 +83,12 @@ where
         let fee_mint = self.get_fee_mint(input_mint, output_mint);
         // if fee_mint is in preset, no need to fetch token program
         let token_program = if self.fee_mints.contains(fee_mint.as_str()) {
-            return Ok(self.get_fee_token_account(options, fee_mint.as_str(), TOKEN_PROGRAM).unwrap_or_default());
+            return Ok(self.get_fee_token_account(options, fee_mint.as_str(), TOKEN_PROGRAM)?.unwrap_or_default());
         } else {
             self.fetch_token_program(&fee_mint).await?
         };
 
-        let mut fee_account = self.get_fee_token_account(options, &fee_mint, &token_program).unwrap_or_default();
+        let mut fee_account = self.get_fee_token_account(options, &fee_mint, &token_program)?.unwrap_or_default();
         if fee_account.is_empty() {
             return Ok(fee_account);
         }