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

bobcat_storage/
lib.rs
