solana_zk_token_sdk/encryption/
auth_encryption.rs

1//! Authenticated encryption implementation.
2//!
3//! This module is a simple wrapper of the `Aes128GcmSiv` implementation specialized for SPL
4//! token-2022 where the plaintext is always `u64`.
5use {
6    crate::errors::AuthenticatedEncryptionError,
7    base64::{prelude::BASE64_STANDARD, Engine},
8    sha3::{Digest, Sha3_512},
9    solana_derivation_path::DerivationPath,
10    solana_seed_derivable::SeedDerivable,
11    solana_seed_phrase::generate_seed_from_seed_phrase_and_passphrase,
12    solana_signature::Signature,
13    solana_signer::{EncodableKey, Signer, SignerError},
14    std::{
15        convert::TryInto,
16        error, fmt,
17        io::{Read, Write},
18    },
19    subtle::ConstantTimeEq,
20    zeroize::Zeroize,
21};
22#[cfg(not(target_os = "solana"))]
23use {
24    aes_gcm_siv::{
25        aead::{Aead, KeyInit},
26        Aes128GcmSiv,
27    },
28    rand::{rngs::OsRng, Rng},
29};
30
31/// Byte length of an authenticated encryption secret key
32pub const AE_KEY_LEN: usize = 16;
33
34/// Byte length of an authenticated encryption nonce component
35const NONCE_LEN: usize = 12;
36
37/// Byte lenth of an authenticated encryption ciphertext component
38const CIPHERTEXT_LEN: usize = 24;
39
40/// Byte length of a complete authenticated encryption ciphertext component that includes the
41/// ciphertext and nonce components
42const AE_CIPHERTEXT_LEN: usize = 36;
43
44struct AuthenticatedEncryption;
45impl AuthenticatedEncryption {
46    /// Generates an authenticated encryption key.
47    ///
48    /// This function is randomized. It internally samples a 128-bit key using `OsRng`.
49    #[cfg(not(target_os = "solana"))]
50    fn keygen() -> AeKey {
51        AeKey(OsRng.gen::<[u8; AE_KEY_LEN]>())
52    }
53
54    /// On input of an authenticated encryption key and an amount, the function returns a
55    /// corresponding authenticated encryption ciphertext.
56    #[cfg(not(target_os = "solana"))]
57    fn encrypt(key: &AeKey, balance: u64) -> AeCiphertext {
58        let mut plaintext = balance.to_le_bytes();
59        let nonce: Nonce = OsRng.gen::<[u8; NONCE_LEN]>();
60
61        // The balance and the nonce have fixed length and therefore, encryption should not fail.
62        let ciphertext = Aes128GcmSiv::new(&key.0.into())
63            .encrypt(&nonce.into(), plaintext.as_ref())
64            .expect("authenticated encryption");
65
66        plaintext.zeroize();
67
68        AeCiphertext {
69            nonce,
70            ciphertext: ciphertext.try_into().unwrap(),
71        }
72    }
73
74    /// On input of an authenticated encryption key and a ciphertext, the function returns the
75    /// originally encrypted amount.
76    #[cfg(not(target_os = "solana"))]
77    fn decrypt(key: &AeKey, ciphertext: &AeCiphertext) -> Option<u64> {
78        let plaintext = Aes128GcmSiv::new(&key.0.into())
79            .decrypt(&ciphertext.nonce.into(), ciphertext.ciphertext.as_ref());
80
81        if let Ok(plaintext) = plaintext {
82            let amount_bytes: [u8; 8] = plaintext.try_into().unwrap();
83            Some(u64::from_le_bytes(amount_bytes))
84        } else {
85            None
86        }
87    }
88}
89
90#[derive(Debug, Zeroize, Eq, PartialEq)]
91pub struct AeKey([u8; AE_KEY_LEN]);
92impl AeKey {
93    /// Deterministically derives an authenticated encryption key from a Solana signer and a public
94    /// seed.
95    ///
96    /// This function exists for applications where a user may not wish to maintain a Solana signer
97    /// and an authenticated encryption key separately. Instead, a user can derive the ElGamal
98    /// keypair on-the-fly whenever encrytion/decryption is needed.
99    pub fn new_from_signer(
100        signer: &dyn Signer,
101        public_seed: &[u8],
102    ) -> Result<Self, Box<dyn error::Error>> {
103        let seed = Self::seed_from_signer(signer, public_seed)?;
104        Self::from_seed(&seed)
105    }
106
107    /// Derive a seed from a Solana signer used to generate an authenticated encryption key.
108    ///
109    /// The seed is derived as the hash of the signature of a public seed.
110    pub fn seed_from_signer(
111        signer: &dyn Signer,
112        public_seed: &[u8],
113    ) -> Result<Vec<u8>, SignerError> {
114        let message = [b"AeKey", public_seed].concat();
115        let signature = signer.try_sign_message(&message)?;
116
117        // Some `Signer` implementations return the default signature, which is not suitable for
118        // use as key material
119        if bool::from(signature.as_ref().ct_eq(Signature::default().as_ref())) {
120            return Err(SignerError::Custom("Rejecting default signature".into()));
121        }
122
123        let mut hasher = Sha3_512::new();
124        hasher.update(signature.as_ref());
125        let result = hasher.finalize();
126
127        Ok(result.to_vec())
128    }
129
130    /// Generates a random authenticated encryption key.
131    ///
132    /// This function is randomized. It internally samples a scalar element using `OsRng`.
133    pub fn new_rand() -> Self {
134        AuthenticatedEncryption::keygen()
135    }
136
137    /// Encrypts an amount under the authenticated encryption key.
138    pub fn encrypt(&self, amount: u64) -> AeCiphertext {
139        AuthenticatedEncryption::encrypt(self, amount)
140    }
141
142    pub fn decrypt(&self, ciphertext: &AeCiphertext) -> Option<u64> {
143        AuthenticatedEncryption::decrypt(self, ciphertext)
144    }
145}
146
147impl EncodableKey for AeKey {
148    fn read<R: Read>(reader: &mut R) -> Result<Self, Box<dyn error::Error>> {
149        let bytes: [u8; AE_KEY_LEN] = serde_json::from_reader(reader)?;
150        Ok(Self(bytes))
151    }
152
153    fn write<W: Write>(&self, writer: &mut W) -> Result<String, Box<dyn error::Error>> {
154        let bytes = self.0;
155        let json = serde_json::to_string(&bytes.to_vec())?;
156        writer.write_all(&json.clone().into_bytes())?;
157        Ok(json)
158    }
159}
160
161impl SeedDerivable for AeKey {
162    fn from_seed(seed: &[u8]) -> Result<Self, Box<dyn error::Error>> {
163        const MINIMUM_SEED_LEN: usize = AE_KEY_LEN;
164        const MAXIMUM_SEED_LEN: usize = 65535;
165
166        if seed.len() < MINIMUM_SEED_LEN {
167            return Err(AuthenticatedEncryptionError::SeedLengthTooShort.into());
168        }
169        if seed.len() > MAXIMUM_SEED_LEN {
170            return Err(AuthenticatedEncryptionError::SeedLengthTooLong.into());
171        }
172
173        let mut hasher = Sha3_512::new();
174        hasher.update(seed);
175        let result = hasher.finalize();
176
177        Ok(Self(result[..AE_KEY_LEN].try_into()?))
178    }
179
180    fn from_seed_and_derivation_path(
181        _seed: &[u8],
182        _derivation_path: Option<DerivationPath>,
183    ) -> Result<Self, Box<dyn error::Error>> {
184        Err(AuthenticatedEncryptionError::DerivationMethodNotSupported.into())
185    }
186
187    fn from_seed_phrase_and_passphrase(
188        seed_phrase: &str,
189        passphrase: &str,
190    ) -> Result<Self, Box<dyn error::Error>> {
191        Self::from_seed(&generate_seed_from_seed_phrase_and_passphrase(
192            seed_phrase,
193            passphrase,
194        ))
195    }
196}
197
198impl From<[u8; AE_KEY_LEN]> for AeKey {
199    fn from(bytes: [u8; AE_KEY_LEN]) -> Self {
200        Self(bytes)
201    }
202}
203
204impl From<AeKey> for [u8; AE_KEY_LEN] {
205    fn from(key: AeKey) -> Self {
206        key.0
207    }
208}
209
210impl TryFrom<&[u8]> for AeKey {
211    type Error = AuthenticatedEncryptionError;
212    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
213        if bytes.len() != AE_KEY_LEN {
214            return Err(AuthenticatedEncryptionError::Deserialization);
215        }
216        bytes
217            .try_into()
218            .map(Self)
219            .map_err(|_| AuthenticatedEncryptionError::Deserialization)
220    }
221}
222
223/// For the purpose of encrypting balances for the spl token accounts, the nonce and ciphertext
224/// sizes should always be fixed.
225type Nonce = [u8; NONCE_LEN];
226type Ciphertext = [u8; CIPHERTEXT_LEN];
227
228/// Authenticated encryption nonce and ciphertext
229#[derive(Debug, Default, Clone)]
230pub struct AeCiphertext {
231    nonce: Nonce,
232    ciphertext: Ciphertext,
233}
234impl AeCiphertext {
235    pub fn decrypt(&self, key: &AeKey) -> Option<u64> {
236        AuthenticatedEncryption::decrypt(key, self)
237    }
238
239    pub fn to_bytes(&self) -> [u8; AE_CIPHERTEXT_LEN] {
240        let mut buf = [0_u8; AE_CIPHERTEXT_LEN];
241        buf[..NONCE_LEN].copy_from_slice(&self.nonce);
242        buf[NONCE_LEN..].copy_from_slice(&self.ciphertext);
243        buf
244    }
245
246    pub fn from_bytes(bytes: &[u8]) -> Option<AeCiphertext> {
247        if bytes.len() != AE_CIPHERTEXT_LEN {
248            return None;
249        }
250
251        let nonce = bytes[..NONCE_LEN].try_into().ok()?;
252        let ciphertext = bytes[NONCE_LEN..].try_into().ok()?;
253
254        Some(AeCiphertext { nonce, ciphertext })
255    }
256}
257
258impl fmt::Display for AeCiphertext {
259    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
260        write!(f, "{}", BASE64_STANDARD.encode(self.to_bytes()))
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use {
267        super::*, solana_keypair::Keypair, solana_pubkey::Pubkey,
268        solana_signer::null_signer::NullSigner,
269    };
270
271    #[test]
272    fn test_aes_encrypt_decrypt_correctness() {
273        let key = AeKey::new_rand();
274        let amount = 55;
275
276        let ciphertext = key.encrypt(amount);
277        let decrypted_amount = ciphertext.decrypt(&key).unwrap();
278
279        assert_eq!(amount, decrypted_amount);
280    }
281
282    #[test]
283    fn test_aes_new() {
284        let keypair1 = Keypair::new();
285        let keypair2 = Keypair::new();
286
287        assert_ne!(
288            AeKey::new_from_signer(&keypair1, Pubkey::default().as_ref())
289                .unwrap()
290                .0,
291            AeKey::new_from_signer(&keypair2, Pubkey::default().as_ref())
292                .unwrap()
293                .0,
294        );
295
296        let null_signer = NullSigner::new(&Pubkey::default());
297        assert!(AeKey::new_from_signer(&null_signer, Pubkey::default().as_ref()).is_err());
298    }
299
300    #[test]
301    fn test_aes_key_from_seed() {
302        let good_seed = vec![0; 32];
303        assert!(AeKey::from_seed(&good_seed).is_ok());
304
305        let too_short_seed = vec![0; 15];
306        assert!(AeKey::from_seed(&too_short_seed).is_err());
307
308        let too_long_seed = vec![0; 65536];
309        assert!(AeKey::from_seed(&too_long_seed).is_err());
310    }
311
312    #[test]
313    fn test_aes_key_from() {
314        let key = AeKey::from_seed(&[0; 32]).unwrap();
315        let key_bytes: [u8; AE_KEY_LEN] = AeKey::from_seed(&[0; 32]).unwrap().into();
316
317        assert_eq!(key, AeKey::from(key_bytes));
318    }
319
320    #[test]
321    fn test_aes_key_try_from() {
322        let key = AeKey::from_seed(&[0; 32]).unwrap();
323        let key_bytes: [u8; AE_KEY_LEN] = AeKey::from_seed(&[0; 32]).unwrap().into();
324
325        assert_eq!(key, AeKey::try_from(key_bytes.as_slice()).unwrap());
326    }
327
328    #[test]
329    fn test_aes_key_try_from_error() {
330        let too_many_bytes = vec![0_u8; 32];
331        assert!(AeKey::try_from(too_many_bytes.as_slice()).is_err());
332    }
333}
solana_zk_token_sdk/encryption/auth_encryption.rs

solana_zk_token_sdk/encryption/
auth_encryption.rs
