1
use nomiscript::SymbolTable;
2
use tracing::debug;
3
use wasmtime::{Engine, Linker, Module, Store};
4

            
5
use crate::error::HookError;
6
use crate::format::{BASE_OFFSET, GlobalHeader, OUTPUT_HEADER_SIZE, OutputHeader};
7
use crate::host::{WasmHost, define_host_functions};
8
use crate::parser::{OutputParser, ParsedEntity};
9
use crate::runtime::{EngineOpts, build_engine};
10

            
11
const DEFAULT_OUTPUT_SIZE: u32 = 64 * 1024;
12
const WASM_PAGE_SIZE: u32 = 65536;
13
/// Wall-clock budget for one script execution, in [`EPOCH_TICK_INTERVAL`] ticks.
14
///
15
/// This is a real-time bound, not a fuel bound, so it is sensitive to how fast
16
/// the host happens to be. Coverage-instrumented builds run several times
17
/// slower, which is why [`ScriptExecutor::with_epoch_deadline_ticks`] exists:
18
/// tests that only care that a script *completes* raise the budget rather than
19
/// racing the clock.
20
const DEFAULT_EPOCH_DEADLINE_TICKS: u64 = 5;
21
const EPOCH_TICK_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1);
22

            
23
/// Drives the engine's epoch clock for the lifetime of one execution.
24
///
25
/// Stopping matters: the epoch counter is shared by every execution on the same
26
/// engine, so a ticker that outlives its call keeps consuming the budget of
27
/// whatever runs next. A `JoinHandle` does not stop a thread when dropped, it
28
/// detaches it — so the previous `drop(ticker)` left it ticking for up to
29
/// `EPOCH_DEADLINE_TICKS` seconds, and every `?` between the spawn and the call
30
/// skipped even that. Executors are reused (the account-script batch path runs
31
/// one across many transactions), which is exactly where a stray bump lands.
32
///
33
/// The stop channel doubles as the sleep: `recv_timeout` returns `Timeout` for a
34
/// real tick and `Disconnected` the moment the guard drops, so teardown is
35
/// prompt rather than waiting out the current second.
36
struct EpochTicker {
37
    stop: Option<std::sync::mpsc::Sender<()>>,
38
    handle: Option<std::thread::JoinHandle<()>>,
39
}
40

            
41
impl EpochTicker {
42
3111
    fn spawn(engine: Engine, ticks: u64) -> Self {
43
3111
        let (stop, rx) = std::sync::mpsc::channel::<()>();
44
3111
        let handle = std::thread::spawn(move || {
45
3111
            for _ in 0..ticks {
46
3111
                match rx.recv_timeout(EPOCH_TICK_INTERVAL) {
47
                    Err(std::sync::mpsc::RecvTimeoutError::Timeout) => engine.increment_epoch(),
48
3111
                    _ => return,
49
                }
50
            }
51
3111
        });
52
3111
        Self {
53
3111
            stop: Some(stop),
54
3111
            handle: Some(handle),
55
3111
        }
56
3111
    }
57
}
58

            
59
impl Drop for EpochTicker {
60
3111
    fn drop(&mut self) {
61
3111
        self.stop.take();
62
3111
        if let Some(handle) = self.handle.take() {
63
3111
            let _ = handle.join();
64
3111
        }
65
3111
    }
66
}
67

            
68
pub struct ScriptExecutor {
69
    host: WasmHost,
70
    epoch_deadline_ticks: u64,
71
}
72

            
73
impl ScriptExecutor {
74
    /// Builds an executor on the baseline engine.
75
    ///
76
    /// Fallible because engine construction is: `wasmtime` validates the
77
    /// configuration and rejects unsupported combinations. There is no
78
    /// infallible `new`/`Default` for the same reason — a constructor that
79
    /// cannot report that failure can only paper over it.
80
1276
    pub fn try_new() -> Result<Self, HookError> {
81
        Ok(Self {
82
1276
            host: WasmHost::new(build_engine(EngineOpts::baseline())?, SymbolTable::new()),
83
            epoch_deadline_ticks: DEFAULT_EPOCH_DEADLINE_TICKS,
84
        })
85
1276
    }
86

            
87
    #[must_use]
88
1428
    pub fn with_engine(engine: Engine) -> Self {
89
1428
        Self {
90
1428
            host: WasmHost::new(engine, SymbolTable::new()),
91
1428
            epoch_deadline_ticks: DEFAULT_EPOCH_DEADLINE_TICKS,
92
1428
        }
93
1428
    }
94

            
95
    /// Overrides the per-execution wall-clock budget.
96
    ///
97
    /// Production uses [`DEFAULT_EPOCH_DEADLINE_TICKS`]; a caller that runs
98
    /// under coverage instrumentation, a debugger, or a loaded CI node raises
99
    /// it so a slow-but-correct script is not cut off mid-execution.
100
    #[must_use]
101
51
    pub fn with_epoch_deadline_ticks(mut self, ticks: u64) -> Self {
102
51
        self.epoch_deadline_ticks = ticks;
103
51
        self
104
51
    }
105

            
106
3213
    fn get_or_compile_module(&self, bytecode: &[u8]) -> Result<Module, HookError> {
107
3213
        self.host
108
3213
            .module_cache()
109
3213
            .get_or_compile(self.host.engine(), bytecode)
110
3213
            .map_err(HookError::from)
111
3213
    }
112

            
113
3213
    pub fn execute(
114
3213
        &self,
115
3213
        bytecode: &[u8],
116
3213
        input: &[u8],
117
3213
        output_size: Option<u32>,
118
3213
    ) -> Result<Vec<ParsedEntity>, HookError> {
119
3213
        debug!(bytecode_size = bytecode.len(), "script execution start");
120
3213
        let output_size = output_size.unwrap_or(DEFAULT_OUTPUT_SIZE);
121
3213
        let module = self.get_or_compile_module(bytecode)?;
122
3111
        debug!("module compiled");
123

            
124
3111
        let header = GlobalHeader::from_bytes(input)
125
3111
            .ok_or_else(|| HookError::Parse("Invalid input header".to_string()))?;
126

            
127
3111
        let input_offset = BASE_OFFSET;
128
3111
        let output_offset = input_offset + input.len() as u32;
129
3111
        let strings_offset = header.strings_pool_offset;
130

            
131
3111
        let exec_state = self
132
3111
            .host
133
3111
            .execution_state(input_offset, output_offset, strings_offset);
134
3111
        let mut store = Store::new(self.host.engine(), exec_state);
135
3111
        store.set_epoch_deadline(self.epoch_deadline_ticks);
136

            
137
3111
        let _ticker = EpochTicker::spawn(self.host.engine().clone(), self.epoch_deadline_ticks);
138

            
139
3111
        let mut linker = Linker::new(self.host.engine());
140
3111
        define_host_functions(&mut linker)?;
141

            
142
3111
        let instance = linker.instantiate(&mut store, &module)?;
143

            
144
3111
        let memory = instance
145
3111
            .get_memory(&mut store, "memory")
146
3111
            .ok_or(HookError::WASMMem)?;
147

            
148
3111
        store.data_mut().memory = Some(memory);
149

            
150
3111
        let total_size = input.len() + output_size as usize;
151
3111
        let required_pages = (BASE_OFFSET as usize + total_size).div_ceil(WASM_PAGE_SIZE as usize);
152
3111
        let current_pages = memory.size(&store) as usize;
153

            
154
3111
        if required_pages > current_pages {
155
            memory.grow(&mut store, (required_pages - current_pages) as u64)?;
156
3111
        }
157

            
158
3111
        let mem_data = memory.data_mut(&mut store);
159
3111
        let input_start = BASE_OFFSET as usize;
160
3111
        mem_data[input_start..input_start + input.len()].copy_from_slice(input);
161

            
162
3111
        let output_start = output_offset as usize;
163
3111
        let input_entity_count = header.input_entity_count;
164
3111
        let output_header = OutputHeader::new(input_entity_count);
165
3111
        mem_data[output_start..output_start + OUTPUT_HEADER_SIZE]
166
3111
            .copy_from_slice(&output_header.to_bytes());
167

            
168
3111
        let should_apply = instance
169
3111
            .get_typed_func::<(), i32>(&mut store, "should_apply")
170
3111
            .map_err(|e| HookError::Script(format!("Missing should_apply export: {e}")))?;
171

            
172
3111
        let result = should_apply.call(&mut store, ())?;
173
3111
        debug!(should_apply = result, "should_apply result");
174
3111
        if result == 0 {
175
867
            return Ok(Vec::new());
176
2244
        }
177

            
178
2244
        let process = instance
179
2244
            .get_typed_func::<(), ()>(&mut store, "process")
180
2244
            .map_err(|e| HookError::Script(format!("Missing process export: {e}")))?;
181

            
182
2244
        debug!("calling process");
183
2244
        process.call(&mut store, ())?;
184

            
185
2193
        let mem_data = memory.data(&store);
186
2193
        let output_data = &mem_data[output_start..output_start + output_size as usize];
187

            
188
2193
        let output_header = OutputHeader::from_bytes(output_data)
189
2193
            .ok_or_else(|| HookError::Parse("Invalid output header".to_string()))?;
190

            
191
2193
        let output_strings_offset = { output_header.strings_offset } as usize;
192

            
193
2193
        let parser = OutputParser::new(output_data, output_strings_offset)?;
194
2193
        let entities: Result<Vec<ParsedEntity>, HookError> = parser.entities().collect();
195
2193
        debug!(
196
255
            entity_count = entities.as_ref().map_or(0, std::vec::Vec::len),
197
            "output parse complete"
198
        );
199
2193
        entities
200
3213
    }
201
}
202

            
203
#[cfg(test)]
204
mod tests {
205
    use super::*;
206

            
207
    #[test]
208
1
    fn test_executor_creation() {
209
1
        let executor = ScriptExecutor::try_new().expect("baseline engine");
210
1
        assert!(
211
1
            executor
212
1
                .host
213
1
                .module_cache()
214
1
                .is_empty()
215
1
                .expect("cache lock must not be poisoned in fresh executor")
216
        );
217
1
    }
218
}