spl_token_2022/extension/confidential_transfer/
account_info.rs

1use {
2    crate::{
3        error::TokenError,
4        extension::confidential_transfer::{
5            ConfidentialTransferAccount, DecryptableBalance, EncryptedBalance,
6            PENDING_BALANCE_LO_BIT_LENGTH,
7        },
8    },
9    bytemuck::{Pod, Zeroable},
10    solana_zk_sdk::{
11        encryption::{
12            auth_encryption::{AeCiphertext, AeKey},
13            elgamal::{ElGamalKeypair, ElGamalPubkey, ElGamalSecretKey},
14        },
15        zk_elgamal_proof_program::proof_data::ZeroCiphertextProofData,
16    },
17    spl_pod::primitives::PodU64,
18    spl_token_confidential_transfer_proof_generation::{
19        transfer::{transfer_split_proof_data, TransferProofData},
20        transfer_with_fee::{transfer_with_fee_split_proof_data, TransferWithFeeProofData},
21        withdraw::{withdraw_proof_data, WithdrawProofData},
22    },
23};
24
25/// Confidential transfer extension information needed to construct an
26/// `EmptyAccount` instruction.
27#[repr(C)]
28#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
29pub struct EmptyAccountAccountInfo {
30    /// The available balance
31    pub(crate) available_balance: EncryptedBalance,
32}
33impl EmptyAccountAccountInfo {
34    /// Create the `EmptyAccount` instruction account information from
35    /// `ConfidentialTransferAccount`.
36    pub fn new(account: &ConfidentialTransferAccount) -> Self {
37        Self {
38            available_balance: account.available_balance,
39        }
40    }
41
42    /// Create an empty account proof data.
43    pub fn generate_proof_data(
44        &self,
45        elgamal_keypair: &ElGamalKeypair,
46    ) -> Result<ZeroCiphertextProofData, TokenError> {
47        let available_balance = self
48            .available_balance
49            .try_into()
50            .map_err(|_| TokenError::MalformedCiphertext)?;
51
52        ZeroCiphertextProofData::new(elgamal_keypair, &available_balance)
53            .map_err(|_| TokenError::ProofGeneration)
54    }
55}
56
57/// Confidential Transfer extension information needed to construct an
58/// `ApplyPendingBalance` instruction.
59#[repr(C)]
60#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
61pub struct ApplyPendingBalanceAccountInfo {
62    /// The total number of `Deposit` and `Transfer` instructions that have
63    /// credited `pending_balance`
64    pub(crate) pending_balance_credit_counter: PodU64,
65    /// The low 16 bits of the pending balance (encrypted by `elgamal_pubkey`)
66    pub(crate) pending_balance_lo: EncryptedBalance,
67    /// The high 48 bits of the pending balance (encrypted by `elgamal_pubkey`)
68    pub(crate) pending_balance_hi: EncryptedBalance,
69    /// The decryptable available balance
70    pub(crate) decryptable_available_balance: DecryptableBalance,
71}
72impl ApplyPendingBalanceAccountInfo {
73    /// Create the `ApplyPendingBalance` instruction account information from
74    /// `ConfidentialTransferAccount`.
75    pub fn new(account: &ConfidentialTransferAccount) -> Self {
76        Self {
77            pending_balance_credit_counter: account.pending_balance_credit_counter,
78            pending_balance_lo: account.pending_balance_lo,
79            pending_balance_hi: account.pending_balance_hi,
80            decryptable_available_balance: account.decryptable_available_balance,
81        }
82    }
83
84    /// Return the pending balance credit counter of the account.
85    pub fn pending_balance_credit_counter(&self) -> u64 {
86        self.pending_balance_credit_counter.into()
87    }
88
89    fn decrypted_pending_balance_lo(
90        &self,
91        elgamal_secret_key: &ElGamalSecretKey,
92    ) -> Result<u64, TokenError> {
93        let pending_balance_lo = self
94            .pending_balance_lo
95            .try_into()
96            .map_err(|_| TokenError::MalformedCiphertext)?;
97        elgamal_secret_key
98            .decrypt_u32(&pending_balance_lo)
99            .ok_or(TokenError::AccountDecryption)
100    }
101
102    fn decrypted_pending_balance_hi(
103        &self,
104        elgamal_secret_key: &ElGamalSecretKey,
105    ) -> Result<u64, TokenError> {
106        let pending_balance_hi = self
107            .pending_balance_hi
108            .try_into()
109            .map_err(|_| TokenError::MalformedCiphertext)?;
110        elgamal_secret_key
111            .decrypt_u32(&pending_balance_hi)
112            .ok_or(TokenError::AccountDecryption)
113    }
114
115    fn decrypted_available_balance(&self, aes_key: &AeKey) -> Result<u64, TokenError> {
116        let decryptable_available_balance = self
117            .decryptable_available_balance
118            .try_into()
119            .map_err(|_| TokenError::MalformedCiphertext)?;
120        aes_key
121            .decrypt(&decryptable_available_balance)
122            .ok_or(TokenError::AccountDecryption)
123    }
124
125    /// Update the decryptable available balance.
126    pub fn new_decryptable_available_balance(
127        &self,
128        elgamal_secret_key: &ElGamalSecretKey,
129        aes_key: &AeKey,
130    ) -> Result<AeCiphertext, TokenError> {
131        let decrypted_pending_balance_lo = self.decrypted_pending_balance_lo(elgamal_secret_key)?;
132        let decrypted_pending_balance_hi = self.decrypted_pending_balance_hi(elgamal_secret_key)?;
133        let pending_balance =
134            combine_balances(decrypted_pending_balance_lo, decrypted_pending_balance_hi)
135                .ok_or(TokenError::AccountDecryption)?;
136        let current_available_balance = self.decrypted_available_balance(aes_key)?;
137        let new_decrypted_available_balance = current_available_balance
138            .checked_add(pending_balance)
139            .unwrap(); // total balance cannot exceed `u64`
140
141        Ok(aes_key.encrypt(new_decrypted_available_balance))
142    }
143}
144
145/// Confidential Transfer extension information needed to construct a `Withdraw`
146/// instruction.
147#[repr(C)]
148#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
149pub struct WithdrawAccountInfo {
150    /// The available balance (encrypted by `encryption_pubkey`)
151    pub available_balance: EncryptedBalance,
152    /// The decryptable available balance
153    pub decryptable_available_balance: DecryptableBalance,
154}
155impl WithdrawAccountInfo {
156    /// Create the `ApplyPendingBalance` instruction account information from
157    /// `ConfidentialTransferAccount`.
158    pub fn new(account: &ConfidentialTransferAccount) -> Self {
159        Self {
160            available_balance: account.available_balance,
161            decryptable_available_balance: account.decryptable_available_balance,
162        }
163    }
164
165    fn decrypted_available_balance(&self, aes_key: &AeKey) -> Result<u64, TokenError> {
166        let decryptable_available_balance = self
167            .decryptable_available_balance
168            .try_into()
169            .map_err(|_| TokenError::MalformedCiphertext)?;
170        aes_key
171            .decrypt(&decryptable_available_balance)
172            .ok_or(TokenError::AccountDecryption)
173    }
174
175    /// Create a withdraw proof data.
176    pub fn generate_proof_data(
177        &self,
178        withdraw_amount: u64,
179        elgamal_keypair: &ElGamalKeypair,
180        aes_key: &AeKey,
181    ) -> Result<WithdrawProofData, TokenError> {
182        let current_available_balance = self
183            .available_balance
184            .try_into()
185            .map_err(|_| TokenError::MalformedCiphertext)?;
186        let current_decrypted_available_balance = self.decrypted_available_balance(aes_key)?;
187
188        withdraw_proof_data(
189            &current_available_balance,
190            current_decrypted_available_balance,
191            withdraw_amount,
192            elgamal_keypair,
193        )
194        .map_err(|e| -> TokenError { e.into() })
195    }
196
197    /// Update the decryptable available balance.
198    pub fn new_decryptable_available_balance(
199        &self,
200        withdraw_amount: u64,
201        aes_key: &AeKey,
202    ) -> Result<AeCiphertext, TokenError> {
203        let current_decrypted_available_balance = self.decrypted_available_balance(aes_key)?;
204        let new_decrypted_available_balance = current_decrypted_available_balance
205            .checked_sub(withdraw_amount)
206            .ok_or(TokenError::InsufficientFunds)?;
207
208        Ok(aes_key.encrypt(new_decrypted_available_balance))
209    }
210}
211
212/// Confidential Transfer extension information needed to construct a `Transfer`
213/// instruction.
214#[repr(C)]
215#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
216pub struct TransferAccountInfo {
217    /// The available balance (encrypted by `encryption_pubkey`)
218    pub available_balance: EncryptedBalance,
219    /// The decryptable available balance
220    pub decryptable_available_balance: DecryptableBalance,
221}
222impl TransferAccountInfo {
223    /// Create the `Transfer` instruction account information from
224    /// `ConfidentialTransferAccount`.
225    pub fn new(account: &ConfidentialTransferAccount) -> Self {
226        Self {
227            available_balance: account.available_balance,
228            decryptable_available_balance: account.decryptable_available_balance,
229        }
230    }
231
232    fn decrypted_available_balance(&self, aes_key: &AeKey) -> Result<u64, TokenError> {
233        let decryptable_available_balance = self
234            .decryptable_available_balance
235            .try_into()
236            .map_err(|_| TokenError::MalformedCiphertext)?;
237        aes_key
238            .decrypt(&decryptable_available_balance)
239            .ok_or(TokenError::AccountDecryption)
240    }
241
242    /// Create a transfer proof data that is split into equality, ciphertext
243    /// validity, and range proofs.
244    pub fn generate_split_transfer_proof_data(
245        &self,
246        transfer_amount: u64,
247        source_elgamal_keypair: &ElGamalKeypair,
248        aes_key: &AeKey,
249        destination_elgamal_pubkey: &ElGamalPubkey,
250        auditor_elgamal_pubkey: Option<&ElGamalPubkey>,
251    ) -> Result<TransferProofData, TokenError> {
252        let current_available_balance = self
253            .available_balance
254            .try_into()
255            .map_err(|_| TokenError::MalformedCiphertext)?;
256        let current_decryptable_available_balance = self
257            .decryptable_available_balance
258            .try_into()
259            .map_err(|_| TokenError::MalformedCiphertext)?;
260
261        transfer_split_proof_data(
262            &current_available_balance,
263            &current_decryptable_available_balance,
264            transfer_amount,
265            source_elgamal_keypair,
266            aes_key,
267            destination_elgamal_pubkey,
268            auditor_elgamal_pubkey,
269        )
270        .map_err(|e| -> TokenError { e.into() })
271    }
272
273    /// Create a transfer proof data that is split into equality, ciphertext
274    /// validity (transfer amount), percentage-with-cap, ciphertext validity
275    /// (fee), and range proofs.
276    #[allow(clippy::too_many_arguments)]
277    pub fn generate_split_transfer_with_fee_proof_data(
278        &self,
279        transfer_amount: u64,
280        source_elgamal_keypair: &ElGamalKeypair,
281        aes_key: &AeKey,
282        destination_elgamal_pubkey: &ElGamalPubkey,
283        auditor_elgamal_pubkey: Option<&ElGamalPubkey>,
284        withdraw_withheld_authority_elgamal_pubkey: &ElGamalPubkey,
285        fee_rate_basis_points: u16,
286        maximum_fee: u64,
287    ) -> Result<TransferWithFeeProofData, TokenError> {
288        let current_available_balance = self
289            .available_balance
290            .try_into()
291            .map_err(|_| TokenError::MalformedCiphertext)?;
292        let current_decryptable_available_balance = self
293            .decryptable_available_balance
294            .try_into()
295            .map_err(|_| TokenError::MalformedCiphertext)?;
296
297        transfer_with_fee_split_proof_data(
298            &current_available_balance,
299            &current_decryptable_available_balance,
300            transfer_amount,
301            source_elgamal_keypair,
302            aes_key,
303            destination_elgamal_pubkey,
304            auditor_elgamal_pubkey,
305            withdraw_withheld_authority_elgamal_pubkey,
306            fee_rate_basis_points,
307            maximum_fee,
308        )
309        .map_err(|e| -> TokenError { e.into() })
310    }
311
312    /// Update the decryptable available balance.
313    pub fn new_decryptable_available_balance(
314        &self,
315        transfer_amount: u64,
316        aes_key: &AeKey,
317    ) -> Result<AeCiphertext, TokenError> {
318        let current_decrypted_available_balance = self.decrypted_available_balance(aes_key)?;
319        let new_decrypted_available_balance = current_decrypted_available_balance
320            .checked_sub(transfer_amount)
321            .ok_or(TokenError::InsufficientFunds)?;
322
323        Ok(aes_key.encrypt(new_decrypted_available_balance))
324    }
325}
326
327/// Combines pending balances low and high bits into singular pending balance
328pub fn combine_balances(balance_lo: u64, balance_hi: u64) -> Option<u64> {
329    balance_hi
330        .checked_shl(PENDING_BALANCE_LO_BIT_LENGTH)?
331        .checked_add(balance_lo)
332}