light_compressed_token/
token_data.rs

1use std::vec;
2
3use anchor_lang::{
4    prelude::borsh, solana_program::pubkey::Pubkey, AnchorDeserialize, AnchorSerialize,
5};
6use light_compressed_account::hash_to_bn254_field_size_be;
7use light_hasher::{errors::HasherError, Hasher, Poseidon};
8
9#[derive(Clone, Copy, Debug, PartialEq, Eq, AnchorSerialize, AnchorDeserialize)]
10#[repr(u8)]
11pub enum AccountState {
12    Initialized,
13    Frozen,
14}
15
16#[derive(Debug, PartialEq, Eq, AnchorSerialize, AnchorDeserialize, Clone)]
17pub struct TokenData {
18    /// The mint associated with this account
19    pub mint: Pubkey,
20    /// The owner of this account.
21    pub owner: Pubkey,
22    /// The amount of tokens this account holds.
23    pub amount: u64,
24    /// If `delegate` is `Some` then `delegated_amount` represents
25    /// the amount authorized by the delegate
26    pub delegate: Option<Pubkey>,
27    /// The account's state
28    pub state: AccountState,
29    /// Placeholder for TokenExtension tlv data (unimplemented)
30    pub tlv: Option<Vec<u8>>,
31}
32
33/// Hashing schema: H(mint, owner, amount, delegate, delegated_amount,
34/// is_native, state)
35///
36/// delegate, delegated_amount, is_native and state have dynamic positions.
37/// Always hash mint, owner and amount If delegate hash delegate and
38/// delegated_amount together. If is native hash is_native else is omitted.
39/// If frozen hash AccountState::Frozen else is omitted.
40///
41/// Security: to prevent the possibility that different fields with the same
42/// value to result in the same hash we add a prefix to the delegated amount, is
43/// native and state fields. This way we can have a dynamic hashing schema and
44/// hash only used values.
45impl TokenData {
46    /// Only the spl representation of native tokens (wrapped SOL) is
47    /// compressed.
48    /// The sol value is stored in the token pool account.
49    /// The sol value in the compressed account is independent from
50    /// the wrapped sol amount.
51    pub fn is_native(&self) -> bool {
52        self.mint == spl_token::native_mint::id()
53    }
54    pub fn hash_with_hashed_values(
55        hashed_mint: &[u8; 32],
56        hashed_owner: &[u8; 32],
57        amount_bytes: &[u8; 32],
58        hashed_delegate: &Option<&[u8; 32]>,
59    ) -> std::result::Result<[u8; 32], HasherError> {
60        Self::hash_inputs_with_hashed_values::<false>(
61            hashed_mint,
62            hashed_owner,
63            amount_bytes,
64            hashed_delegate,
65        )
66    }
67
68    pub fn hash_frozen_with_hashed_values(
69        hashed_mint: &[u8; 32],
70        hashed_owner: &[u8; 32],
71        amount_bytes: &[u8; 32],
72        hashed_delegate: &Option<&[u8; 32]>,
73    ) -> std::result::Result<[u8; 32], HasherError> {
74        Self::hash_inputs_with_hashed_values::<true>(
75            hashed_mint,
76            hashed_owner,
77            amount_bytes,
78            hashed_delegate,
79        )
80    }
81
82    /// We should not hash pubkeys multiple times. For all we can assume mints
83    /// are equal. For all input compressed accounts we assume owners are
84    /// equal.
85    pub fn hash_inputs_with_hashed_values<const FROZEN_INPUTS: bool>(
86        mint: &[u8; 32],
87        owner: &[u8; 32],
88        amount_bytes: &[u8],
89        hashed_delegate: &Option<&[u8; 32]>,
90    ) -> std::result::Result<[u8; 32], HasherError> {
91        let mut hash_inputs = vec![mint.as_slice(), owner.as_slice(), amount_bytes];
92        if let Some(hashed_delegate) = hashed_delegate {
93            hash_inputs.push(hashed_delegate.as_slice());
94        }
95        let mut state_bytes = [0u8; 32];
96        if FROZEN_INPUTS {
97            state_bytes[31] = AccountState::Frozen as u8;
98            hash_inputs.push(&state_bytes[..]);
99        }
100        Poseidon::hashv(hash_inputs.as_slice())
101    }
102}
103
104impl TokenData {
105    /// Hashes token data of token accounts.
106    ///
107    /// Note, hashing changed for token account data in batched Merkle trees.
108    /// For hashing of token account data stored in concurrent Merkle trees use hash_legacy().
109    pub fn hash(&self) -> std::result::Result<[u8; 32], HasherError> {
110        self._hash::<true>()
111    }
112
113    /// Hashes token data of token accounts stored in concurrent Merkle trees.
114    pub fn hash_legacy(&self) -> std::result::Result<[u8; 32], HasherError> {
115        self._hash::<false>()
116    }
117
118    fn _hash<const BATCHED: bool>(&self) -> std::result::Result<[u8; 32], HasherError> {
119        let hashed_mint = hash_to_bn254_field_size_be(self.mint.to_bytes().as_slice());
120        let hashed_owner = hash_to_bn254_field_size_be(self.owner.to_bytes().as_slice());
121        let mut amount_bytes = [0u8; 32];
122        if BATCHED {
123            amount_bytes[24..].copy_from_slice(self.amount.to_be_bytes().as_slice());
124        } else {
125            amount_bytes[24..].copy_from_slice(self.amount.to_le_bytes().as_slice());
126        }
127        let hashed_delegate;
128        let hashed_delegate_option = if let Some(delegate) = self.delegate {
129            hashed_delegate = hash_to_bn254_field_size_be(delegate.to_bytes().as_slice());
130            Some(&hashed_delegate)
131        } else {
132            None
133        };
134        if self.state != AccountState::Initialized {
135            Self::hash_inputs_with_hashed_values::<true>(
136                &hashed_mint,
137                &hashed_owner,
138                &amount_bytes,
139                &hashed_delegate_option,
140            )
141        } else {
142            Self::hash_inputs_with_hashed_values::<false>(
143                &hashed_mint,
144                &hashed_owner,
145                &amount_bytes,
146                &hashed_delegate_option,
147            )
148        }
149    }
150}
151
152#[cfg(test)]
153pub mod test {
154
155    use num_bigint::BigUint;
156    use rand::Rng;
157
158    use super::*;
159
160    #[test]
161    fn equivalency_of_hash_functions() {
162        let token_data = TokenData {
163            mint: Pubkey::new_unique(),
164            owner: Pubkey::new_unique(),
165            amount: 100,
166            delegate: Some(Pubkey::new_unique()),
167            state: AccountState::Initialized,
168            tlv: None,
169        };
170        let hashed_token_data = token_data.hash_legacy().unwrap();
171        let hashed_mint = hash_to_bn254_field_size_be(token_data.mint.to_bytes().as_slice());
172        let hashed_owner = hash_to_bn254_field_size_be(token_data.owner.to_bytes().as_slice());
173        let hashed_delegate =
174            hash_to_bn254_field_size_be(token_data.delegate.unwrap().to_bytes().as_slice());
175        let mut amount_bytes = [0u8; 32];
176        amount_bytes[24..].copy_from_slice(token_data.amount.to_le_bytes().as_slice());
177        let hashed_token_data_with_hashed_values =
178            TokenData::hash_inputs_with_hashed_values::<false>(
179                &hashed_mint,
180                &hashed_owner,
181                &amount_bytes,
182                &Some(&hashed_delegate),
183            )
184            .unwrap();
185        assert_eq!(hashed_token_data, hashed_token_data_with_hashed_values);
186
187        let token_data = TokenData {
188            mint: Pubkey::new_unique(),
189            owner: Pubkey::new_unique(),
190            amount: 101,
191            delegate: None,
192            state: AccountState::Initialized,
193            tlv: None,
194        };
195        let hashed_token_data = token_data.hash_legacy().unwrap();
196        let hashed_mint = hash_to_bn254_field_size_be(token_data.mint.to_bytes().as_slice());
197        let hashed_owner = hash_to_bn254_field_size_be(token_data.owner.to_bytes().as_slice());
198        let mut amount_bytes = [0u8; 32];
199        amount_bytes[24..].copy_from_slice(token_data.amount.to_le_bytes().as_slice());
200        let hashed_token_data_with_hashed_values =
201            TokenData::hash_with_hashed_values(&hashed_mint, &hashed_owner, &amount_bytes, &None)
202                .unwrap();
203        assert_eq!(hashed_token_data, hashed_token_data_with_hashed_values);
204    }
205
206    impl TokenData {
207        fn legacy_hash(&self) -> std::result::Result<[u8; 32], HasherError> {
208            let hashed_mint = hash_to_bn254_field_size_be(self.mint.to_bytes().as_slice());
209            let hashed_owner = hash_to_bn254_field_size_be(self.owner.to_bytes().as_slice());
210            let amount_bytes = self.amount.to_le_bytes();
211            let hashed_delegate;
212            let hashed_delegate_option = if let Some(delegate) = self.delegate {
213                hashed_delegate = hash_to_bn254_field_size_be(delegate.to_bytes().as_slice());
214                Some(&hashed_delegate)
215            } else {
216                None
217            };
218            if self.state != AccountState::Initialized {
219                Self::hash_inputs_with_hashed_values::<true>(
220                    &hashed_mint,
221                    &hashed_owner,
222                    &amount_bytes,
223                    &hashed_delegate_option,
224                )
225            } else {
226                Self::hash_inputs_with_hashed_values::<false>(
227                    &hashed_mint,
228                    &hashed_owner,
229                    &amount_bytes,
230                    &hashed_delegate_option,
231                )
232            }
233        }
234    }
235    fn equivalency_of_hash_functions_rnd_iters<const ITERS: usize>() {
236        let mut rng = rand::thread_rng();
237
238        for _ in 0..ITERS {
239            let token_data = TokenData {
240                mint: Pubkey::new_unique(),
241                owner: Pubkey::new_unique(),
242                amount: rng.gen(),
243                delegate: Some(Pubkey::new_unique()),
244                state: AccountState::Initialized,
245                tlv: None,
246            };
247            let hashed_token_data = token_data.hash_legacy().unwrap();
248            let hashed_mint = hash_to_bn254_field_size_be(token_data.mint.to_bytes().as_slice());
249            let hashed_owner = hash_to_bn254_field_size_be(token_data.owner.to_bytes().as_slice());
250            let hashed_delegate =
251                hash_to_bn254_field_size_be(token_data.delegate.unwrap().to_bytes().as_slice());
252            let mut amount_bytes = [0u8; 32];
253            amount_bytes[24..].copy_from_slice(token_data.amount.to_le_bytes().as_slice());
254            let hashed_token_data_with_hashed_values = TokenData::hash_with_hashed_values(
255                &hashed_mint,
256                &hashed_owner,
257                &amount_bytes,
258                &Some(&hashed_delegate),
259            )
260            .unwrap();
261            assert_eq!(hashed_token_data, hashed_token_data_with_hashed_values);
262            let legacy_hash = token_data.legacy_hash().unwrap();
263            assert_eq!(hashed_token_data, legacy_hash);
264
265            let token_data = TokenData {
266                mint: Pubkey::new_unique(),
267                owner: Pubkey::new_unique(),
268                amount: rng.gen(),
269                delegate: None,
270                state: AccountState::Initialized,
271                tlv: None,
272            };
273            let hashed_token_data = token_data.hash_legacy().unwrap();
274            let hashed_mint = hash_to_bn254_field_size_be(token_data.mint.to_bytes().as_slice());
275            let hashed_owner = hash_to_bn254_field_size_be(token_data.owner.to_bytes().as_slice());
276            let mut amount_bytes = [0u8; 32];
277            amount_bytes[24..].copy_from_slice(token_data.amount.to_le_bytes().as_slice());
278            let hashed_token_data_with_hashed_values: [u8; 32] =
279                TokenData::hash_with_hashed_values(
280                    &hashed_mint,
281                    &hashed_owner,
282                    &amount_bytes,
283                    &None,
284                )
285                .unwrap();
286            assert_eq!(hashed_token_data, hashed_token_data_with_hashed_values);
287            let legacy_hash = token_data.legacy_hash().unwrap();
288            assert_eq!(hashed_token_data, legacy_hash);
289        }
290    }
291
292    #[test]
293    fn equivalency_of_hash_functions_iters_poseidon() {
294        equivalency_of_hash_functions_rnd_iters::<10_000>();
295    }
296
297    #[test]
298    fn test_circuit_equivalence() {
299        // Convert hex strings to Pubkeys
300        let mint_pubkey = Pubkey::new_from_array([
301            0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
302            0, 0, 0,
303        ]);
304        let owner_pubkey = Pubkey::new_from_array([
305            0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
306            0, 0, 0,
307        ]);
308        let delegate_pubkey = Pubkey::new_from_array([
309            0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
310            0, 0, 0,
311        ]);
312
313        let token_data = TokenData {
314            mint: mint_pubkey,
315            owner: owner_pubkey,
316            amount: 1000000u64,
317            delegate: Some(delegate_pubkey),
318            state: AccountState::Initialized, // Using Frozen state to match our circuit test
319            tlv: None,
320        };
321
322        // Calculate the hash with the Rust code
323        let rust_hash = token_data.hash().unwrap();
324
325        let circuit_hash_str =
326            "12698830169693734517877055378728747723888091986541703429186543307137690361131";
327        use std::str::FromStr;
328        let circuit_hash = BigUint::from_str(circuit_hash_str).unwrap().to_bytes_be();
329        let rust_hash_string = BigUint::from_bytes_be(rust_hash.as_slice()).to_string();
330        println!("Circuit hash string: {}", circuit_hash_str);
331        println!("rust_hash_string {}", rust_hash_string);
332        assert_eq!(rust_hash.to_vec(), circuit_hash);
333    }
334
335    #[test]
336    fn test_frozen_equivalence() {
337        let token_data = TokenData {
338            mint: Pubkey::new_unique(),
339            owner: Pubkey::new_unique(),
340            amount: 100,
341            delegate: Some(Pubkey::new_unique()),
342            state: AccountState::Initialized,
343            tlv: None,
344        };
345        let hashed_mint = hash_to_bn254_field_size_be(token_data.mint.to_bytes().as_slice());
346        let hashed_owner = hash_to_bn254_field_size_be(token_data.owner.to_bytes().as_slice());
347        let hashed_delegate =
348            hash_to_bn254_field_size_be(token_data.delegate.unwrap().to_bytes().as_slice());
349        let mut amount_bytes = [0u8; 32];
350        amount_bytes[24..].copy_from_slice(token_data.amount.to_le_bytes().as_slice());
351        let hash = TokenData::hash_with_hashed_values(
352            &hashed_mint,
353            &hashed_owner,
354            &amount_bytes,
355            &Some(&hashed_delegate),
356        )
357        .unwrap();
358        let other_hash = token_data.hash_legacy().unwrap();
359        assert_eq!(hash, other_hash);
360    }
361
362    #[test]
363    fn failing_tests_hashing() {
364        let mut vec_previous_hashes = Vec::new();
365        let token_data = TokenData {
366            mint: Pubkey::new_unique(),
367            owner: Pubkey::new_unique(),
368            amount: 100,
369            delegate: None,
370            state: AccountState::Initialized,
371            tlv: None,
372        };
373        let hashed_mint = hash_to_bn254_field_size_be(token_data.mint.to_bytes().as_slice());
374        let hashed_owner = hash_to_bn254_field_size_be(token_data.owner.to_bytes().as_slice());
375        let mut amount_bytes = [0u8; 32];
376        amount_bytes[24..].copy_from_slice(token_data.amount.to_le_bytes().as_slice());
377        let hash =
378            TokenData::hash_with_hashed_values(&hashed_mint, &hashed_owner, &amount_bytes, &None)
379                .unwrap();
380        vec_previous_hashes.push(hash);
381        // different mint
382        let hashed_mint_2 = hash_to_bn254_field_size_be(Pubkey::new_unique().to_bytes().as_slice());
383        let mut amount_bytes = [0u8; 32];
384        amount_bytes[24..].copy_from_slice(token_data.amount.to_le_bytes().as_slice());
385        let hash2 =
386            TokenData::hash_with_hashed_values(&hashed_mint_2, &hashed_owner, &amount_bytes, &None)
387                .unwrap();
388        assert_to_previous_hashes(hash2, &mut vec_previous_hashes);
389
390        // different owner
391        let hashed_owner_2 =
392            hash_to_bn254_field_size_be(Pubkey::new_unique().to_bytes().as_slice());
393        let mut amount_bytes = [0u8; 32];
394        amount_bytes[24..].copy_from_slice(token_data.amount.to_le_bytes().as_slice());
395        let hash3 =
396            TokenData::hash_with_hashed_values(&hashed_mint, &hashed_owner_2, &amount_bytes, &None)
397                .unwrap();
398        assert_to_previous_hashes(hash3, &mut vec_previous_hashes);
399
400        // different amount
401        let different_amount: u64 = 101;
402        let mut different_amount_bytes = [0u8; 32];
403        different_amount_bytes[24..].copy_from_slice(different_amount.to_le_bytes().as_slice());
404        let hash4 = TokenData::hash_with_hashed_values(
405            &hashed_mint,
406            &hashed_owner,
407            &different_amount_bytes,
408            &None,
409        )
410        .unwrap();
411        assert_to_previous_hashes(hash4, &mut vec_previous_hashes);
412
413        // different delegate
414        let delegate = Pubkey::new_unique();
415        let hashed_delegate = hash_to_bn254_field_size_be(delegate.to_bytes().as_slice());
416        let mut amount_bytes = [0u8; 32];
417        amount_bytes[24..].copy_from_slice(token_data.amount.to_le_bytes().as_slice());
418        let hash7 = TokenData::hash_with_hashed_values(
419            &hashed_mint,
420            &hashed_owner,
421            &amount_bytes,
422            &Some(&hashed_delegate),
423        )
424        .unwrap();
425
426        assert_to_previous_hashes(hash7, &mut vec_previous_hashes);
427        // different account state
428        let mut token_data = token_data;
429        token_data.state = AccountState::Frozen;
430        let hash9 = token_data.hash_legacy().unwrap();
431        assert_to_previous_hashes(hash9, &mut vec_previous_hashes);
432        // different account state with delegate
433        token_data.delegate = Some(delegate);
434        let hash10 = token_data.hash_legacy().unwrap();
435        assert_to_previous_hashes(hash10, &mut vec_previous_hashes);
436    }
437
438    fn assert_to_previous_hashes(hash: [u8; 32], previous_hashes: &mut Vec<[u8; 32]>) {
439        for previous_hash in previous_hashes.iter() {
440            assert_ne!(hash, *previous_hash);
441        }
442        println!("len previous hashes: {}", previous_hashes.len());
443        previous_hashes.push(hash);
444    }
445}