1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

use {
    crate::{error::TokenError, extension::confidential_transfer_fee::EncryptedWithheldAmount},
    bytemuck::{Pod, Zeroable},
    solana_zk_token_sdk::{
        encryption::{
            elgamal::{ElGamalCiphertext, ElGamalKeypair, ElGamalPubkey},
            pedersen::PedersenOpening,
        },
        instruction::ciphertext_ciphertext_equality::CiphertextCiphertextEqualityProofData,
    },
};

/// Confidential transfer fee extension information needed to construct a
/// `WithdrawWithheldTokensFromMint` or `WithdrawWithheldTokensFromAccounts`
/// instruction.
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
pub struct WithheldTokensInfo {
    /// The available balance
    pub(crate) withheld_amount: EncryptedWithheldAmount,
}
impl WithheldTokensInfo {
    /// Create a `WithheldTokensInfo` from an ElGamal ciphertext.
    pub fn new(withheld_amount: &EncryptedWithheldAmount) -> Self {
        Self {
            withheld_amount: *withheld_amount,
        }
    }

    /// Create withdraw withheld proof data.
    pub fn generate_proof_data(
        &self,
        withdraw_withheld_authority_elgamal_keypair: &ElGamalKeypair,
        destination_elgamal_pubkey: &ElGamalPubkey,
    ) -> Result<CiphertextCiphertextEqualityProofData, TokenError> {
        let withheld_amount_in_mint: ElGamalCiphertext = self
            .withheld_amount
            .try_into()
            .map_err(|_| TokenError::AccountDecryption)?;

        let decrypted_withheld_amount_in_mint = withheld_amount_in_mint
            .decrypt_u32(withdraw_withheld_authority_elgamal_keypair.secret())
            .ok_or(TokenError::AccountDecryption)?;

        let destination_opening = PedersenOpening::new_rand();

        let destination_ciphertext = destination_elgamal_pubkey
            .encrypt_with(decrypted_withheld_amount_in_mint, &destination_opening);

        CiphertextCiphertextEqualityProofData::new(
            withdraw_withheld_authority_elgamal_keypair,
            destination_elgamal_pubkey,
            &withheld_amount_in_mint,
            &destination_ciphertext,
            &destination_opening,
            decrypted_withheld_amount_in_mint,
        )
        .map_err(|_| TokenError::ProofGeneration)
    }
}