From 3ea32c4cbdf2e2bc2d1b50b1822c114c9582f8a5 Mon Sep 17 00:00:00 2001 From: Ivan Krivosheev Date: Wed, 21 Jul 2021 11:14:22 +0300 Subject: [PATCH] feat: Scanner callback API --- src/internals/scan.rs | 96 ++++++++++++++++++------------------------- src/scanner.rs | 90 +++++++++++++++++++++++++++++++++------- 2 files changed, 117 insertions(+), 69 deletions(-) diff --git a/src/internals/scan.rs b/src/internals/scan.rs index 2fb6637..fbf0d7a 100644 --- a/src/internals/scan.rs +++ b/src/internals/scan.rs @@ -68,9 +68,7 @@ pub fn rules_scan_mem<'a>( flags: i32, callback: impl FnMut(CallbackMsg<'a>) -> CallbackReturn, ) -> Result<(), YaraError> { - let p_callback: Box) -> CallbackReturn>> = - Box::new(Box::new(callback)); - let user_data = Box::into_raw(p_callback) as *mut c_void; + let p_callback: Box) -> CallbackReturn> = Box::new(callback); let result = unsafe { yara_sys::yr_rules_scan_mem( rules, @@ -78,12 +76,10 @@ pub fn rules_scan_mem<'a>( mem.len().try_into().unwrap(), flags, Some(scan_callback), - user_data, + &p_callback as *const Box) -> CallbackReturn> as *mut _, timeout, ) }; - let _: Box) -> CallbackReturn>> = - unsafe { Box::from_raw(user_data as *mut _) }; yara_sys::Error::from_code(result) .map_err(|e| e.into()) @@ -99,16 +95,15 @@ pub fn scanner_scan_mem<'a>( mem: &[u8], callback: impl FnMut(CallbackMsg<'a>) -> CallbackReturn, ) -> Result<(), YaraError> { - let p_callback: Box) -> CallbackReturn>> = - Box::new(Box::new(callback)); - let user_data = Box::into_raw(p_callback) as *mut c_void; + let p_callback: Box) -> CallbackReturn> = Box::new(callback); let result = unsafe { - yara_sys::yr_scanner_set_callback(scanner, Some(scan_callback), user_data); + yara_sys::yr_scanner_set_callback( + scanner, + Some(scan_callback), + &p_callback as *const Box) -> CallbackReturn> as *mut _, + ); yara_sys::yr_scanner_scan_mem(scanner, mem.as_ptr(), mem.len().try_into().unwrap()) }; - let _: Box) -> CallbackReturn>> = - unsafe { Box::from_raw(user_data as *mut _) }; - yara_sys::Error::from_code(result) .map_err(|e| e.into()) .map(|_| ()) @@ -153,15 +148,17 @@ pub fn rules_scan_raw<'a>( callback: impl FnMut(CallbackMsg<'a>) -> CallbackReturn, ) -> i32 { let fd = file.as_raw_fd(); - let p_callback: Box) -> CallbackReturn>> = - Box::new(Box::new(callback)); - let user_data = Box::into_raw(p_callback) as *mut c_void; - let result = unsafe { - yara_sys::yr_rules_scan_fd(rules, fd, flags, Some(scan_callback), user_data, timeout) - }; - let _: Box) -> CallbackReturn>> = - unsafe { Box::from_raw(user_data as *mut _) }; - result + let p_callback: Box) -> CallbackReturn> = Box::new(callback); + unsafe { + yara_sys::yr_rules_scan_fd( + rules, + fd, + flags, + Some(scan_callback), + &p_callback as *const Box) -> CallbackReturn> as *mut _, + timeout, + ) + } } #[cfg(windows)] @@ -173,23 +170,18 @@ pub fn rules_scan_raw<'a>( callback: impl FnMut(CallbackMsg<'a>) -> CallbackReturn, ) -> i32 { let handle = file.as_raw_handle(); - let p_callback: Box) -> CallbackReturn>> = - Box::new(Box::new(callback)); - let user_data = Box::into_raw(p_callback) as *mut c_void; + let p_callback: Box) -> CallbackReturn> = Box::new(callback); - let result = unsafe { + unsafe { yara_sys::yr_rules_scan_fd( rules, handle, flags, Some(scan_callback), - user_data, + &p_callback as *const Box) -> CallbackReturn> as *mut _, timeout, ) - }; - let _: Box) -> CallbackReturn>> = - unsafe { Box::from_raw(user_data as *mut _) }; - result + } } #[cfg(unix)] @@ -203,16 +195,15 @@ pub fn scanner_scan_raw<'a>( callback: impl FnMut(CallbackMsg<'a>) -> CallbackReturn, ) -> i32 { let fd = file.as_raw_fd(); - let p_callback: Box) -> CallbackReturn>> = - Box::new(Box::new(callback)); - let user_data = Box::into_raw(p_callback) as *mut c_void; - let result = unsafe { - yara_sys::yr_scanner_set_callback(scanner, Some(scan_callback), user_data); + let p_callback: Box) -> CallbackReturn> = Box::new(callback); + unsafe { + yara_sys::yr_scanner_set_callback( + scanner, + Some(scan_callback), + &p_callback as *const Box) -> CallbackReturn> as *mut _, + ); yara_sys::yr_scanner_scan_fd(scanner, fd) - }; - let _: Box) -> CallbackReturn>> = - unsafe { Box::from_raw(user_data as *mut _) }; - result + } } #[cfg(windows)] @@ -226,16 +217,15 @@ pub fn scanner_scan_raw<'a>( callback: impl FnMut(CallbackMsg<'a>) -> CallbackReturn, ) -> i32 { let handle = file.as_raw_handle(); - let p_callback: Box) -> CallbackReturn>> = - Box::new(Box::new(callback)); - let user_data = Box::into_raw(p_callback) as *mut c_void; - let result = unsafe { - yara_sys::yr_scanner_set_callback(scanner, Some(scan_callback), user_data); + let p_callback: Box) -> CallbackReturn> = Box::new(callback); + unsafe { + yara_sys::yr_scanner_set_callback( + scanner, + Some(scan_callback), + &p_callback as *const Box) -> CallbackReturn> as *mut _, + ); yara_sys::yr_scanner_scan_fd(scanner, handle) - }; - let _: Box) -> CallbackReturn>> = - unsafe { Box::from_raw(user_data as *mut _) }; - result + } } /// Attach a process, pause it, and scan its memory. @@ -246,21 +236,17 @@ pub fn rules_scan_proc<'a>( flags: i32, callback: impl FnMut(CallbackMsg<'a>) -> CallbackReturn, ) -> Result<(), YaraError> { - let p_callback: Box) -> CallbackReturn>> = - Box::new(Box::new(callback)); - let user_data = Box::into_raw(p_callback) as *mut c_void; + let p_callback: Box) -> CallbackReturn> = Box::new(callback); let result = unsafe { yara_sys::yr_rules_scan_proc( rules, pid as i32, flags, Some(scan_callback), - user_data, + &p_callback as *const Box) -> CallbackReturn> as *mut _, timeout, ) }; - let _: Box) -> CallbackReturn>> = - unsafe { Box::from_raw(user_data as *mut _) }; yara_sys::Error::from_code(result) .map_err(|e| e.into()) diff --git a/src/scanner.rs b/src/scanner.rs index 9b3247e..beeca5e 100644 --- a/src/scanner.rs +++ b/src/scanner.rs @@ -1,9 +1,11 @@ use std::fs::File; use std::marker::PhantomData; use std::path::Path; + pub use yara_sys::scan_flags::*; -use crate::{compiler::CompilerVariableValue, errors::*, internals, rules::Rules, Rule}; +use crate::{compiler::CompilerVariableValue, errors::*, internals, Rule, rules::Rules}; +pub use crate::internals::{CallbackMsg, CallbackReturn}; /// A wrapper around compiled [Rules], with its own set of external variables, flags and timeout. /// @@ -131,7 +133,26 @@ impl<'rules> Scanner<'rules> { } internals::CallbackReturn::Continue }; - internals::scanner_scan_mem(self.inner, mem, callback).map(|_| results) + self.scan_mem_callback(mem, callback).map(|_| results) + } + + /// Scan memory with custom callback + /// + /// Returns + /// + /// * `mem` - Slice to scan + /// * `callback` - YARA callback mor read [here](https://yara.readthedocs.io/en/stable/capi.html#scanning-data) + /// + /// # Ownership + /// + /// This funciton takes the Scanner as `&mut` because it modifies the + /// `scanner->callback` and `scanner->user_data`, which are not behind a Mutex. + pub fn scan_mem_callback<'r>( + &mut self, + mem: &[u8], + callback: impl FnMut(CallbackMsg<'r>) -> CallbackReturn, + ) -> Result<(), YaraError> { + internals::scanner_scan_mem(self.inner, mem, callback) } /// Scan a file. @@ -142,20 +163,38 @@ impl<'rules> Scanner<'rules> { /// /// This function takes the Scanner as `&mut` because it modifies the /// `scanner->callback` and `scanner->user_data`, which are not behind a Mutex. - pub fn scan_file<'r, P: AsRef>(&self, path: P) -> Result>, Error> { + pub fn scan_file<'r, P: AsRef>(&mut self, path: P) -> Result>, Error> { + let mut results: Vec = Vec::new(); + let callback = |message| { + if let internals::CallbackMsg::RuleMatching(rule) = message { + results.push(rule) + } + internals::CallbackReturn::Continue + }; + + self.scan_file_callback(path, callback).map(|_| results) + } + + /// Scan file with custom callback + /// + /// Returns + /// + /// * `path` - Path to file + /// * `callback` - YARA callback mor read [here](https://yara.readthedocs.io/en/stable/capi.html#scanning-data) + /// + /// # Ownership + /// + /// This function takes the Scanner as `&mut` because it modifies the + /// `scanner->callback` and `scanner->user_data`, which are not behind a Mutex. + pub fn scan_file_callback<'r, P: AsRef>( + &mut self, + path: P, + callback: impl FnMut(CallbackMsg<'r>) -> CallbackReturn, + ) -> Result<(), Error> { File::open(path) .map_err(|e| IoError::new(e, IoErrorKind::OpenScanFile).into()) .and_then(|file| { - let mut results: Vec = Vec::new(); - let callback = |message| { - if let internals::CallbackMsg::RuleMatching(rule) = message { - results.push(rule) - } - internals::CallbackReturn::Continue - }; - internals::scanner_scan_file(self.inner, &file, callback) - .map_err(|e| e.into()) - .map(|_| results) + internals::scanner_scan_file(self.inner, &file, callback).map_err(|e| e.into()) }) } @@ -179,7 +218,30 @@ impl<'rules> Scanner<'rules> { } internals::CallbackReturn::Continue }; - internals::scanner_scan_proc(self.inner, pid, callback).map(|_| results) + self.scan_process_callback(pid, callback).map(|_| results) + } + + /// Attach a process, pause it, and scan its memory. + /// + /// Returns + /// + /// * `pid` - Process id + /// * `callback` - YARA callback mor read [here](https://yara.readthedocs.io/en/stable/capi.html#scanning-data) + /// + /// # Permissions + /// + /// You need to be able to attach to process `pid`. + /// + /// # Ownership + /// + /// This function takes the Scanner as `&mut` because it modifies the + /// `scanner->callback` and `scanner->user_data`, which are not behind a Mutex. + pub fn scan_process_callback<'r>( + &mut self, + pid: u32, + callback: impl FnMut(CallbackMsg<'r>) -> CallbackReturn, + ) -> Result<(), YaraError> { + internals::scanner_scan_proc(self.inner, pid, callback) } /// Set the maximum number of seconds that the scanner will spend in any call