safe_zk_token_sdk/encryption/
auth_encryption.rs

1//! Authenticated encryption implementation.
2//!
3//! This module is a simple wrapper of the `Aes128GcmSiv` implementation.
4#[cfg(not(target_os = "solana"))]
5use {
6    aes_gcm_siv::{
7        aead::{Aead, NewAead},
8        Aes128GcmSiv,
9    },
10    rand::{rngs::OsRng, CryptoRng, Rng, RngCore},
11};
12use {
13    arrayref::{array_ref, array_refs},
14    solana_sdk::{
15        instruction::Instruction,
16        message::Message,
17        pubkey::Pubkey,
18        signature::Signature,
19        signer::{Signer, SignerError},
20    },
21    std::{convert::TryInto, fmt},
22    subtle::ConstantTimeEq,
23    zeroize::Zeroize,
24};
25
26struct AuthenticatedEncryption;
27impl AuthenticatedEncryption {
28    #[cfg(not(target_os = "solana"))]
29    #[allow(clippy::new_ret_no_self)]
30    fn keygen<T: RngCore + CryptoRng>(rng: &mut T) -> AeKey {
31        AeKey(rng.gen::<[u8; 16]>())
32    }
33
34    #[cfg(not(target_os = "solana"))]
35    fn encrypt(key: &AeKey, balance: u64) -> AeCiphertext {
36        let mut plaintext = balance.to_le_bytes();
37        let nonce: Nonce = OsRng.gen::<[u8; 12]>();
38
39        // The balance and the nonce have fixed length and therefore, encryption should not fail.
40        let ciphertext = Aes128GcmSiv::new(&key.0.into())
41            .encrypt(&nonce.into(), plaintext.as_ref())
42            .expect("authenticated encryption");
43
44        plaintext.zeroize();
45
46        AeCiphertext {
47            nonce,
48            ciphertext: ciphertext.try_into().unwrap(),
49        }
50    }
51
52    #[cfg(not(target_os = "solana"))]
53    fn decrypt(key: &AeKey, ct: &AeCiphertext) -> Option<u64> {
54        let plaintext =
55            Aes128GcmSiv::new(&key.0.into()).decrypt(&ct.nonce.into(), ct.ciphertext.as_ref());
56
57        if let Ok(plaintext) = plaintext {
58            let amount_bytes: [u8; 8] = plaintext.try_into().unwrap();
59            Some(u64::from_le_bytes(amount_bytes))
60        } else {
61            None
62        }
63    }
64}
65
66#[derive(Debug, Zeroize)]
67pub struct AeKey([u8; 16]);
68impl AeKey {
69    pub fn new(signer: &dyn Signer, address: &Pubkey) -> Result<Self, SignerError> {
70        let message = Message::new(
71            &[Instruction::new_with_bytes(*address, b"AeKey", vec![])],
72            Some(&signer.try_pubkey()?),
73        );
74        let signature = signer.try_sign_message(&message.serialize())?;
75
76        // Some `Signer` implementations return the default signature, which is not suitable for
77        // use as key material
78        if bool::from(signature.as_ref().ct_eq(Signature::default().as_ref())) {
79            Err(SignerError::Custom("Rejecting default signature".into()))
80        } else {
81            Ok(AeKey(signature.as_ref()[..16].try_into().unwrap()))
82        }
83    }
84
85    pub fn random<T: RngCore + CryptoRng>(rng: &mut T) -> Self {
86        AuthenticatedEncryption::keygen(rng)
87    }
88
89    pub fn encrypt(&self, amount: u64) -> AeCiphertext {
90        AuthenticatedEncryption::encrypt(self, amount)
91    }
92
93    pub fn decrypt(&self, ct: &AeCiphertext) -> Option<u64> {
94        AuthenticatedEncryption::decrypt(self, ct)
95    }
96}
97
98/// For the purpose of encrypting balances for the spl token accounts, the nonce and ciphertext
99/// sizes should always be fixed.
100pub type Nonce = [u8; 12];
101pub type Ciphertext = [u8; 24];
102
103/// Authenticated encryption nonce and ciphertext
104#[derive(Debug, Default, Clone)]
105pub struct AeCiphertext {
106    pub nonce: Nonce,
107    pub ciphertext: Ciphertext,
108}
109impl AeCiphertext {
110    pub fn decrypt(&self, key: &AeKey) -> Option<u64> {
111        AuthenticatedEncryption::decrypt(key, self)
112    }
113
114    pub fn to_bytes(&self) -> [u8; 36] {
115        let mut buf = [0_u8; 36];
116        buf[..12].copy_from_slice(&self.nonce);
117        buf[12..].copy_from_slice(&self.ciphertext);
118        buf
119    }
120
121    pub fn from_bytes(bytes: &[u8]) -> Option<AeCiphertext> {
122        if bytes.len() != 36 {
123            return None;
124        }
125
126        let bytes = array_ref![bytes, 0, 36];
127        let (nonce, ciphertext) = array_refs![bytes, 12, 24];
128
129        Some(AeCiphertext {
130            nonce: *nonce,
131            ciphertext: *ciphertext,
132        })
133    }
134}
135
136impl fmt::Display for AeCiphertext {
137    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
138        write!(f, "{}", base64::encode(self.to_bytes()))
139    }
140}
141
142#[cfg(test)]
143mod tests {
144    use {
145        super::*,
146        solana_sdk::{signature::Keypair, signer::null_signer::NullSigner},
147    };
148
149    #[test]
150    fn test_aes_encrypt_decrypt_correctness() {
151        let key = AeKey::random(&mut OsRng);
152        let amount = 55;
153
154        let ct = key.encrypt(amount);
155        let decrypted_amount = ct.decrypt(&key).unwrap();
156
157        assert_eq!(amount, decrypted_amount);
158    }
159
160    #[test]
161    fn test_aes_new() {
162        let keypair1 = Keypair::new();
163        let keypair2 = Keypair::new();
164
165        assert_ne!(
166            AeKey::new(&keypair1, &Pubkey::default()).unwrap().0,
167            AeKey::new(&keypair2, &Pubkey::default()).unwrap().0,
168        );
169
170        let null_signer = NullSigner::new(&Pubkey::default());
171        assert!(AeKey::new(&null_signer, &Pubkey::default()).is_err());
172    }
173}
safe_zk_token_sdk/encryption/auth_encryption.rs

safe_zk_token_sdk/encryption/
auth_encryption.rs
