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