bobcat_storage/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2
3use keccak_const::Keccak256;
4
5use array_concat::concat_arrays;
6
7pub use bobcat_maths::U;
8
9#[cfg(all(target_family = "wasm", target_os = "unknown"))]
10#[link(wasm_import_module = "vm_hooks")]
11unsafe extern "C" {
12    fn storage_load_bytes32(key: *const u8, out: *mut u8);
13    fn storage_cache_bytes32(key: *const u8, from: *const u8);
14    fn transient_load_bytes32(key: *const u8, dest: *mut u8);
15    fn transient_store_bytes32(key: *const u8, value: *const u8);
16    fn native_keccak256(bytes: *const u8, len: usize, output: *mut u8);
17    pub fn storage_flush_cache(clear: bool);
18}
19
20#[cfg(all(
21    not(all(target_family = "wasm", target_os = "unknown")),
22    feature = "std"
23))]
24pub mod storage_host {
25    use super::*;
26
27    use std::{cell::RefCell, collections::HashMap, ptr::copy_nonoverlapping};
28
29    type WordHashMap = HashMap<U, U>;
30
31    thread_local! {
32        pub static STORAGE: RefCell<WordHashMap> = RefCell::default();
33        pub static TRANSIENT: RefCell<WordHashMap> = RefCell::default();
34    }
35
36    pub fn storage_clear() {
37        STORAGE.with(|s| s.borrow_mut().clear())
38    }
39
40    pub fn transient_clear() {
41        TRANSIENT.with(|s| s.borrow_mut().clear())
42    }
43
44    unsafe fn read_word(key: *const u8) -> U {
45        let mut r = [0u8; 32];
46        unsafe {
47            copy_nonoverlapping(key, r.as_mut_ptr(), 32);
48        }
49        U(r)
50    }
51
52    unsafe fn write_word(key: *mut u8, val: U) {
53        unsafe {
54            copy_nonoverlapping(val.as_ptr(), key, 32);
55        }
56    }
57
58    pub(crate) unsafe fn storage_load_bytes32(key: *const u8, out: *mut u8) {
59        let k = unsafe { read_word(key) };
60        let value = STORAGE.with(|s| match s.borrow().get(&k) {
61            Some(v) => *v,
62            None => U::ZERO,
63        });
64        unsafe { write_word(out, value) };
65    }
66
67    pub(crate) unsafe fn storage_cache_bytes32(key: *const u8, value: *const u8) {
68        let k = unsafe { read_word(key) };
69        let v = unsafe { read_word(value) };
70        STORAGE.with(|s| s.borrow_mut().insert(k, v));
71    }
72
73    pub(crate) unsafe fn transient_load_bytes32(key: *const u8, out: *mut u8) {
74        let k = unsafe { read_word(key) };
75        let value = TRANSIENT.with(|s| match s.borrow().get(&k) {
76            Some(v) => *v,
77            None => U::ZERO,
78        });
79        unsafe { write_word(out, value) };
80    }
81
82    pub(crate) unsafe fn transient_store_bytes32(key: *const u8, value: *const u8) {
83        let k = unsafe { read_word(key) };
84        let v = unsafe { read_word(value) };
85        TRANSIENT.with(|s| s.borrow_mut().insert(k, v));
86    }
87
88    pub unsafe fn storage_flush_cache(_: bool) {}
89}
90
91#[cfg(all(
92    not(all(target_family = "wasm", target_os = "unknown")),
93    not(feature = "std")
94))]
95mod storage_host {
96    pub(crate) unsafe fn storage_load_bytes32(_: *const u8, _: *mut u8) {}
97
98    pub(crate) unsafe fn storage_cache_bytes32(_: *const u8, _: *const u8) {}
99
100    pub(crate) unsafe fn transient_load_bytes32(_: *const u8, _: *mut u8) {}
101
102    pub(crate) unsafe fn transient_store_bytes32(_: *const u8, _: *const u8) {}
103
104    pub unsafe fn storage_flush_cache(_: bool) {}
105}
106
107#[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
108use storage_host::*;
109
110#[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
111pub use storage_host::storage_flush_cache;
112
113macro_rules! storage_ops {
114    ($($prefix:ident),* $(,)?) => {
115        $(
116            paste::paste! {
117                pub fn [<$prefix _load>](x: &U) -> U {
118                    let mut b = [0u8; 32];
119                    unsafe { [<$prefix _load_bytes32>](x.as_ptr(), b.as_mut_ptr()) }
120                    U(b)
121                }
122
123                pub fn [<$prefix _load_bool>](x: &U) -> bool {
124                    [<$prefix _load>](x).into()
125                }
126
127                /// Attempt to "exchange" a value, returning whether the expected value was set.
128                pub fn [<$prefix _exchange>](k: &U, exp: &U, new: &U) -> bool {
129                    let t = [<$prefix _load>](k);
130                    if &t != exp {
131                        return false;
132                    }
133                    [<$prefix _store>](k, new);
134                    true
135                }
136
137                pub fn [<$prefix _exchange_res>](k: &U, exp: &U, new: &U) -> Result<(), U> {
138                    let t = [<$prefix _load>](k);
139                    if &t != exp {
140                        return Err(t);
141                    }
142                    [<$prefix _store>](k, new);
143                    Ok(())
144                }
145
146                /// Set the value given, checking that the value passed has the inverse
147                /// set set currently. So, passing true would check if false is set.
148                pub fn [<$prefix _exchange_bool>](k: &U, new: bool) -> bool {
149                    [<$prefix _exchange>](k, &U::from(!new), &U::from(new))
150                }
151
152                pub fn [<$prefix _exchange_bool_res>](k: &U, new: bool) -> Result<(), bool> {
153                   let x = [<$prefix _exchange_bool>](k, new);
154                   if x == !new {
155                       Ok(())
156                   } else {
157                       Err(x)
158                   }
159                }
160            }
161        )*
162    };
163}
164
165pub fn storage_store(x: &U, y: &U) {
166    unsafe { storage_cache_bytes32(x.as_ptr(), y.as_ptr()) }
167}
168
169pub fn storage_store_bool(x: &U, y: bool) {
170    storage_store(x, &U::from(y))
171}
172
173pub fn transient_store(x: &U, y: &U) {
174    unsafe { transient_store_bytes32(x.as_ptr(), y.as_ptr()) }
175}
176
177pub fn transient_store_bool(x: &U, y: bool) {
178    transient_store(x, &U::from(y))
179}
180
181pub fn flush_cache() {
182    unsafe { storage_flush_cache(false) }
183}
184
185pub fn flush_guard<R, F: FnOnce() -> R>(f: F) -> R {
186    let r = f();
187    flush_cache();
188    r
189}
190
191storage_ops!(storage, transient);
192
193macro_rules! storage_mutate_ops {
194    ($prefix:ident, $($op:expr),* $(,)?) => {
195        $(
196            paste::paste! {
197                pub fn [<$prefix _wrapping_ $op>](x: &U, new: &U) {
198                    [<$prefix _store>](x, &bobcat_maths::[<wrapping_ $op>](&[<$prefix _load>](x), new))
199                }
200
201                pub fn [<$prefix _saturating_ $op>](x: &U, new: &U) {
202                    [<$prefix _store>](x, &bobcat_maths::[<saturating_ $op>](&[<$prefix _load>](x), new))
203                }
204
205                pub fn [<$prefix _checked_ $op>](x: &U, new: &U) -> Option<()> {
206                    let y = [<$prefix _load>](x);
207                    let v = bobcat_maths::[<checked_ $op>](&y, new);
208                    [<$prefix _store>](x, &v);
209                    Some(())
210                }
211
212                pub fn [<$prefix _checked_ $op _res>](x: &U, new: &U) -> U {
213                    let y = [<$prefix _load>](x);
214                    let v = bobcat_maths::[<checked_ $op>](&y, new);
215                    [<$prefix _store>](x, &v);
216                    v
217                }
218            }
219        )*
220    };
221}
222
223storage_mutate_ops!(storage, add, sub, mul, div);
224storage_mutate_ops!(transient, add, sub, mul, div);
225
226#[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
227pub fn slot_map_slot(k: &U, p: &U) -> U {
228    const_slot_map(k, p)
229}
230
231pub fn reentrancy_guard_entry(x: &U) {
232    assert!(x.len() <= 32, "too large");
233    assert!(transient_exchange_bool(x, true), "reentrancy alarm")
234}
235
236pub fn reentrancy_guard_exit(x: &U) {
237    assert!(x.len() <= 32, "too large");
238    transient_store(x, &U::ZERO);
239}
240
241pub fn reentrancy_guard<R>(k: &U, f: impl FnOnce() -> R) -> R {
242    reentrancy_guard_entry(k);
243    let v = f();
244    reentrancy_guard_exit(k);
245    v
246}
247
248pub fn reentrancy_guard_sel<R>(k: &[u8; 4], f: impl FnOnce() -> R) -> R {
249    reentrancy_guard::<R>(&U::from(k), f)
250}
251
252/// Compute the slot for a slice, and take it off the curve. Useful for
253/// storage slot accesses (and more).
254pub const fn const_slot_off_curve(b: &[u8]) -> U {
255    bobcat_maths::wrapping_sub(&const_keccak256(b), &U::ONE)
256}
257
258pub fn slot_off_curve(b: &[u8]) -> U {
259    // This won't result in 0 from the keccak, so we can use checked_sub to
260    // use the code the host gives us for a slightly lower codesize profile.
261    bobcat_maths::checked_sub_opt(&keccak256(b), &U::ONE).unwrap()
262}
263
264#[cfg(all(target_family = "wasm", target_os = "unknown"))]
265pub fn keccak256(b: &[u8]) -> U {
266    let mut out = [0u8; 32];
267    unsafe {
268        native_keccak256(b.as_ptr(), b.len(), out.as_mut_ptr());
269    }
270    U(out)
271}
272
273pub const fn const_keccak256(b: &[u8]) -> U {
274    U(Keccak256::new().update(b).finalize())
275}
276
277#[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
278pub fn keccak256(b: &[u8]) -> U {
279    const_keccak256(b)
280}
281
282pub fn reentrancy_guard_const_keccak<R>(k: &[u8], f: impl FnOnce() -> R) -> R {
283    reentrancy_guard(&const_keccak256(k), f)
284}
285
286pub fn reentrancy_guard_keccak<R>(k: &[u8], f: impl FnOnce() -> R) -> R {
287    reentrancy_guard(&keccak256(k), f)
288}
289
290/// Find the storage map slot using keccak_const. Don't do this during
291/// your runtime code, unless you want to pay the codesize price.
292pub const fn const_slot_map(k: &U, p: &U) -> U {
293    let a: [u8; 32 * 2] = concat_arrays!(k.0, p.0);
294    const_keccak256(&a)
295}
296
297#[cfg(all(target_family = "wasm", target_os = "unknown"))]
298pub fn slot_map(k: &U, p: &U) -> U {
299    let b: [u8; 32 * 2] = concat_arrays!(k.0, p.0);
300    keccak256(&b)
301}
302
303#[cfg(not(all(target_family = "wasm", target_os = "unknown")))]
304pub fn slot_map(k: &U, p: &U) -> U {
305    const_slot_map(k, p)
306}
307
308#[test]
309fn test_slot_edd25519_count() {
310    assert_eq!(
311        U::from(
312            const_hex::const_decode_to_array::<32>(
313                b"709318ac04e7c3155ef66c30be7220b3243d7e2378fa4153b5f14ebd3ea771ab"
314            )
315            .unwrap()
316        ),
317        const_slot_off_curve(b"superposition.passport.ed25519_count")
318    );
319}
320
321#[cfg(all(feature = "std", test))]
322mod test {
323    use super::*;
324
325    use proptest::prelude::*;
326
327    proptest! {
328        #[test]
329        fn test_reentrancy_guard(x in any::<[u8; 8]>()) {
330            reentrancy_guard(&U::from(x), || {
331                assert!(transient_load(&U::from(x)).is_true());
332            });
333            assert!(transient_load(&U::from(x)).is_zero());
334        }
335
336        #[test]
337        fn test_reentrancy_guard_bad(x in any::<[u8; 8]>()) {
338             let x = U::from(x);
339             transient_store(&x, &U::from(false));
340             assert!(transient_exchange_bool(&x, true));
341             assert!(!transient_exchange_bool(&x, true));
342            assert!(transient_load(&x).is_some());
343        }
344
345        #[test]
346        fn test_reentrancy_guard_sel(x in any::<[u8; 4]>()) {
347            reentrancy_guard_sel(&x, || {
348                assert!(transient_load(&U::from(x)).is_true());
349            });
350            assert!(transient_load(&U::from(x)).is_zero());
351        }
352    }
353}
bobcat_storage/lib.rs

bobcat_storage/
lib.rs
