Skip to main content

jito_restaking_cli/
lib.rs

1use ::log::info;
2use anyhow::anyhow;
3use base64::{engine::general_purpose, Engine};
4use borsh::BorshDeserialize;
5use cli_config::CliConfig;
6use cli_signer::CliSigner;
7use jito_restaking_client_common::log::PrettyDisplay;
8use log::print_base58_tx;
9use serde::Serialize;
10use solana_account_decoder::{UiAccountEncoding, UiDataSliceConfig};
11use solana_rpc_client::nonblocking::rpc_client::RpcClient;
12use solana_rpc_client_api::{
13    config::{RpcAccountInfoConfig, RpcProgramAccountsConfig},
14    filter::{Memcmp, MemcmpEncodedBytes, RpcFilterType},
15};
16use solana_sdk::{
17    instruction::Instruction, pubkey::Pubkey, signers::Signers, transaction::Transaction,
18};
19
20pub mod cli_args;
21pub mod cli_config;
22pub mod cli_signer;
23pub mod log;
24pub mod restaking;
25pub mod restaking_handler;
26pub mod vault;
27pub mod vault_handler;
28
29pub(crate) trait CliHandler {
30    fn cli_config(&self) -> &CliConfig;
31
32    fn print_tx(&self) -> bool;
33
34    fn print_json(&self) -> bool;
35
36    fn print_json_with_reserves(&self) -> bool;
37
38    fn signer(&self) -> anyhow::Result<&CliSigner> {
39        self.cli_config()
40            .signer
41            .as_ref()
42            .ok_or_else(|| anyhow!("Signer not provided"))
43    }
44
45    /// Creates a new RPC client using the configuration from the CLI handler.
46    ///
47    /// This method constructs an RPC client with the URL and commitment level specified in the
48    /// CLI configuration. The client can be used to communicate with a Solana node for
49    /// submitting transactions, querying account data, and other RPC operations.
50    fn get_rpc_client(&self) -> RpcClient {
51        RpcClient::new_with_commitment(
52            self.cli_config().rpc_url.clone(),
53            self.cli_config().commitment,
54        )
55    }
56
57    /// Resolves a signer from a keypair path, creating a new file signer if needed
58    ///
59    /// This function:
60    /// 1. Checks if the keypair path starts with "usb://"
61    /// 2. If it does, returns the existing CLI signer
62    /// 3. If not, creates a new CliSigner from the file path
63    fn resolve_keypair<'a>(
64        &'a self,
65        keypair_path: &str,
66        owned_signer: &'a mut Option<CliSigner>,
67    ) -> anyhow::Result<&'a CliSigner> {
68        if keypair_path.starts_with("usb://") {
69            let signer = self.signer()?;
70            match signer.remote_keypair {
71                Some(_) => Ok(signer),
72                None => {
73                    let signer = CliSigner::new_ledger(keypair_path);
74                    *owned_signer = Some(signer);
75                    Ok(owned_signer.as_ref().unwrap())
76                }
77            }
78        } else {
79            let signer = CliSigner::new_keypair_from_path(keypair_path)?;
80            *owned_signer = Some(signer);
81
82            Ok(owned_signer.as_ref().unwrap())
83        }
84    }
85
86    /// Creates an RPC program accounts configuration for fetching accounts of type `T` with an optional public key filter.
87    ///
88    /// This method constructs a configuration that can be used with RPC methods to fetch program accounts
89    /// that match specific criteria. It automatically adds filters for the account data size and the discriminator
90    /// of type `T` to ensure only accounts of the expected type are returned.
91    fn get_rpc_program_accounts_config<T: jito_bytemuck::Discriminator>(
92        &self,
93        filter_pubkey: Option<(&Pubkey, usize)>,
94    ) -> anyhow::Result<RpcProgramAccountsConfig> {
95        let data_size = std::mem::size_of::<T>()
96            .checked_add(8)
97            .ok_or_else(|| anyhow!("Failed to add"))?;
98
99        let encoded_discriminator =
100            general_purpose::STANDARD.encode(vec![T::DISCRIMINATOR, 0, 0, 0, 0, 0, 0, 0]);
101        let discriminator_filter = RpcFilterType::Memcmp(Memcmp::new(
102            0,
103            MemcmpEncodedBytes::Base64(encoded_discriminator),
104        ));
105
106        let mut filters = vec![
107            RpcFilterType::DataSize(data_size as u64),
108            discriminator_filter,
109        ];
110
111        if let Some((pubkey, offset)) = filter_pubkey {
112            let pubkey_filter = RpcFilterType::Memcmp(Memcmp::new(
113                offset,
114                MemcmpEncodedBytes::Base64(general_purpose::STANDARD.encode(pubkey.to_bytes())),
115            ));
116
117            filters.push(pubkey_filter);
118        }
119
120        let config = RpcProgramAccountsConfig {
121            filters: Some(filters),
122            account_config: RpcAccountInfoConfig {
123                encoding: Some(UiAccountEncoding::Base64),
124                data_slice: Some(UiDataSliceConfig {
125                    offset: 0,
126                    length: data_size,
127                }),
128                commitment: None,
129                min_context_slot: None,
130            },
131            with_context: Some(false),
132            sort_results: Some(false),
133        };
134
135        Ok(config)
136    }
137    /// Fetches and deserializes an account
138    ///
139    /// This method retrieves account data using the configured RPC client,
140    /// then deserializes it into the specified account type using Borsh deserialization.
141    async fn get_account<T: BorshDeserialize + PrettyDisplay>(
142        &self,
143        account_pubkey: &Pubkey,
144    ) -> anyhow::Result<T> {
145        let rpc_client = self.get_rpc_client();
146
147        let account = rpc_client.get_account(account_pubkey).await?;
148        let account = T::deserialize(&mut account.data.as_slice())?;
149
150        Ok(account)
151    }
152
153    /// Processes a transaction by either printing it as Base58 or sending it.
154    ///
155    /// This method handles the logic for processing a set of instructions as a transaction.
156    /// If `print_tx` is enabled in the CLI handler (helpful for running commands in Squads), it will print the transaction in Base58 format
157    /// without sending it. Otherwise, it will submit and confirm the transaction.
158    async fn process_transaction<T>(
159        &self,
160        ixs: &[Instruction],
161        payer: &Pubkey,
162        signers: &T,
163    ) -> anyhow::Result<()>
164    where
165        T: Signers + ?Sized,
166    {
167        let rpc_client = self.get_rpc_client();
168
169        if self.print_tx() {
170            print_base58_tx(ixs);
171        } else {
172            let blockhash = rpc_client.get_latest_blockhash().await?;
173            let tx = Transaction::new_signed_with_payer(ixs, Some(payer), signers, blockhash);
174            let result = rpc_client.send_and_confirm_transaction(&tx).await?;
175
176            info!("Transaction confirmed: {:?}", result);
177        }
178
179        Ok(())
180    }
181
182    /// Prints a value either as JSON or using its pretty display format.
183    ///
184    /// This function provides flexible output formatting for any type that implements both
185    /// [`Serialize`] and [`PrettyDisplay`]. It determines the output format based on the
186    /// configuration of the containing struct.
187    ///
188    /// # Format options:
189    /// - Default: Uses the [`PrettyDisplay`] trait to format output.
190    /// - `--print-json`: Prints account information in JSON format but automatically
191    ///   filters out the `reserved` fields.
192    /// - `--print-json-with-reserves`: Prints the full account information in JSON format with
193    ///   reserved space.
194    fn print_out<T>(
195        &self,
196        index: Option<usize>,
197        address: Option<&Pubkey>,
198        value: &T,
199    ) -> anyhow::Result<()>
200    where
201        T: ?Sized + Serialize + PrettyDisplay,
202    {
203        match (self.print_json(), self.print_json_with_reserves()) {
204            (true, true) => {
205                return Err(anyhow!("Conflicting flags: both --print-json and --print-json-with-reserves are enabled. Please enable only one of these flags."));
206            }
207            (true, false) => {
208                let mut json_value = serde_json::to_value(value)?;
209                self.remove_reserved_fields(&mut json_value);
210
211                let mut account_obj = serde_json::Map::new();
212                if let Some(index) = index {
213                    account_obj.insert(
214                        "index".to_string(),
215                        serde_json::Value::String(index.to_string()),
216                    );
217                }
218                if let Some(address) = address {
219                    account_obj.insert(
220                        "address".to_string(),
221                        serde_json::Value::String(address.to_string()),
222                    );
223                }
224                account_obj.insert("data".to_string(), json_value);
225
226                let json_string = serde_json::to_string_pretty(&account_obj)?;
227
228                println!("{json_string}");
229            }
230            (false, true) => {
231                let json_string = serde_json::to_string_pretty(&value)?;
232
233                println!("{json_string}");
234            }
235            (false, false) => {
236                let type_name = std::any::type_name::<T>();
237                let msg = address.map_or("".to_string(), |address| {
238                    format!("{type_name} at {address}")
239                });
240                info!("{msg}");
241                info!("{}", value.pretty_display());
242            }
243        }
244
245        Ok(())
246    }
247
248    /// Recursively removes all "reserved" fields from a JSON value
249    fn remove_reserved_fields(&self, value: &mut serde_json::Value) {
250        if let serde_json::Value::Object(map) = value {
251            map.remove("reserved");
252            map.remove("reserved_space");
253
254            // Recursively process all remaining object values
255            for (_, v) in map.iter_mut() {
256                self.remove_reserved_fields(v);
257            }
258        } else if let serde_json::Value::Array(arr) = value {
259            // Recursively process array elements
260            for item in arr.iter_mut() {
261                self.remove_reserved_fields(item);
262            }
263        }
264    }
265}