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 _exchange>](k: &U, exp: &U, new: &U) -> Result<(), U> {
112                    let t = [<$prefix _load>](k);
113                    if &t != exp {
114                        return Err(t);
115                    }
116                    [<$prefix _store>](k, new);
117                    Ok(())
118                }
119
120                pub fn [<$prefix _exchange_bool>](k: &U, exp: bool) -> Result<(), bool> {
121                    [<$prefix _exchange>](k, &U::from(exp), &U::from(!exp))
122                        .map_err(|x| x.is_true())
123                }
124            }
125        )*
126    };
127}
128
129pub fn storage_store(x: &U, y: &U) {
130    unsafe { storage_cache_bytes32(x.as_ptr(), y.as_ptr()) }
131}
132
133pub fn transient_store(x: &U, y: &U) {
134    unsafe { transient_store_bytes32(x.as_ptr(), y.as_ptr()) }
135}
136
137pub fn flush_cache() {
138    unsafe { storage_flush_cache(false) }
139}
140
141pub fn flush_guard<R, F: FnOnce() -> R>(f: F) -> R {
142    let r = f();
143    flush_cache();
144    r
145}
146
147storage_ops!(storage, transient);
148
149macro_rules! storage_mutate_ops {
150    ($prefix:ident, $($op:expr),* $(,)?) => {
151        $(
152            paste::paste! {
153                pub fn [<$prefix _wrapping_ $op>](x: &U, new: &U) {
154                    [<$prefix _store>](x, &bobcat_maths::[<wrapping_ $op>](&[<$prefix _load>](x), new))
155                }
156
157                pub fn [<$prefix _checked_ $op>](x: &U, new: &U) -> Option<()> {
158                    let y = [<$prefix _load>](x);
159                    let v = bobcat_maths::[<checked_ $op>](&y, new)?;
160                    [<$prefix _store>](x, &v);
161                    Some(())
162                }
163
164                pub fn [<$prefix _checked_ $op _res>](x: &U, new: &U) -> Result<U, (U, U)> {
165                    let y = [<$prefix _load>](x);
166                    let Some(v) = bobcat_maths::[<checked_ $op>](&y, new) else {
167                        return Err((y, *new));
168                    };
169                    [<$prefix _store>](x, &v);
170                    Ok(v)
171                }
172            }
173        )*
174    };
175}
176
177storage_mutate_ops!(storage, add, sub, mul, div);
178storage_mutate_ops!(transient, add, sub, mul, div);
179
180#[cfg(not(target_arch = "wasm32"))]
181pub fn slot_map_slot(k: &U, p: &U) -> U {
182    const_slot_map(k, p)
183}
184
185pub fn reentrancy_guard_entry(x: &U) -> Result<(), bool> {
186    assert!(x.len() <= 32, "too large");
187    transient_exchange_bool(x, false)
188}
189
190pub fn reentrancy_guard_exit(x: &U) {
191    assert!(x.len() <= 32, "too large");
192    transient_store(x, &U::ZERO);
193}
194
195pub fn reentrancy_guard<R>(k: &U, f: impl FnOnce() -> R) -> Result<R, bool> {
196    reentrancy_guard_entry(k)?;
197    let v = f();
198    reentrancy_guard_exit(k);
199    Ok(v)
200}
201
202/// Compute the slot for a slice, and take it off the curve. Useful for
203/// storage slot accesses (and more).
204pub const fn const_slot_off_curve(b: &[u8]) -> U {
205    bobcat_maths::wrapping_sub(&const_keccak256(b), &U::ONE)
206}
207
208pub fn slot_off_curve(b: &[u8]) -> U {
209    // This won't result in 0 from the keccak, so we can use checked_sub to
210    // use the code the host gives us for a slightly lower codesize profile.
211    bobcat_maths::checked_sub(&keccak256(b), &U::ONE).unwrap()
212}
213
214#[cfg(target_arch = "wasm32")]
215pub fn keccak256(b: &[u8]) -> U {
216    let mut out = [0u8; 32];
217    unsafe {
218        native_keccak256(b.as_ptr(), b.len(), out.as_mut_ptr());
219    }
220    U(out)
221}
222
223pub const fn const_keccak256(b: &[u8]) -> U {
224    U(Keccak256::new().update(b).finalize())
225}
226
227#[cfg(not(target_arch = "wasm32"))]
228pub fn keccak256(b: &[u8]) -> U {
229    const_keccak256(b)
230}
231
232pub fn reentrancy_guard_const_keccak<R>(k: &[u8], f: impl FnOnce() -> R) -> Result<R, bool> {
233    reentrancy_guard(&const_keccak256(k), f)
234}
235
236pub fn reentrancy_guard_keccak<R>(k: &[u8], f: impl FnOnce() -> R) -> Result<R, bool> {
237    reentrancy_guard(&keccak256(k), f)
238}
239
240/// Find the storage map slot using keccak_const. Don't do this during
241/// your runtime code, unless you want to pay the codesize price.
242pub const fn const_slot_map(k: &U, p: &U) -> U {
243    let a: [u8; 32 * 2] = concat_arrays!(k.0, p.0);
244    const_keccak256(&a)
245}
246
247#[cfg(target_arch = "wasm32")]
248pub fn slot_map(k: &U, p: &U) -> U {
249    let b: [u8; 32 * 2] = concat_arrays!(k.0, p.0);
250    keccak256(&b)
251}
252
253#[cfg(not(target_arch = "wasm32"))]
254pub fn slot_map(k: &U, p: &U) -> U {
255    const_slot_map(k, p)
256}
257
258#[test]
259fn test_slot_edd25519_count() {
260    assert_eq!(
261        U::from(
262            const_hex::const_decode_to_array::<32>(
263                b"709318ac04e7c3155ef66c30be7220b3243d7e2378fa4153b5f14ebd3ea771ab"
264            )
265            .unwrap()
266        ),
267        const_slot_off_curve(b"superposition.passport.ed25519_count")
268    );
269}
270
271#[cfg(all(feature = "std", test))]
272mod test {
273    use super::*;
274
275    use proptest::prelude::*;
276
277    proptest! {
278        #[test]
279        fn test_reentrancy_guard(x in any::<[u8; 8]>()) {
280            reentrancy_guard(&U::from(x), || {
281                assert!(transient_load(&U::from(x)).is_true());
282            })
283            .unwrap();
284            assert!(transient_load(&U::from(x)).is_zero());
285        }
286    }
287}
bobcat_storage/lib.rs

bobcat_storage/
lib.rs
