1use crate::error::Error;
2
3use std::borrow::{Borrow, BorrowMut};
4use std::ptr::NonNull;
5use std::sync::{Arc, RwLock};
6use wasmer::{Instance, Memory, WasmerEnv};
7use wasmer_middlewares::metering::{get_remaining_points, set_remaining_points, MeteringPoints};
8
9pub trait Querier {
10 fn get_span_size(&self) -> i64;
12 fn get_calldata(&self) -> Result<Vec<u8>, Error>;
14 fn set_return_data(&self, data: &[u8]) -> Result<(), Error>;
16 fn get_ask_count(&self) -> i64;
18 fn get_min_count(&self) -> i64;
20 fn get_prepare_time(&self) -> i64;
22 fn get_execute_time(&self) -> Result<i64, Error>;
24 fn get_ans_count(&self) -> Result<i64, Error>;
26 fn ask_external_data(&self, eid: i64, did: i64, data: &[u8]) -> Result<(), Error>;
28 fn get_external_data_status(&self, eid: i64, vid: i64) -> Result<i64, Error>;
30 fn get_external_data(&self, eid: i64, vid: i64) -> Result<Vec<u8>, Error>;
32}
33
34pub struct ContextData<Q: Querier> {
35 querier: Q,
36 wasmer_instance: Option<NonNull<Instance>>,
38}
39
40impl<Q: Querier> ContextData<Q> {
41 pub fn new(querier: Q) -> Self {
42 ContextData::<Q> { wasmer_instance: None, querier }
43 }
44}
45
46#[derive(WasmerEnv)]
47pub struct Environment<Q>
48where
49 Q: Querier + 'static,
50{
51 data: Arc<RwLock<ContextData<Q>>>,
52}
53
54impl<Q: Querier + 'static> Clone for Environment<Q> {
55 fn clone(&self) -> Self {
56 Self { data: self.data.clone() }
57 }
58}
59unsafe impl<Q: Querier> Send for Environment<Q> {}
60unsafe impl<Q: Querier> Sync for Environment<Q> {}
61
62impl<Q> Environment<Q>
63where
64 Q: Querier + 'static,
65{
66 pub fn new(q: Q) -> Self {
67 Self { data: Arc::new(RwLock::new(ContextData::new(q))) }
68 }
69
70 pub fn with_querier_from_context<C, R>(&self, callback: C) -> R
71 where
72 C: FnOnce(&Q) -> R,
73 {
74 self.with_context_data(|context_data| callback(&context_data.querier))
75 }
76
77 pub fn set_wasmer_instance(&self, instance: Option<NonNull<Instance>>) {
79 self.with_context_data_mut(|data| {
80 data.wasmer_instance = instance;
81 })
82 }
83
84 pub fn with_wasmer_instance<C, R>(&self, callback: C) -> Result<R, Error>
85 where
86 C: FnOnce(&Instance) -> Result<R, Error>,
87 {
88 self.with_context_data(|context_data| match context_data.wasmer_instance {
89 Some(instance_ptr) => {
90 let instance_ref = unsafe { instance_ptr.as_ref() };
91 callback(instance_ref)
92 }
93 None => Err(Error::UninitializedContextData),
94 })
95 }
96
97 fn with_context_data<C, R>(&self, callback: C) -> R
98 where
99 C: FnOnce(&ContextData<Q>) -> R,
100 {
101 let guard = self.data.as_ref().read().unwrap();
102 let context_data = guard.borrow();
103 callback(context_data)
104 }
105
106 fn with_context_data_mut<C, R>(&self, callback: C) -> R
107 where
108 C: FnOnce(&mut ContextData<Q>) -> R,
109 {
110 let mut guard = self.data.as_ref().write().unwrap();
111 let context_data = guard.borrow_mut();
112 callback(context_data)
113 }
114
115 pub fn get_gas_left(&self) -> u64 {
116 self.with_wasmer_instance(|instance| {
117 Ok(match get_remaining_points(instance) {
118 MeteringPoints::Remaining(count) => count,
119 MeteringPoints::Exhausted => 0,
120 })
121 })
122 .expect("Wasmer instance is not set. This is a bug in the lifecycle.")
123 }
124
125 pub fn set_gas_left(&self, new_value: u64) {
126 self.with_wasmer_instance(|instance| {
127 set_remaining_points(instance, new_value);
128 Ok(())
129 })
130 .expect("Wasmer instance is not set. This is a bug in the lifecycle.")
131 }
132
133 pub fn decrease_gas_left(&self, gas: u64) -> Result<(), Error> {
134 let gas_left = self.get_gas_left();
135 if gas > gas_left {
136 Err(Error::OutOfGasError)
137 } else {
138 self.set_gas_left(gas_left.saturating_sub(gas));
139 Ok(())
140 }
141 }
142
143 pub fn memory(&self) -> Result<Memory, Error> {
144 self.with_context_data(|data| match data.wasmer_instance {
145 Some(instance_ptr) => {
146 let instance_ref = unsafe { instance_ptr.as_ref() };
147 let mut memories: Vec<Memory> =
148 instance_ref.exports.iter().memories().map(|pair| pair.1.clone()).collect();
149
150 match memories.pop() {
151 Some(m) => Ok(m),
152 None => Err(Error::MemoryOutOfBoundError),
153 }
154 }
155 _ => Err(Error::BadMemorySectionError),
156 })
157 }
158}
159
160#[cfg(test)]
161mod tests {
162 use std::{
163 io::{Read, Write},
164 process::Command,
165 };
166
167 use tempfile::NamedTempFile;
168 use wasmer::{imports, Singlepass, Store, Universal};
169
170 use crate::{
171 cache::{Cache, CacheOptions},
172 store::make_store,
173 };
174
175 use super::*;
176
177 pub struct MockQuerier {}
178
179 impl Querier for MockQuerier {
180 fn get_span_size(&self) -> i64 {
181 300
182 }
183 fn get_calldata(&self) -> Result<Vec<u8>, Error> {
184 Ok(vec![1])
185 }
186 fn set_return_data(&self, _: &[u8]) -> Result<(), Error> {
187 Ok(())
188 }
189 fn get_ask_count(&self) -> i64 {
190 10
191 }
192 fn get_min_count(&self) -> i64 {
193 8
194 }
195 fn get_prepare_time(&self) -> i64 {
196 100_000
197 }
198 fn get_execute_time(&self) -> Result<i64, Error> {
199 Ok(100_000)
200 }
201 fn get_ans_count(&self) -> Result<i64, Error> {
202 Ok(8)
203 }
204 fn ask_external_data(&self, _: i64, _: i64, _: &[u8]) -> Result<(), Error> {
205 Ok(())
206 }
207 fn get_external_data_status(&self, _: i64, _: i64) -> Result<i64, Error> {
208 Ok(1)
209 }
210 fn get_external_data(&self, _: i64, _: i64) -> Result<Vec<u8>, Error> {
211 Ok(vec![1])
212 }
213 }
214
215 fn wat2wasm(wat: impl AsRef<[u8]>) -> Vec<u8> {
216 let mut input_file = NamedTempFile::new().unwrap();
217 let mut output_file = NamedTempFile::new().unwrap();
218 input_file.write_all(wat.as_ref()).unwrap();
219 Command::new("wat2wasm")
220 .args(&[
221 input_file.path().to_str().unwrap(),
222 "-o",
223 output_file.path().to_str().unwrap(),
224 ])
225 .output()
226 .unwrap();
227 let mut wasm = Vec::new();
228 output_file.read_to_end(&mut wasm).unwrap();
229 wasm
230 }
231
232 #[test]
233 fn test_env_querier() {
234 let env = Environment::new(MockQuerier {});
235 assert_eq!(300, env.with_querier_from_context(|querier| querier.get_span_size()));
236 }
237
238 #[test]
239 fn test_env_wasmer_instance() {
240 let env = Environment::new(MockQuerier {});
241 assert_eq!(
242 Error::UninitializedContextData,
243 env.with_wasmer_instance(|_| { Ok(()) }).unwrap_err()
244 );
245
246 let wasm = wat2wasm(
247 r#"(module
248 (func $execute (export "execute"))
249 (func $prepare (export "prepare"))
250 )"#,
251 );
252 let compiler = Singlepass::new();
253 let store = Store::new(&Universal::new(compiler).engine());
254 let import_object = imports! {};
255 let mut cache = Cache::new(CacheOptions { cache_size: 10000 });
256 let (instance, _) = cache.get_instance(&wasm, &store, &import_object).unwrap();
257 env.set_wasmer_instance(Some(NonNull::from(&instance)));
258 assert_eq!(Ok(()), env.with_wasmer_instance(|_| { Ok(()) }));
259 }
260
261 #[test]
262 fn test_env_gas() {
263 let env = Environment::new(MockQuerier {});
264 let wasm = wat2wasm(
265 r#"(module
266 (func $execute (export "execute"))
267 (func $prepare (export "prepare"))
268 )"#,
269 );
270 let store = make_store();
271 let import_object = imports! {};
272 let mut cache = Cache::new(CacheOptions { cache_size: 10000 });
273 let (instance, _) = cache.get_instance(&wasm, &store, &import_object).unwrap();
274 env.set_wasmer_instance(Some(NonNull::from(&instance)));
275
276 assert_eq!(0, env.get_gas_left());
277
278 env.set_gas_left(10);
279 assert_eq!(10, env.get_gas_left());
280
281 assert_eq!(Error::OutOfGasError, env.decrease_gas_left(11).unwrap_err());
282 assert_eq!(Ok(()), env.decrease_gas_left(3));
283 assert_eq!(7, env.get_gas_left());
284 }
285}