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