Skip to content

Commit

Permalink
refactor: simplify the logic for returning timeouts.
Browse files Browse the repository at this point in the history
Remove the global WASM variable `timeout_occurred` which tracked whether a timeout occurred during the execution of `search_for_pattern`. This variable is not really necessary because once `search_for_pattern` returns to WASM code, the wasmtime's timeout mechanism will produce a timeout error.
  • Loading branch information
plusvic committed Oct 18, 2024
1 parent 89d6013 commit 9a92dd0
Show file tree
Hide file tree
Showing 4 changed files with 18 additions and 58 deletions.
32 changes: 11 additions & 21 deletions lib/src/compiler/emit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -938,28 +938,18 @@ fn emit_lazy_pattern_search(
// do here.
},
|_else| {
// Call `search_for_patterns`.
_else.call(
ctx.function_id(
_else
// Call `search_for_patterns`.
.call(ctx.function_id(
wasm::export__search_for_patterns.mangled_name,
),
);
// `search_for_patterns` returns `true` when everything went ok, and
// `false` when a timeout occurs.
_else.if_else(
None,
|_then| {
// Everything ok, set pattern_search_done to true.
_then.i32_const(1);
_then.global_set(ctx.wasm_symbols.pattern_search_done);
},
|_else| {
// A timeout occurred, set the global variable
// `timeout_occurred` to true.
_else.i32_const(1);
_else.global_set(ctx.wasm_symbols.timeout_occurred);
},
);
))
// Remove the result of `search_for_patterns` from the stack.
// The result is `true` if everything went fine and `false`
// in case of timeout, but we don't use this result.
.drop()
// Set `pattern_search_done` to true.
.i32_const(1)
.global_set(ctx.wasm_symbols.pattern_search_done);
},
);
}
Expand Down
24 changes: 2 additions & 22 deletions lib/src/scanner/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -236,15 +236,6 @@ impl<'r> Scanner<'r> {
)
.unwrap();

// Global variable that is set to `true` when a timeout occurs during
// the scanning phase.
let timeout_occurred = Global::new(
wasm_store.as_context_mut(),
GlobalType::new(ValType::I32, Mutability::Var),
Val::I32(0),
)
.unwrap();

// Compute the base offset for the bitmap that contains matching
// information for patterns. This bitmap has 1 bit per pattern, the
// N-th bit is set if pattern with PatternId = N matched. The bitmap
Expand Down Expand Up @@ -286,13 +277,6 @@ impl<'r> Scanner<'r> {
pattern_search_done,
)
.unwrap()
.define(
wasm_store.as_context(),
"yara_x",
"timeout_occurred",
timeout_occurred,
)
.unwrap()
.define(
wasm_store.as_context(),
"yara_x",
Expand Down Expand Up @@ -741,10 +725,7 @@ impl<'r> Scanner<'r> {
// scanning) only if necessary.
//
// This will return Err(ScanError::Timeout), when the scan timeout is
// reached while WASM code is being executed. If the timeout occurs
// while ScanContext::search_for_patterns is being executed, the result
// will be Ok(1). If the scan completes successfully the result is
// Ok(0).`
// reached while WASM code is being executed.
let func_result =
self.wasm_main_func.call(self.wasm_store.as_context_mut(), ());

Expand Down Expand Up @@ -800,8 +781,7 @@ impl<'r> Scanner<'r> {

match func_result {
Ok(0) => Ok(ScanResults::new(self.wasm_store.data(), data)),
Ok(1) => Err(ScanError::Timeout),
Ok(_) => unreachable!(),
Ok(v) => panic!("WASM main returned: {}", v),
Err(err) if err.is::<ScanError>() => {
Err(err.downcast::<ScanError>().unwrap())
}
Expand Down
16 changes: 5 additions & 11 deletions lib/src/wasm/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,6 @@ impl WasmModuleBuilder {
global_const!(module, matching_patterns_bitmap_base, I32);
global_var!(module, filesize, I64);
global_var!(module, pattern_search_done, I32);
global_var!(module, timeout_occurred, I32);

let (main_memory, _) = module.add_import_memory(
"yara_x",
Expand All @@ -182,7 +181,6 @@ impl WasmModuleBuilder {
check_for_pattern_match,
filesize,
pattern_search_done,
timeout_occurred,
i64_tmp_a: module.locals.add(I64),
i64_tmp_b: module.locals.add(I64),
i32_tmp: module.locals.add(I32),
Expand All @@ -203,11 +201,9 @@ impl WasmModuleBuilder {
FunctionBuilder::new(&mut module.types, &[], &[I32]);

// The first instructions in the main function initialize the global
// variables `pattern_search_done` and `timeout_occurred` to 0 (false).
// variables `pattern_search_done`.
main_func.func_body().i32_const(0);
main_func.func_body().global_set(pattern_search_done);
main_func.func_body().i32_const(0);
main_func.func_body().global_set(timeout_occurred);

let namespace_block = namespace_func.dangling_instr_seq(None).id();

Expand Down Expand Up @@ -346,12 +342,10 @@ impl WasmModuleBuilder {
self.finish_namespace_block();
self.finish_namespace_func();

// Emit the last few instructions for the main function, which consist
// in putting the return value in the stack. The return value is 0 if
// everything went ok and 1 if a timeout occurred.
self.main_func
.func_body()
.global_get(self.wasm_symbols.timeout_occurred);
// Emit the last instruction for the main function, which consist
// in putting the return value in the stack. The return value is
// always 0.
self.main_func.func_body().i32_const(0);

let main_func =
self.main_func.finish(Vec::new(), &mut self.module.funcs);
Expand Down
4 changes: 0 additions & 4 deletions lib/src/wasm/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -693,10 +693,6 @@ pub(crate) struct WasmSymbols {
/// evaluated and some of them needs to know if a pattern matched or not.
pub pattern_search_done: walrus::GlobalId,

/// Global variable that is set to true when a timeout during the scanning
/// phase.
pub timeout_occurred: walrus::GlobalId,

/// Local variables used for temporary storage.
pub i64_tmp_a: walrus::LocalId,
pub i64_tmp_b: walrus::LocalId,
Expand Down

0 comments on commit 9a92dd0

Please sign in to comment.