fix: use the config arg in the get_balance rpc method (#547)

0xzrf · web-flow · commit f7386ab43854 · 2026-03-10T10:54:52.000-04:00
diff --git a/crates/core/src/rpc/full.rs b/crates/core/src/rpc/full.rs
@@ -2292,16 +2292,38 @@ impl Full for SurfpoolFullRpc {
         &self,
         meta: Self::Metadata,
         encoded: String,
-        _config: Option<RpcContextConfig>, // TODO: use config
+        config: Option<RpcContextConfig>,
     ) -> Result<RpcResponse<Option<u64>>> {
         let (_, message) =
             decode_and_deserialize::<VersionedMessage>(encoded, TransactionBinaryEncoding::Base64)?;
 
-        meta.with_svm_reader(|svm_reader| RpcResponse {
-            context: RpcResponseContext::new(svm_reader.get_latest_absolute_slot()),
+        let RpcContextConfig {
+            commitment,
+            min_context_slot,
+        } = config.unwrap_or_default();
+        let min_ctx_slot = min_context_slot.unwrap_or_default();
+
+        let svm_locker = meta.get_svm_locker()?;
+
+        let slot = if let Some(commitment_config) = commitment {
+            svm_locker.get_slot_for_commitment(&commitment_config)
+        } else {
+            svm_locker.get_latest_absolute_slot()
+        };
+
+        if let Some(min_slot) = min_context_slot
+            && slot < min_slot
+        {
+            return Err(RpcCustomError::MinContextSlotNotReached {
+                context_slot: min_ctx_slot,
+            }
+            .into());
+        }
+
+        Ok(RpcResponse {
+            context: RpcResponseContext::new(slot),
             value: Some((message.header().num_required_signatures as u64) * 5000),
         })
-        .map_err(Into::into)
     }
 
     fn get_stake_minimum_delegation(
@@ -2509,6 +2531,7 @@ mod tests {
     use std::thread::JoinHandle;
 
     use base64::{Engine, prelude::BASE64_STANDARD};
+    use bincode::Options;
     use crossbeam_channel::Receiver;
     use solana_account_decoder::{UiAccount, UiAccountData, UiAccountEncoding};
     use solana_client::rpc_config::RpcSimulateTransactionAccountsConfig;
@@ -2521,7 +2544,10 @@ mod tests {
     };
     use solana_pubkey::Pubkey;
     use solana_signer::Signer;
-    use solana_system_interface::{instruction as system_instruction, program as system_program};
+    use solana_system_interface::{
+        instruction::{self as system_instruction, transfer},
+        program as system_program,
+    };
     use solana_transaction::{
         Transaction,
         versioned::{Legacy, TransactionVersion},
@@ -2645,6 +2671,115 @@ mod tests {
         }
     }
 
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_get_fee_for_message() {
+        let setup = TestSetup::new(SurfpoolFullRpc);
+        let runloop_context = setup.context;
+        let rpc_server = setup.rpc;
+        let payer = Keypair::new();
+        let recipient = Pubkey::new_unique();
+        let lamports_to_send = 5 * LAMPORTS_PER_SOL;
+        let commitment_config_to_use = CommitmentConfig::confirmed();
+
+        let wrong_comm_min_ctx_slot = runloop_context
+            .svm_locker
+            .get_slot_for_commitment(&commitment_config_to_use)
+            + 10;
+
+        let wrong_min_slot = runloop_context.svm_locker.get_latest_absolute_slot() + 10;
+        let rpc_ctx_config_with_wrong_commitment = RpcContextConfig {
+            commitment: Some(commitment_config_to_use),
+            min_context_slot: Some(wrong_comm_min_ctx_slot),
+        };
+        let rpc_ctx_config_with_wrong_min_slot = RpcContextConfig {
+            commitment: None,
+            min_context_slot: Some(wrong_min_slot),
+        };
+
+        let instruction = transfer(&payer.pubkey(), &recipient, lamports_to_send);
+
+        let latest_blockhash = runloop_context
+            .svm_locker
+            .with_svm_reader(|svm| svm.latest_blockhash());
+        let message = solana_message::Message::new_with_blockhash(
+            &[instruction],
+            Some(&payer.pubkey()),
+            &latest_blockhash,
+        );
+        let num_required_signatures = message.header.num_required_signatures as u64;
+        let transaction =
+            VersionedTransaction::try_new(VersionedMessage::Legacy(message), &[&payer]).unwrap();
+
+        let message_bytes = bincode::options()
+            .with_fixint_encoding()
+            .serialize(&transaction.message)
+            .expect("message serialization");
+        let encoded_message = base64::engine::general_purpose::STANDARD.encode(&message_bytes);
+
+        let get_fee_with_correct_config_pass_result = rpc_server.get_fee_for_message(
+            Some(runloop_context.clone()),
+            encoded_message.clone(),
+            None,
+        );
+
+        assert!(
+            get_fee_with_correct_config_pass_result.is_ok(),
+            "Expected get_fee_for_message to pass with correct configs"
+        );
+        assert_eq!(
+            get_fee_with_correct_config_pass_result
+                .unwrap()
+                .value
+                .unwrap(),
+            (num_required_signatures as u64) * 5_000,
+            "Invalid return value"
+        );
+
+        let get_fee_with_wrong_commitment_fail_result = rpc_server.get_fee_for_message(
+            Some(runloop_context.clone()),
+            encoded_message.clone(),
+            Some(rpc_ctx_config_with_wrong_commitment),
+        );
+
+        let wrong_comm_expected_err: Result<()> = Result::Err(
+            RpcCustomError::MinContextSlotNotReached {
+                context_slot: wrong_comm_min_ctx_slot,
+            }
+            .into(),
+        );
+
+        assert!(
+            get_fee_with_wrong_commitment_fail_result.is_err(),
+            "expected this txn to fail when min_ctx_slot > slot_for_commitment"
+        );
+
+        assert_eq!(
+            get_fee_with_wrong_commitment_fail_result.err().unwrap(),
+            wrong_comm_expected_err.err().unwrap()
+        );
+
+        let get_fee_with_wrong_mint_slot_fail_result = rpc_server.get_fee_for_message(
+            Some(runloop_context.clone()),
+            encoded_message,
+            Some(rpc_ctx_config_with_wrong_min_slot),
+        );
+
+        let wrong_min_slot_expected_err: Result<()> = Result::Err(
+            RpcCustomError::MinContextSlotNotReached {
+                context_slot: wrong_min_slot,
+            }
+            .into(),
+        );
+        assert!(
+            get_fee_with_wrong_mint_slot_fail_result.is_err(),
+            "expected this txn to fail when min_ctx_slot > absolute_latest_slot"
+        );
+        assert_eq!(
+            get_fee_with_wrong_mint_slot_fail_result.err().unwrap(),
+            wrong_min_slot_expected_err.err().unwrap()
+        );
+    }
+
     #[tokio::test(flavor = "multi_thread")]
     async fn test_get_signature_statuses() {
         let pks = (0..10).map(|_| Pubkey::new_unique());
diff --git a/crates/core/src/rpc/minimal.rs b/crates/core/src/rpc/minimal.rs
@@ -12,7 +12,7 @@ use solana_client::{
     },
 };
 use solana_clock::Slot;
-use solana_commitment_config::{CommitmentConfig, CommitmentLevel};
+use solana_commitment_config::CommitmentLevel;
 use solana_epoch_info::EpochInfo;
 use solana_rpc_client_api::response::Response as RpcResponse;
 
@@ -88,7 +88,7 @@ pub trait Minimal {
         &self,
         meta: Self::Metadata,
         pubkey_str: String,
-        _config: Option<RpcContextConfig>,
+        config: Option<RpcContextConfig>,
     ) -> BoxFuture<Result<RpcResponse<u64>>>;
 
     /// Returns information about the current epoch.
@@ -586,17 +586,21 @@ impl Minimal for SurfpoolMinimalRpc {
         &self,
         meta: Self::Metadata,
         pubkey_str: String,
-        _config: Option<RpcContextConfig>, // TODO: use config
+        config: Option<RpcContextConfig>,
     ) -> BoxFuture<Result<RpcResponse<u64>>> {
         let pubkey = match verify_pubkey(&pubkey_str) {
             Ok(res) => res,
             Err(e) => return e.into(),
         };
 
+        let config = config.unwrap_or_default();
+        let commitment_config = config.commitment.unwrap_or_default();
+        let min_ctx_slot = config.min_context_slot;
+
         let SurfnetRpcContext {
             svm_locker,
             remote_ctx,
-        } = match meta.get_rpc_context(CommitmentConfig::confirmed()) {
+        } = match meta.get_rpc_context(commitment_config) {
             Ok(res) => res,
             Err(e) => return e.into(),
         };
@@ -608,6 +612,15 @@ impl Minimal for SurfpoolMinimalRpc {
                 ..
             } = svm_locker.get_account(&remote_ctx, &pubkey, None).await?;
 
+            if let Some(min_slot) = min_ctx_slot
+                && slot < min_slot
+            {
+                return Err(RpcCustomError::MinContextSlotNotReached {
+                    context_slot: min_slot,
+                }
+                .into());
+            }
+
             let balance = match &account_update {
                 GetAccountResult::FoundAccount(_, account, _)
                 | GetAccountResult::FoundProgramAccount((_, account), _)
@@ -972,6 +985,72 @@ mod tests {
         assert_eq!(result.unwrap(), "ok");
     }
 
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_get_balance() {
+        let setup = TestSetup::new(SurfpoolMinimalRpc);
+
+        let airdrop_amount = 5 * 1_000_000_000u64;
+        let to_airdrop_pubkey = Pubkey::new_unique();
+
+        setup
+            .context
+            .svm_locker
+            .airdrop(&to_airdrop_pubkey, airdrop_amount)
+            .unwrap()
+            .unwrap();
+
+        let pass_if_correct_config_result = setup
+            .rpc
+            .get_balance(
+                Some(setup.context.clone()),
+                to_airdrop_pubkey.to_string(),
+                None,
+            )
+            .await;
+
+        assert!(
+            pass_if_correct_config_result.is_ok(),
+            "Expected the operation to pass"
+        );
+
+        assert_eq!(
+            pass_if_correct_config_result.unwrap().value,
+            airdrop_amount,
+            "Invalid returned lamports for the account"
+        );
+
+        let wrong_min_slot = setup.context.svm_locker.get_latest_absolute_slot() + 100;
+
+        let fail_if_latest_slot_lt_min_ctx_slot_result = setup
+            .rpc
+            .get_balance(
+                Some(setup.context.clone()),
+                Pubkey::new_unique().to_string(),
+                Some(RpcContextConfig {
+                    commitment: None,
+                    min_context_slot: Some(wrong_min_slot),
+                }),
+            )
+            .await;
+
+        let expected_err: Result<()> = Result::Err(
+            RpcCustomError::MinContextSlotNotReached {
+                context_slot: wrong_min_slot,
+            }
+            .into(),
+        );
+
+        assert!(
+            fail_if_latest_slot_lt_min_ctx_slot_result.is_err(),
+            "Expected get_balance rpc method to fail when latest_absolute_slot < min_context_slot"
+        );
+
+        assert_eq!(
+            fail_if_latest_slot_lt_min_ctx_slot_result.err().unwrap(),
+            expected_err.err().unwrap()
+        );
+    }
+
     #[test]
     fn test_get_transaction_count() {
         let setup = TestSetup::new(SurfpoolMinimalRpc);
diff --git a/crates/core/src/tests/integration.rs b/crates/core/src/tests/integration.rs
@@ -1,6 +1,7 @@
 use std::{str::FromStr, sync::Arc, time::Duration};
 
 use base64::Engine;
+use bincode::Options;
 use crossbeam_channel::{unbounded, unbounded as crossbeam_unbounded};
 use jsonrpc_core::{
     Error, Result as JsonRpcResult,
@@ -11,7 +12,8 @@ use solana_account::Account;
 use solana_account_decoder::{UiAccountData, UiAccountEncoding, parse_account_data::ParsedAccount};
 use solana_address_lookup_table_interface::state::{AddressLookupTable, LookupTableMeta};
 use solana_client::{
-    nonblocking::rpc_client::RpcClient, rpc_config::RpcSimulateTransactionConfig,
+    nonblocking::rpc_client::RpcClient,
+    rpc_config::{RpcContextConfig, RpcSimulateTransactionConfig},
     rpc_response::RpcLogsResponse,
 };
 use solana_clock::{Clock, Slot};
@@ -49,8 +51,8 @@ use crate::{
     error::SurfpoolError,
     rpc::{
         RunloopContext,
-        full::FullClient,
-        minimal::MinimalClient,
+        full::{Full, FullClient, SurfpoolFullRpc},
+        minimal::{Minimal, MinimalClient, SurfpoolMinimalRpc},
         surfnet_cheatcodes::{SurfnetCheatcodes, SurfnetCheatcodesRpc},
     },
     runloops::start_local_surfnet_runloop,
