defuse_core/nonce/
salted.rs

1use core::mem;
2use hex::FromHex;
3use near_sdk::{
4    IntoStorageKey,
5    borsh::{BorshDeserialize, BorshSerialize},
6    env::{self, sha256_array},
7    near,
8    store::{IterableMap, key::Identity},
9};
10use serde_with::{DeserializeFromStr, SerializeDisplay};
11use std::{
12    fmt::{self, Debug},
13    str::FromStr,
14};
15
16use crate::{DefuseError, Result};
17
18#[cfg_attr(any(feature = "arbitrary", test), derive(arbitrary::Arbitrary))]
19#[derive(PartialEq, PartialOrd, Ord, Eq, Copy, Clone, SerializeDisplay, DeserializeFromStr)]
20#[near(serializers = [borsh])]
21pub struct Salt([u8; 4]);
22
23impl Salt {
24    pub fn derive(num: u8) -> Self {
25        const SIZE: usize = size_of::<Salt>();
26
27        let seed = env::random_seed_array();
28        let mut input = [0u8; 33];
29        input[..32].copy_from_slice(&seed);
30        input[32] = num;
31
32        Self(
33            sha256_array(input)[..SIZE]
34                .try_into()
35                .unwrap_or_else(|_| unreachable!()),
36        )
37    }
38}
39
40impl fmt::Debug for Salt {
41    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
42        write!(f, "{}", hex::encode(self.0))
43    }
44}
45
46impl fmt::Display for Salt {
47    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
48        Debug::fmt(self, f)
49    }
50}
51
52impl FromStr for Salt {
53    type Err = hex::FromHexError;
54
55    fn from_str(s: &str) -> Result<Self, Self::Err> {
56        FromHex::from_hex(s).map(Self)
57    }
58}
59
60#[cfg(feature = "abi")]
61const _: () = {
62    use near_sdk::schemars::{
63        JsonSchema,
64        r#gen::SchemaGenerator,
65        schema::{InstanceType, Schema, SchemaObject},
66    };
67
68    impl JsonSchema for Salt {
69        fn schema_name() -> String {
70            String::schema_name()
71        }
72
73        fn is_referenceable() -> bool {
74            false
75        }
76
77        fn json_schema(_gen: &mut SchemaGenerator) -> Schema {
78            SchemaObject {
79                instance_type: Some(InstanceType::String.into()),
80                extensions: std::iter::once(("contentEncoding", "hex".into()))
81                    .map(|(k, v)| (k.to_string(), v))
82                    .collect(),
83                ..Default::default()
84            }
85            .into()
86        }
87    }
88};
89
90/// Contains current valid salt and set of previous
91/// salts that can be valid or invalid.
92#[near(serializers = [borsh])]
93#[derive(Debug)]
94pub struct SaltRegistry {
95    previous: IterableMap<Salt, bool, Identity>,
96    current: Salt,
97}
98
99impl SaltRegistry {
100    /// There can be only one valid salt at the beginning
101    #[inline]
102    pub fn new<S>(prefix: S) -> Self
103    where
104        S: IntoStorageKey,
105    {
106        Self {
107            previous: IterableMap::with_hasher(prefix),
108            current: Salt::derive(0),
109        }
110    }
111
112    fn derive_next_salt(&self) -> Result<Salt> {
113        (0..=u8::MAX)
114            .map(Salt::derive)
115            .find(|s| !self.is_used(*s))
116            .ok_or(DefuseError::SaltGenerationFailed)
117    }
118
119    /// Rotates the current salt, making it previous and keeping it valid.
120    #[inline]
121    pub fn set_new(&mut self) -> Result<Salt> {
122        let salt = self.derive_next_salt()?;
123
124        let previous = mem::replace(&mut self.current, salt);
125        self.previous.insert(previous, true);
126
127        Ok(previous)
128    }
129
130    /// Deactivates the previous salt, making it invalid.
131    #[inline]
132    pub fn invalidate(&mut self, salt: Salt) -> Result<()> {
133        if salt == self.current {
134            self.set_new()?;
135        }
136
137        self.previous
138            .get_mut(&salt)
139            .map(|v| *v = false)
140            .ok_or(DefuseError::InvalidSalt)
141    }
142
143    #[inline]
144    pub fn is_valid(&self, salt: Salt) -> bool {
145        salt == self.current || self.previous.get(&salt).is_some_and(|v| *v)
146    }
147
148    #[inline]
149    fn is_used(&self, salt: Salt) -> bool {
150        salt == self.current || self.previous.contains_key(&salt)
151    }
152
153    #[inline]
154    pub const fn current(&self) -> Salt {
155        self.current
156    }
157}
158
159#[derive(Clone, Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
160#[borsh(crate = "::near_sdk::borsh")]
161pub struct SaltedNonce<T>
162where
163    T: BorshSerialize + BorshDeserialize,
164{
165    pub salt: Salt,
166    pub nonce: T,
167}
168
169impl<T> SaltedNonce<T>
170where
171    T: BorshSerialize + BorshDeserialize,
172{
173    pub const fn new(salt: Salt, nonce: T) -> Self {
174        Self { salt, nonce }
175    }
176}
177
178#[cfg(test)]
179mod tests {
180    use super::*;
181
182    use arbitrary::Unstructured;
183    use defuse_test_utils::random::{Rng, RngExt, random_bytes, rng};
184    use near_sdk::{test_utils::VMContextBuilder, testing_env};
185
186    use rstest::rstest;
187
188    impl From<&[u8]> for Salt {
189        fn from(value: &[u8]) -> Self {
190            let mut result = [0u8; 4];
191            result.copy_from_slice(&value[..4]);
192            Self(result)
193        }
194    }
195
196    fn seed_to_salt(seed: &[u8; 32], attempts: u8) -> Salt {
197        let seed = [seed, attempts.to_be_bytes().as_ref()].concat();
198        let hash = sha256_array(&seed);
199
200        hash[..4].into()
201    }
202
203    fn set_random_seed(rng: &mut impl Rng) -> [u8; 32] {
204        let seed = rng.random();
205        let context = VMContextBuilder::new().random_seed(seed).build();
206        testing_env!(context);
207
208        seed
209    }
210
211    #[rstest]
212    fn contains_salt_test(random_bytes: Vec<u8>) {
213        let random_salt: Salt = Unstructured::new(&random_bytes).arbitrary().unwrap();
214        let salts = SaltRegistry::new(random_bytes);
215
216        assert!(salts.is_valid(salts.current));
217        assert!(!salts.is_valid(random_salt));
218    }
219
220    #[rstest]
221    fn update_current_salt_test(random_bytes: Vec<u8>, mut rng: impl Rng) {
222        let mut salts = SaltRegistry::new(random_bytes);
223
224        let seed = set_random_seed(&mut rng);
225        let previous_salt = salts.set_new().expect("should set new salt");
226
227        assert!(salts.is_valid(seed_to_salt(&seed, 0)));
228        assert!(salts.is_valid(previous_salt));
229
230        let previous_salt = salts.set_new().expect("should set new salt");
231        assert!(salts.is_valid(seed_to_salt(&seed, 1)));
232        assert!(salts.is_valid(previous_salt));
233    }
234
235    #[rstest]
236    fn reset_salt_test(random_bytes: Vec<u8>, mut rng: impl Rng) {
237        let mut salts = SaltRegistry::new(random_bytes);
238        let random_salt = rng.random::<[u8; 4]>().as_slice().into();
239
240        let seed = set_random_seed(&mut rng);
241        let current = seed_to_salt(&seed, 0);
242        let previous_salt = salts.set_new().expect("should set new salt");
243
244        assert!(salts.invalidate(previous_salt).is_ok());
245        assert!(!salts.is_valid(previous_salt));
246        assert!(matches!(
247            salts.invalidate(random_salt).unwrap_err(),
248            DefuseError::InvalidSalt
249        ));
250
251        let seed = set_random_seed(&mut rng);
252        let new_salt = seed_to_salt(&seed, 0);
253
254        assert!(salts.invalidate(current).is_ok());
255        assert!(!salts.is_valid(current));
256        assert_eq!(salts.current(), new_salt);
257    }
258
259    #[rstest]
260    fn derive_next_test(random_bytes: Vec<u8>) {
261        let mut salt_registry = SaltRegistry::new(random_bytes);
262
263        let prev = salt_registry.set_new().unwrap();
264
265        salt_registry.invalidate(prev).unwrap();
266        salt_registry.set_new().unwrap();
267
268        assert!(!salt_registry.is_valid(prev));
269        assert!(salt_registry.is_used(prev));
270    }
271}
defuse_core/nonce/salted.rs

defuse_core/nonce/
salted.rs
