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

defuse_core/nonce/
salted.rs
