1use {
2 crate::cli::CliError, solana_cli_output::display::build_balance_message,
3 solana_commitment_config::CommitmentConfig, solana_message::Message, solana_pubkey::Pubkey,
4 solana_rpc_client::nonblocking::rpc_client::RpcClient,
5 solana_rpc_client_api::client_error::Result as ClientResult,
6};
7
8pub async fn check_account_for_fee(
9 rpc_client: &RpcClient,
10 account_pubkey: &Pubkey,
11 message: &Message,
12) -> Result<(), CliError> {
13 check_account_for_multiple_fees(rpc_client, account_pubkey, &[message]).await
14}
15
16pub async fn check_account_for_fee_with_commitment(
17 rpc_client: &RpcClient,
18 account_pubkey: &Pubkey,
19 message: &Message,
20 commitment: CommitmentConfig,
21) -> Result<(), CliError> {
22 check_account_for_multiple_fees_with_commitment(
23 rpc_client,
24 account_pubkey,
25 &[message],
26 commitment,
27 )
28 .await
29}
30
31pub async fn check_account_for_multiple_fees(
32 rpc_client: &RpcClient,
33 account_pubkey: &Pubkey,
34 messages: &[&Message],
35) -> Result<(), CliError> {
36 check_account_for_multiple_fees_with_commitment(
37 rpc_client,
38 account_pubkey,
39 messages,
40 CommitmentConfig::default(),
41 )
42 .await
43}
44
45pub async fn check_account_for_multiple_fees_with_commitment(
46 rpc_client: &RpcClient,
47 account_pubkey: &Pubkey,
48 messages: &[&Message],
49 commitment: CommitmentConfig,
50) -> Result<(), CliError> {
51 check_account_for_spend_multiple_fees_with_commitment(
52 rpc_client,
53 account_pubkey,
54 0,
55 messages,
56 commitment,
57 )
58 .await
59}
60
61pub async fn check_account_for_spend_multiple_fees_with_commitment(
62 rpc_client: &RpcClient,
63 account_pubkey: &Pubkey,
64 balance: u64,
65 messages: &[&Message],
66 commitment: CommitmentConfig,
67) -> Result<(), CliError> {
68 let fee = get_fee_for_messages(rpc_client, messages).await?;
69 check_account_for_spend_and_fee_with_commitment(
70 rpc_client,
71 account_pubkey,
72 balance,
73 fee,
74 commitment,
75 )
76 .await
77}
78
79pub async fn check_account_for_spend_and_fee_with_commitment(
80 rpc_client: &RpcClient,
81 account_pubkey: &Pubkey,
82 balance: u64,
83 fee: u64,
84 commitment: CommitmentConfig,
85) -> Result<(), CliError> {
86 let required_balance =
87 balance
88 .checked_add(fee)
89 .ok_or(CliError::InsufficientFundsForSpendAndFee(
90 build_balance_message(balance, false, false),
91 build_balance_message(fee, false, false),
92 *account_pubkey,
93 ))?;
94
95 if !check_account_for_balance_with_commitment(
96 rpc_client,
97 account_pubkey,
98 required_balance,
99 commitment,
100 )
101 .await?
102 {
103 if balance > 0 {
104 return Err(CliError::InsufficientFundsForSpendAndFee(
105 build_balance_message(balance, false, false),
106 build_balance_message(fee, false, false),
107 *account_pubkey,
108 ));
109 } else {
110 return Err(CliError::InsufficientFundsForFee(
111 build_balance_message(fee, false, false),
112 *account_pubkey,
113 ));
114 }
115 }
116 Ok(())
117}
118
119pub async fn get_fee_for_messages(
120 rpc_client: &RpcClient,
121 messages: &[&Message],
122) -> Result<u64, CliError> {
123 let mut total_fee = 0u64;
124 for message in messages {
125 let fee = rpc_client.get_fee_for_message(*message).await?;
126 total_fee = total_fee
127 .checked_add(fee)
128 .ok_or(CliError::BadParameter("Fee overflow".to_string()))?;
129 }
130 Ok(total_fee)
131}
132
133pub async fn check_account_for_balance(
134 rpc_client: &RpcClient,
135 account_pubkey: &Pubkey,
136 balance: u64,
137) -> ClientResult<bool> {
138 check_account_for_balance_with_commitment(
139 rpc_client,
140 account_pubkey,
141 balance,
142 CommitmentConfig::default(),
143 )
144 .await
145}
146
147pub async fn check_account_for_balance_with_commitment(
148 rpc_client: &RpcClient,
149 account_pubkey: &Pubkey,
150 balance: u64,
151 commitment: CommitmentConfig,
152) -> ClientResult<bool> {
153 let lamports = rpc_client
154 .get_balance_with_commitment(account_pubkey, commitment)
155 .await?
156 .value;
157 if lamports != 0 && lamports >= balance {
158 return Ok(true);
159 }
160 Ok(false)
161}
162
163pub fn check_unique_pubkeys(
164 pubkey0: (&Pubkey, String),
165 pubkey1: (&Pubkey, String),
166) -> Result<(), CliError> {
167 if pubkey0.0 == pubkey1.0 {
168 Err(CliError::BadParameter(format!(
169 "Identical pubkeys found: `{}` and `{}` must be unique",
170 pubkey0.1, pubkey1.1
171 )))
172 } else {
173 Ok(())
174 }
175}
176
177#[cfg(test)]
178mod tests {
179 use {
180 super::*,
181 serde_json::json,
182 solana_rpc_client_api::{
183 request::RpcRequest,
184 response::{Response, RpcResponseContext},
185 },
186 solana_system_interface::instruction as system_instruction,
187 std::collections::HashMap,
188 };
189
190 #[tokio::test]
191 async fn test_check_account_for_fees() {
192 let account_balance = 1;
193 let account_balance_response = json!(Response {
194 context: RpcResponseContext {
195 slot: 1,
196 api_version: None
197 },
198 value: json!(account_balance),
199 });
200 let pubkey = solana_pubkey::new_rand();
201
202 let pubkey0 = Pubkey::from([0; 32]);
203 let pubkey1 = Pubkey::from([1; 32]);
204 let ix0 = system_instruction::transfer(&pubkey0, &pubkey1, 1);
205 let message0 = Message::new(&[ix0], Some(&pubkey0));
206
207 let ix0 = system_instruction::transfer(&pubkey0, &pubkey1, 1);
208 let ix1 = system_instruction::transfer(&pubkey1, &pubkey0, 1);
209 let message1 = Message::new(&[ix0, ix1], Some(&pubkey0));
210
211 let mut mocks = HashMap::new();
212 mocks.insert(RpcRequest::GetBalance, account_balance_response.clone());
213 let rpc_client = RpcClient::new_mock_with_mocks("".to_string(), mocks);
214 check_account_for_fee(&rpc_client, &pubkey, &message0)
215 .await
216 .expect("unexpected result");
217
218 let check_fee_response = json!(Response {
219 context: RpcResponseContext {
220 slot: 1,
221 api_version: None
222 },
223 value: json!(2),
224 });
225 let mut mocks = HashMap::new();
226 mocks.insert(RpcRequest::GetFeeForMessage, check_fee_response);
227 mocks.insert(RpcRequest::GetBalance, account_balance_response.clone());
228 let rpc_client = RpcClient::new_mock_with_mocks("".to_string(), mocks);
229 assert!(
230 check_account_for_fee(&rpc_client, &pubkey, &message1)
231 .await
232 .is_err()
233 );
234
235 let check_fee_response = json!(Response {
236 context: RpcResponseContext {
237 slot: 1,
238 api_version: None
239 },
240 value: json!(2),
241 });
242 let mut mocks = HashMap::new();
243 mocks.insert(RpcRequest::GetFeeForMessage, check_fee_response);
244 mocks.insert(RpcRequest::GetBalance, account_balance_response);
245 let rpc_client = RpcClient::new_mock_with_mocks("".to_string(), mocks);
246 assert!(
247 check_account_for_multiple_fees(&rpc_client, &pubkey, &[&message0, &message0])
248 .await
249 .is_err()
250 );
251
252 let account_balance = 2;
253 let account_balance_response = json!(Response {
254 context: RpcResponseContext {
255 slot: 1,
256 api_version: None
257 },
258 value: json!(account_balance),
259 });
260 let check_fee_response = json!(Response {
261 context: RpcResponseContext {
262 slot: 1,
263 api_version: None
264 },
265 value: json!(1),
266 });
267
268 let mut mocks = HashMap::new();
269 mocks.insert(RpcRequest::GetFeeForMessage, check_fee_response);
270 mocks.insert(RpcRequest::GetBalance, account_balance_response);
271 let rpc_client = RpcClient::new_mock_with_mocks("".to_string(), mocks);
272
273 check_account_for_multiple_fees(&rpc_client, &pubkey, &[&message0, &message0])
274 .await
275 .expect("unexpected result");
276 }
277
278 #[tokio::test]
279 async fn test_check_account_for_balance() {
280 let account_balance = 50;
281 let account_balance_response = json!(Response {
282 context: RpcResponseContext {
283 slot: 1,
284 api_version: None
285 },
286 value: json!(account_balance),
287 });
288 let pubkey = solana_pubkey::new_rand();
289
290 let mut mocks = HashMap::new();
291 mocks.insert(RpcRequest::GetBalance, account_balance_response);
292 let rpc_client = RpcClient::new_mock_with_mocks("".to_string(), mocks);
293
294 assert!(
295 check_account_for_balance(&rpc_client, &pubkey, 1)
296 .await
297 .unwrap()
298 );
299 assert!(
300 check_account_for_balance(&rpc_client, &pubkey, account_balance)
301 .await
302 .unwrap()
303 );
304 assert!(
305 !check_account_for_balance(&rpc_client, &pubkey, account_balance + 1)
306 .await
307 .unwrap()
308 );
309 }
310
311 #[tokio::test]
312 async fn test_get_fee_for_messages() {
313 let check_fee_response = json!(Response {
314 context: RpcResponseContext {
315 slot: 1,
316 api_version: None
317 },
318 value: json!(1),
319 });
320 let mut mocks = HashMap::new();
321 mocks.insert(RpcRequest::GetFeeForMessage, check_fee_response);
322 let rpc_client = RpcClient::new_mock_with_mocks("".to_string(), mocks);
323
324 assert_eq!(get_fee_for_messages(&rpc_client, &[]).await.unwrap(), 0);
326
327 let pubkey0 = Pubkey::from([0; 32]);
329 let pubkey1 = Pubkey::from([1; 32]);
330 let ix0 = system_instruction::transfer(&pubkey0, &pubkey1, 1);
331 let message0 = Message::new(&[ix0], Some(&pubkey0));
332 assert_eq!(
333 get_fee_for_messages(&rpc_client, &[&message0])
334 .await
335 .unwrap(),
336 1
337 );
338
339 let check_fee_response = json!(Response {
341 context: RpcResponseContext {
342 slot: 1,
343 api_version: None
344 },
345 value: json!(0),
346 });
347 let mut mocks = HashMap::new();
348 mocks.insert(RpcRequest::GetFeeForMessage, check_fee_response);
349 let rpc_client = RpcClient::new_mock_with_mocks("".to_string(), mocks);
350 let message = Message::default();
351 assert_eq!(
352 get_fee_for_messages(&rpc_client, &[&message, &message])
353 .await
354 .unwrap(),
355 0
356 );
357 }
358
359 #[test]
360 fn test_check_unique_pubkeys() {
361 let pubkey0 = solana_pubkey::new_rand();
362 let pubkey_clone = pubkey0;
363 let pubkey1 = solana_pubkey::new_rand();
364
365 check_unique_pubkeys((&pubkey0, "foo".to_string()), (&pubkey1, "bar".to_string()))
366 .expect("unexpected result");
367 check_unique_pubkeys((&pubkey0, "foo".to_string()), (&pubkey1, "foo".to_string()))
368 .expect("unexpected result");
369
370 assert!(
371 check_unique_pubkeys(
372 (&pubkey0, "foo".to_string()),
373 (&pubkey_clone, "bar".to_string())
374 )
375 .is_err()
376 );
377 }
378}