Merge pull request #69 from avast/fix/sigsegv-on-module-import

Fix/sigsegv on module import
avast · Jan 9, 2024 · 5cfe552 · 5cfe552
2 parents 65ab2d7 + d270c90
commit 5cfe552
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 23 deletions.
diff --git a/yari-cli/src/main.rs b/yari-cli/src/main.rs
@@ -108,7 +108,7 @@ fn main() -> Result<()> {
     match matches.subcommand() {
         Some(("dump", sub_matches)) => {
             let module = Module::from_str(sub_matches.get_one::<String>("MODULE").unwrap())?;
-            context.dump_module(module);
+            context.dump_module(module)?;
         }
         _ => {
             // Start interactive shell

diff --git a/yari-sys/README.md b/yari-sys/README.md
diff --git a/yari-sys/src/lib.rs b/yari-sys/src/lib.rs
@@ -730,9 +730,9 @@ impl Context {
     }
 
     /// Import and initialize `module`.
-    fn import_module(&mut self, module: Module) {
+    fn import_module(&mut self, module: Module) -> Result<(), YariError> {
         if self.modules.contains_key(&module) {
-            return;
+            return Ok(());
         }
 
         debug!("Importing module {:?}", module);
@@ -749,8 +749,14 @@ impl Context {
         }
         .cast();
 
+        if new_module.is_null() {
+            return Err(YariError::SymbolNotFound(module.to_string()));
+        }
+
         self.modules.insert(module, new_module);
         self.init_objects_cache(new_module);
+
+        Ok(())
     }
 
     pub fn get_object(&self, path: &str) -> Option<&*mut YR_OBJECT> {
@@ -1148,7 +1154,7 @@ impl Context {
         // Import module used in expression
         let expr_module = expr.get_module();
         if let Some(module) = expr_module {
-            self.import_module(module);
+            self.import_module(module)?;
         }
 
         let rule_ctx = self.get_rule_context(rule_name)?;
@@ -1329,14 +1335,16 @@ impl Context {
         ) as *mut YR_MATCHES;
     }
 
-    pub fn dump_module(&mut self, module: Module) {
-        self.import_module(module);
+    pub fn dump_module(&mut self, module: Module) -> Result<(), YariError> {
+        self.import_module(module)?;
         match self.modules.get(&module) {
             Some(module) => {
                 self.visit_structure(module.cast::<YR_OBJECT>(), 0);
             }
             None => error!("Module '{}' not found", module),
         }
+
+        Ok(())
     }
 
     fn visit_structure(&self, structure_ptr: *const YR_OBJECT, depth: usize) {

diff --git a/yari-sys/tests/tests_dump_module.rs b/yari-sys/tests/tests_dump_module.rs
@@ -7,7 +7,7 @@ fn test_function_dump_plain() {
     let mut context = common::context();
 
     for module in MODULES {
-        context.dump_module(*module);
+        context.dump_module(*module).unwrap();
     }
 }
 
@@ -16,7 +16,7 @@ fn test_function_dump_cuckoo() {
     let mut context = common::context_with_cuckoo();
 
     for module in MODULES {
-        context.dump_module(*module);
+        context.dump_module(*module).unwrap();
     }
 }
 
@@ -25,7 +25,7 @@ fn test_function_dump_pe() {
     let mut context = common::context_with_pe_sample_and_rule();
 
     for module in MODULES {
-        context.dump_module(*module);
+        context.dump_module(*module).unwrap();
     }
 }
 
@@ -34,6 +34,6 @@ fn test_function_dump_elf() {
     let mut context = common::context_with_elf_sample();
 
     for module in MODULES {
-        context.dump_module(*module);
+        context.dump_module(*module).unwrap();
     }
 }