solana_zk_token_sdk/zk_token_elgamal/pod/
grouped_elgamal.rs

1//! Plain Old Data types for the Grouped ElGamal encryption scheme.
2
3#[cfg(not(target_os = "solana"))]
4use crate::encryption::grouped_elgamal::GroupedElGamalCiphertext;
5use {
6    crate::{
7        errors::ElGamalError,
8        zk_token_elgamal::pod::{
9            elgamal::{ElGamalCiphertext, DECRYPT_HANDLE_LEN, ELGAMAL_CIPHERTEXT_LEN},
10            pedersen::{PedersenCommitment, PEDERSEN_COMMITMENT_LEN},
11        },
12    },
13    bytemuck::Zeroable,
14    std::fmt,
15};
16
17macro_rules! impl_extract {
18    (TYPE = $type:ident) => {
19        impl $type {
20            /// Extract the commitment component from a grouped ciphertext
21            pub fn extract_commitment(&self) -> PedersenCommitment {
22                // `GROUPED_ELGAMAL_CIPHERTEXT_2_HANDLES` guaranteed to be at least `PEDERSEN_COMMITMENT_LEN`
23                let commitment = self.0[..PEDERSEN_COMMITMENT_LEN].try_into().unwrap();
24                PedersenCommitment(commitment)
25            }
26
27            /// Extract a regular ElGamal ciphertext using the decrypt handle at a specified index.
28            pub fn try_extract_ciphertext(
29                &self,
30                index: usize,
31            ) -> Result<ElGamalCiphertext, ElGamalError> {
32                let mut ciphertext_bytes = [0u8; ELGAMAL_CIPHERTEXT_LEN];
33                ciphertext_bytes[..PEDERSEN_COMMITMENT_LEN]
34                    .copy_from_slice(&self.0[..PEDERSEN_COMMITMENT_LEN]);
35
36                let handle_start = DECRYPT_HANDLE_LEN
37                    .checked_mul(index)
38                    .and_then(|n| n.checked_add(PEDERSEN_COMMITMENT_LEN))
39                    .ok_or(ElGamalError::CiphertextDeserialization)?;
40                let handle_end = handle_start
41                    .checked_add(DECRYPT_HANDLE_LEN)
42                    .ok_or(ElGamalError::CiphertextDeserialization)?;
43                ciphertext_bytes[PEDERSEN_COMMITMENT_LEN..].copy_from_slice(
44                    self.0
45                        .get(handle_start..handle_end)
46                        .ok_or(ElGamalError::CiphertextDeserialization)?,
47                );
48
49                Ok(ElGamalCiphertext(ciphertext_bytes))
50            }
51        }
52    };
53}
54
55/// Byte length of a grouped ElGamal ciphertext with 2 handles
56const GROUPED_ELGAMAL_CIPHERTEXT_2_HANDLES: usize =
57    PEDERSEN_COMMITMENT_LEN + DECRYPT_HANDLE_LEN + DECRYPT_HANDLE_LEN;
58
59/// Byte length of a grouped ElGamal ciphertext with 3 handles
60const GROUPED_ELGAMAL_CIPHERTEXT_3_HANDLES: usize =
61    PEDERSEN_COMMITMENT_LEN + DECRYPT_HANDLE_LEN + DECRYPT_HANDLE_LEN + DECRYPT_HANDLE_LEN;
62
63/// The `GroupedElGamalCiphertext` type with two decryption handles as a `Pod`
64#[derive(Clone, Copy, bytemuck_derive::Pod, bytemuck_derive::Zeroable, PartialEq, Eq)]
65#[repr(transparent)]
66pub struct GroupedElGamalCiphertext2Handles(pub [u8; GROUPED_ELGAMAL_CIPHERTEXT_2_HANDLES]);
67
68impl fmt::Debug for GroupedElGamalCiphertext2Handles {
69    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
70        write!(f, "{:?}", self.0)
71    }
72}
73
74impl Default for GroupedElGamalCiphertext2Handles {
75    fn default() -> Self {
76        Self::zeroed()
77    }
78}
79#[cfg(not(target_os = "solana"))]
80impl From<GroupedElGamalCiphertext<2>> for GroupedElGamalCiphertext2Handles {
81    fn from(decoded_ciphertext: GroupedElGamalCiphertext<2>) -> Self {
82        Self(decoded_ciphertext.to_bytes().try_into().unwrap())
83    }
84}
85
86#[cfg(not(target_os = "solana"))]
87impl TryFrom<GroupedElGamalCiphertext2Handles> for GroupedElGamalCiphertext<2> {
88    type Error = ElGamalError;
89
90    fn try_from(pod_ciphertext: GroupedElGamalCiphertext2Handles) -> Result<Self, Self::Error> {
91        Self::from_bytes(&pod_ciphertext.0).ok_or(ElGamalError::CiphertextDeserialization)
92    }
93}
94
95impl_extract!(TYPE = GroupedElGamalCiphertext2Handles);
96
97/// The `GroupedElGamalCiphertext` type with three decryption handles as a `Pod`
98#[derive(Clone, Copy, bytemuck_derive::Pod, bytemuck_derive::Zeroable, PartialEq, Eq)]
99#[repr(transparent)]
100pub struct GroupedElGamalCiphertext3Handles(pub [u8; GROUPED_ELGAMAL_CIPHERTEXT_3_HANDLES]);
101
102impl fmt::Debug for GroupedElGamalCiphertext3Handles {
103    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
104        write!(f, "{:?}", self.0)
105    }
106}
107
108impl Default for GroupedElGamalCiphertext3Handles {
109    fn default() -> Self {
110        Self::zeroed()
111    }
112}
113
114#[cfg(not(target_os = "solana"))]
115impl From<GroupedElGamalCiphertext<3>> for GroupedElGamalCiphertext3Handles {
116    fn from(decoded_ciphertext: GroupedElGamalCiphertext<3>) -> Self {
117        Self(decoded_ciphertext.to_bytes().try_into().unwrap())
118    }
119}
120
121#[cfg(not(target_os = "solana"))]
122impl TryFrom<GroupedElGamalCiphertext3Handles> for GroupedElGamalCiphertext<3> {
123    type Error = ElGamalError;
124
125    fn try_from(pod_ciphertext: GroupedElGamalCiphertext3Handles) -> Result<Self, Self::Error> {
126        Self::from_bytes(&pod_ciphertext.0).ok_or(ElGamalError::CiphertextDeserialization)
127    }
128}
129
130impl_extract!(TYPE = GroupedElGamalCiphertext3Handles);
131
132#[cfg(test)]
133mod tests {
134    use {
135        super::*,
136        crate::{
137            encryption::{
138                elgamal::ElGamalKeypair, grouped_elgamal::GroupedElGamal, pedersen::Pedersen,
139            },
140            zk_token_elgamal::pod::pedersen::PedersenCommitment,
141        },
142    };
143
144    #[test]
145    fn test_2_handles_ciphertext_extraction() {
146        let elgamal_keypair_0 = ElGamalKeypair::new_rand();
147        let elgamal_keypair_1 = ElGamalKeypair::new_rand();
148
149        let amount: u64 = 10;
150        let (commitment, opening) = Pedersen::new(amount);
151
152        let grouped_ciphertext = GroupedElGamal::encrypt_with(
153            [elgamal_keypair_0.pubkey(), elgamal_keypair_1.pubkey()],
154            amount,
155            &opening,
156        );
157        let pod_grouped_ciphertext: GroupedElGamalCiphertext2Handles = grouped_ciphertext.into();
158
159        let expected_pod_commitment: PedersenCommitment = commitment.into();
160        let actual_pod_commitment = pod_grouped_ciphertext.extract_commitment();
161        assert_eq!(expected_pod_commitment, actual_pod_commitment);
162
163        let expected_ciphertext_0 = elgamal_keypair_0.pubkey().encrypt_with(amount, &opening);
164        let expected_pod_ciphertext_0: ElGamalCiphertext = expected_ciphertext_0.into();
165        let actual_pod_ciphertext_0 = pod_grouped_ciphertext.try_extract_ciphertext(0).unwrap();
166        assert_eq!(expected_pod_ciphertext_0, actual_pod_ciphertext_0);
167
168        let expected_ciphertext_1 = elgamal_keypair_1.pubkey().encrypt_with(amount, &opening);
169        let expected_pod_ciphertext_1: ElGamalCiphertext = expected_ciphertext_1.into();
170        let actual_pod_ciphertext_1 = pod_grouped_ciphertext.try_extract_ciphertext(1).unwrap();
171        assert_eq!(expected_pod_ciphertext_1, actual_pod_ciphertext_1);
172
173        let err = pod_grouped_ciphertext
174            .try_extract_ciphertext(2)
175            .unwrap_err();
176        assert_eq!(err, ElGamalError::CiphertextDeserialization);
177    }
178
179    #[test]
180    fn test_3_handles_ciphertext_extraction() {
181        let elgamal_keypair_0 = ElGamalKeypair::new_rand();
182        let elgamal_keypair_1 = ElGamalKeypair::new_rand();
183        let elgamal_keypair_2 = ElGamalKeypair::new_rand();
184
185        let amount: u64 = 10;
186        let (commitment, opening) = Pedersen::new(amount);
187
188        let grouped_ciphertext = GroupedElGamal::encrypt_with(
189            [
190                elgamal_keypair_0.pubkey(),
191                elgamal_keypair_1.pubkey(),
192                elgamal_keypair_2.pubkey(),
193            ],
194            amount,
195            &opening,
196        );
197        let pod_grouped_ciphertext: GroupedElGamalCiphertext3Handles = grouped_ciphertext.into();
198
199        let expected_pod_commitment: PedersenCommitment = commitment.into();
200        let actual_pod_commitment = pod_grouped_ciphertext.extract_commitment();
201        assert_eq!(expected_pod_commitment, actual_pod_commitment);
202
203        let expected_ciphertext_0 = elgamal_keypair_0.pubkey().encrypt_with(amount, &opening);
204        let expected_pod_ciphertext_0: ElGamalCiphertext = expected_ciphertext_0.into();
205        let actual_pod_ciphertext_0 = pod_grouped_ciphertext.try_extract_ciphertext(0).unwrap();
206        assert_eq!(expected_pod_ciphertext_0, actual_pod_ciphertext_0);
207
208        let expected_ciphertext_1 = elgamal_keypair_1.pubkey().encrypt_with(amount, &opening);
209        let expected_pod_ciphertext_1: ElGamalCiphertext = expected_ciphertext_1.into();
210        let actual_pod_ciphertext_1 = pod_grouped_ciphertext.try_extract_ciphertext(1).unwrap();
211        assert_eq!(expected_pod_ciphertext_1, actual_pod_ciphertext_1);
212
213        let expected_ciphertext_2 = elgamal_keypair_2.pubkey().encrypt_with(amount, &opening);
214        let expected_pod_ciphertext_2: ElGamalCiphertext = expected_ciphertext_2.into();
215        let actual_pod_ciphertext_2 = pod_grouped_ciphertext.try_extract_ciphertext(2).unwrap();
216        assert_eq!(expected_pod_ciphertext_2, actual_pod_ciphertext_2);
217
218        let err = pod_grouped_ciphertext
219            .try_extract_ciphertext(3)
220            .unwrap_err();
221        assert_eq!(err, ElGamalError::CiphertextDeserialization);
222    }
223}