drift_rs/
wallet.rs

1use std::sync::Arc;
2
3use solana_sdk::{
4    hash::Hash,
5    message::VersionedMessage,
6    pubkey::Pubkey,
7    signature::{keypair_from_seed, Keypair, Signature},
8    signer::Signer,
9    transaction::VersionedTransaction,
10};
11
12use crate::{
13    constants::{self},
14    types::{SdkError, SdkResult},
15    utils,
16};
17
18/// Wallet operation mode
19#[derive(Clone, Debug)]
20pub enum Mode {
21    /// Normal tx signing
22    Normal,
23    /// Tx signing will fail
24    ReadOnly,
25    /// Tx signed by delegate
26    Delegated,
27}
28
29/// Drift wallet
30/// ```example(no_run)
31/// // load wallet from base58 private key or path to wallet.json
32/// let wallet = Wallet::try_from_str("PRIVATE_KEY").unwrap();
33///
34/// // construct wallet for delegated signing
35/// let delegated_keypair = solana_sdk::signature::Keypair::new();
36/// let delegated_wallet = Wallet::delegated(delegated_keypair, drift_account_key);
37///
38/// // place holder wallet for readonly apps
39/// let ro_wallet = Wallet::read_only(drift_authority);
40/// ```
41#[derive(Clone, Debug)]
42pub struct Wallet {
43    /// The signing keypair, it could be authority or delegate
44    signer: Arc<Keypair>,
45    /// The drift 'authority' account
46    /// user (sub)accounts are derived from this
47    authority: Pubkey,
48    /// The drift 'stats' account
49    stats: Pubkey,
50    /// Wallet operation mode
51    mode: Mode,
52}
53
54impl Wallet {
55    /// Returns true if the wallet is configured for delegated signing
56    pub fn is_delegated(&self) -> bool {
57        matches!(self.mode, Mode::Delegated)
58    }
59    /// Returns true if the wallet is configured for read-only/emulation mode
60    pub fn is_read_only(&self) -> bool {
61        matches!(self.mode, Mode::ReadOnly)
62    }
63    /// Init wallet from a string of either file path or the encoded private key
64    /// ```example(no_run)
65    /// // from private key str
66    /// let wallet = Wallet::try_from_str("PRIVATE_KEY").unwrap();
67    ///
68    /// // from file path
69    /// let wallet = Wallet::try_from_str("/path/to/wallet.json").unwrap();
70    /// ```
71    pub fn try_from_str(path_or_key: &str) -> SdkResult<Self> {
72        let authority = utils::load_keypair_multi_format(path_or_key)?;
73        Ok(Self::new(authority))
74    }
75    /// Construct a read-only wallet
76    pub fn read_only(authority: Pubkey) -> Self {
77        Self {
78            signer: Arc::new(Keypair::new()),
79            authority,
80            stats: Wallet::derive_stats_account(&authority),
81            mode: Mode::ReadOnly,
82        }
83    }
84    /// Init wallet from base58 encoded seed, uses default sub-account
85    ///
86    /// * `seed` - base58 encoded private key
87    ///
88    /// # panics
89    /// if the key is invalid
90    pub fn from_seed_bs58(seed: &str) -> Self {
91        let authority = Keypair::from_base58_string(seed);
92        Self::new(authority)
93    }
94    /// Init wallet from seed bytes, uses default sub-account
95    pub fn from_seed(seed: &[u8]) -> SdkResult<Self> {
96        let authority = keypair_from_seed(seed).map_err(|_| SdkError::InvalidSeed)?;
97        Ok(Self::new(authority))
98    }
99    /// Init wallet with keypair
100    ///
101    /// * `authority` - keypair for tx signing
102    pub fn new(authority: Keypair) -> Self {
103        Self {
104            stats: Wallet::derive_stats_account(&authority.pubkey()),
105            authority: authority.pubkey(),
106            signer: Arc::new(authority),
107            mode: Mode::Normal,
108        }
109    }
110    /// Convert the wallet into a delegated one by providing the `authority` public key
111    ///
112    /// * `authority` - keypair for tx signing
113    #[deprecated = "use Wallet::delegated"]
114    pub fn to_delegated(&mut self, authority: Pubkey) {
115        self.stats = Wallet::derive_stats_account(&authority);
116        self.authority = authority;
117        self.mode = Mode::Delegated;
118    }
119    /// Create a delegated wallet
120    ///
121    /// * `signer` - the delegated keypair for tx signing
122    /// * `authority` - drift account to sign for (the delegator)
123    pub fn delegated(signer: Keypair, authority: Pubkey) -> Self {
124        Self {
125            signer: Arc::new(signer),
126            stats: Wallet::derive_stats_account(&authority),
127            authority,
128            mode: Mode::Delegated,
129        }
130    }
131    /// Calculate the address of a drift user account/sub-account
132    pub fn derive_user_account(authority: &Pubkey, sub_account_id: u16) -> Pubkey {
133        let (account_drift_pda, _seed) = Pubkey::find_program_address(
134            &[
135                &b"user"[..],
136                authority.as_ref(),
137                &sub_account_id.to_le_bytes(),
138            ],
139            &constants::PROGRAM_ID,
140        );
141        account_drift_pda
142    }
143
144    /// Calculate the address of a drift stats account
145    pub fn derive_stats_account(account: &Pubkey) -> Pubkey {
146        let (account_drift_pda, _seed) = Pubkey::find_program_address(
147            &[&b"user_stats"[..], account.as_ref()],
148            &constants::PROGRAM_ID,
149        );
150        account_drift_pda
151    }
152
153    /// Calculate the address of `authority`s swift (taker) order account
154    pub fn derive_swift_order_account(authority: &Pubkey) -> Pubkey {
155        let (account_drift_pda, _seed) = Pubkey::find_program_address(
156            &[&b"SIGNED_MSG"[..], authority.as_ref()],
157            &constants::PROGRAM_ID,
158        );
159        account_drift_pda
160    }
161
162    /// Signs a solana message (ixs, accounts) and builds a signed tx
163    /// ready for sending over RPC
164    ///
165    /// * `message` - solana VersionedMessage
166    /// * `recent_block_hash` blockhash for  tx longevity
167    ///
168    /// Returns `VersionedTransaction` on success
169    pub fn sign_tx(
170        &self,
171        mut message: VersionedMessage,
172        recent_block_hash: Hash,
173    ) -> SdkResult<VersionedTransaction> {
174        message.set_recent_blockhash(recent_block_hash);
175        let signer: &dyn Signer = self.signer.as_ref();
176        match self.mode {
177            Mode::ReadOnly => Err(SdkError::WalletSigningDisabled),
178            _ => VersionedTransaction::try_new(message, &[signer]).map_err(Into::into),
179        }
180    }
181
182    /// Sign message with the wallet's signer
183    pub fn sign_message(&self, message: &[u8]) -> SdkResult<Signature> {
184        let signer: &dyn Signer = self.signer.as_ref();
185        match self.mode {
186            Mode::ReadOnly => Err(SdkError::WalletSigningDisabled),
187            _ => Ok(signer.sign_message(message)),
188        }
189    }
190    /// Return the wallet authority address
191    pub fn authority(&self) -> &Pubkey {
192        &self.authority
193    }
194    /// Return the wallet signing address
195    pub fn signer(&self) -> Pubkey {
196        self.signer.pubkey()
197    }
198    /// Return the drift user stats address
199    pub fn stats(&self) -> &Pubkey {
200        &self.stats
201    }
202    /// Return the address of the default sub-account (0)
203    pub fn default_sub_account(&self) -> Pubkey {
204        self.sub_account(0)
205    }
206    /// Calculate the drift user address given a `sub_account_id`
207    pub fn sub_account(&self, sub_account_id: u16) -> Pubkey {
208        Self::derive_user_account(self.authority(), sub_account_id)
209    }
210}
211
212impl From<Keypair> for Wallet {
213    fn from(value: Keypair) -> Self {
214        Self::new(value)
215    }
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221
222    #[test]
223    fn wallet_read_only() {
224        let keypair = Keypair::new();
225        let ro = Wallet::read_only(keypair.pubkey());
226
227        let rw = Wallet::new(keypair);
228        assert_eq!(rw.authority, ro.authority);
229        assert_eq!(rw.stats, ro.stats);
230        assert_eq!(rw.default_sub_account(), ro.default_sub_account());
231    }
232}