shadow_drive_sdk/client/
add_storage.rs

1use anchor_lang::{system_program, InstructionData, ToAccountMetas};
2use byte_unit::Byte;
3use shadow_drive_user_staking::accounts as shdw_drive_accounts;
4use shadow_drive_user_staking::instruction as shdw_drive_instructions;
5use solana_client::{
6    client_error::{ClientError, ClientErrorKind},
7    rpc_request::RpcError,
8};
9use solana_sdk::{
10    instruction::Instruction, pubkey::Pubkey, signer::Signer, transaction::Transaction,
11};
12use spl_associated_token_account::get_associated_token_address;
13use spl_token::ID as TokenProgramID;
14
15use super::ShadowDriveClient;
16use crate::constants::UPLOADER;
17use crate::models::storage_acct::{StorageAccount, StorageAccountV2, StorageAcct};
18use crate::serialize_and_encode;
19use crate::{
20    constants::{PROGRAM_ADDRESS, STORAGE_CONFIG_PDA, TOKEN_MINT},
21    derived_addresses,
22    error::Error,
23    models::*,
24};
25
26impl<T> ShadowDriveClient<T>
27where
28    T: Signer,
29{
30    /// Adds storage capacity to the specified [`StorageAccount`](crate::models::StorageAccount).
31    /// * `storage_account_key` - The public key of the [`StorageAccount`](crate::models::StorageAccount).
32    /// * `size` - The additional amount of storage you want to add.
33    /// E.g if you have an existing [`StorageAccount`](crate::models::StorageAccount) with 1MB of storage
34    /// but you need 2MB total, `size` should equal 1MB.
35    /// When specifying size, only KB, MB, and GB storage units are currently supported.
36    /// # Example
37    ///
38    /// ```
39    /// # use byte_unit::Byte;
40    /// # use shadow_drive_rust::{ShadowDriveClient, derived_addresses::storage_account};
41    /// # use solana_client::rpc_client::RpcClient;
42    /// # use solana_sdk::{
43    /// # pubkey::Pubkey,
44    /// # signature::Keypair,
45    /// # signer::{keypair::read_keypair_file, Signer},
46    /// # };
47    /// #
48    /// # let keypair = read_keypair_file(KEYPAIR_PATH).expect("failed to load keypair at path");
49    /// # let user_pubkey = keypair.pubkey();
50    /// # let rpc_client = RpcClient::new("https://ssc-dao.genesysgo.net");
51    /// # let shdw_drive_client = ShadowDriveClient::new(keypair, rpc_client);
52    /// # let (storage_account_key, _) = storage_account(&user_pubkey, 0);
53    /// # let added_bytes = Byte::from_str("1MB").expect("invalid byte string");
54    /// #
55    /// let add_storage_response = shdw_drive_client
56    ///     .add_storage(&storage_account_key, added_bytes)
57    ///     .await?;
58    /// ```
59    pub async fn add_storage(
60        &self,
61        storage_account_key: &Pubkey,
62        size: Byte,
63    ) -> ShadowDriveResult<StorageResponse> {
64        let size_as_bytes: u64 = size
65            .get_bytes()
66            .try_into()
67            .map_err(|_| Error::InvalidStorage)?;
68
69        let wallet_pubkey = self.wallet.pubkey();
70        let (user_info, _) = derived_addresses::user_info(&wallet_pubkey);
71
72        let user_info_acct = self.rpc_client.get_account(&user_info).await;
73        match user_info_acct {
74            Ok(_) => {
75                // the user_info_acct exists. don't need to verify anything about it as
76                // the txn will fail if self.wallet is not the owner of the storage_account
77            }
78            Err(ClientError {
79                kind: ClientErrorKind::RpcError(RpcError::ForUser(_)),
80                ..
81            }) => {
82                // this is what rpc_client.get_account() returns if the account doesn't exist
83                // If userInfo hasn't been initialized, error out
84                return Err(Error::UserInfoNotCreated);
85            }
86            Err(err) => {
87                //a different rpc error occurred
88                return Err(Error::from(err));
89            }
90        }
91
92        let selected_storage_acct = self.get_storage_account(storage_account_key).await?;
93
94        let txn_encoded = match selected_storage_acct {
95            StorageAcct::V1(storage_account) => {
96                self.add_storage_v1(storage_account_key, storage_account, size_as_bytes)
97                    .await?
98            }
99            StorageAcct::V2(storage_account) => {
100                self.add_storage_v2(storage_account_key, storage_account, size_as_bytes)
101                    .await?
102            }
103        };
104
105        self.send_shdw_txn("add-storage", txn_encoded, None).await
106    }
107
108    async fn add_storage_v1(
109        &self,
110        storage_account_key: &Pubkey,
111        storage_account: StorageAccount,
112        size_as_bytes: u64,
113    ) -> ShadowDriveResult<String> {
114        let wallet_pubkey = &self.wallet.pubkey();
115        let owner_ata = get_associated_token_address(wallet_pubkey, &TOKEN_MINT);
116        let (stake_account, _) = derived_addresses::stake_account(storage_account_key);
117
118        let accounts = shdw_drive_accounts::IncreaseStorageV1 {
119            storage_config: *STORAGE_CONFIG_PDA,
120            storage_account: *storage_account_key,
121            owner: storage_account.owner_1,
122            owner_ata,
123            stake_account,
124            uploader: UPLOADER,
125            token_mint: TOKEN_MINT,
126            system_program: system_program::ID,
127            token_program: TokenProgramID,
128        };
129        let args = shdw_drive_instructions::IncreaseStorage {
130            additional_storage: size_as_bytes,
131        };
132
133        let instruction = Instruction {
134            program_id: PROGRAM_ADDRESS,
135            accounts: accounts.to_account_metas(None),
136            data: args.data(),
137        };
138
139        let mut txn = Transaction::new_with_payer(&[instruction], Some(wallet_pubkey));
140
141        txn.try_partial_sign(
142            &[&self.wallet],
143            self.rpc_client.get_latest_blockhash().await?,
144        )?;
145
146        let txn_encoded = serialize_and_encode(&txn)?;
147
148        Ok(txn_encoded)
149    }
150
151    async fn add_storage_v2(
152        &self,
153        storage_account_key: &Pubkey,
154        storage_account: StorageAccountV2,
155        size_as_bytes: u64,
156    ) -> ShadowDriveResult<String> {
157        let wallet_pubkey = &self.wallet.pubkey();
158        let owner_ata = get_associated_token_address(wallet_pubkey, &TOKEN_MINT);
159        let (stake_account, _) = derived_addresses::stake_account(storage_account_key);
160
161        let accounts = shdw_drive_accounts::IncreaseStorageV2 {
162            storage_config: *STORAGE_CONFIG_PDA,
163            storage_account: *storage_account_key,
164            owner: storage_account.owner_1,
165            owner_ata,
166            stake_account,
167            uploader: UPLOADER,
168            token_mint: TOKEN_MINT,
169            system_program: system_program::ID,
170            token_program: TokenProgramID,
171        };
172
173        let args = shdw_drive_instructions::IncreaseStorage2 {
174            additional_storage: size_as_bytes,
175        };
176
177        let instruction = Instruction {
178            program_id: PROGRAM_ADDRESS,
179            accounts: accounts.to_account_metas(None),
180            data: args.data(),
181        };
182
183        let mut txn = Transaction::new_with_payer(&[instruction], Some(wallet_pubkey));
184
185        txn.try_partial_sign(
186            &[&self.wallet],
187            self.rpc_client.get_latest_blockhash().await?,
188        )?;
189
190        let txn_encoded = serialize_and_encode(&txn)?;
191
192        Ok(txn_encoded)
193    }
194}
shadow_drive_sdk/client/add_storage.rs

shadow_drive_sdk/client/
add_storage.rs
