diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..5b5e090 --- /dev/null +++ b/src/error.rs @@ -0,0 +1,51 @@ +use windows::core::PSTR; +use windows::Win32::Foundation::{GetLastError, WIN32_ERROR}; +use windows::Win32::System::Diagnostics::Debug::{ + FormatMessageA, FORMAT_MESSAGE_FROM_SYSTEM, FORMAT_MESSAGE_IGNORE_INSERTS, +}; + +/// The errors that may occur in blondie. +#[derive(Debug)] +pub enum Error { + /// Blondie requires administrator privileges + NotAnAdmin, + /// Error writing to the provided Writer + Write(std::io::Error), + /// Error spawning a suspended process + SpawnErr(std::io::Error), + /// Error waiting for child, abandoned + WaitOnChildErrAbandoned, + /// Error waiting for child, timed out + WaitOnChildErrTimeout, + /// A call to a windows API function returned an error and we didn't know how to handle it + Other(WIN32_ERROR, String, &'static str), + /// We require Windows 7 or greater + UnsupportedOsVersion, + /// This should never happen + UnknownError, +} +/// A [`std::result::Result`] alias where the `Err` case is [`blondie::Error`](Error). +pub type Result = std::result::Result; + +pub(crate) fn get_last_error(extra: &'static str) -> Error { + const BUF_LEN: usize = 1024; + let mut buf = [0u8; BUF_LEN]; + let code = unsafe { GetLastError() }; + let chars_written = unsafe { + FormatMessageA( + FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS, + None, + code.0, + 0, + PSTR(buf.as_mut_ptr()), + BUF_LEN as u32, + None, + ) + }; + assert!(chars_written != 0); + let code_str = std::ffi::CStr::from_bytes_until_nul(&buf) + .unwrap() + .to_str() + .unwrap_or("Invalid utf8 in error"); + Error::Other(code, code_str.to_string(), extra) +} diff --git a/src/lib.rs b/src/lib.rs index 1779b6e..4953b1e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -8,6 +8,9 @@ #![warn(missing_docs)] #![allow(clippy::field_reassign_with_default)] +mod error; +mod util; + use std::ffi::OsString; use std::io::{Read, Write}; use std::mem::size_of; @@ -21,15 +24,7 @@ use pdb_addr2line::pdb::PDB; use pdb_addr2line::ContextPdbData; use windows::core::{GUID, PCSTR, PSTR}; use windows::Win32::Foundation::{ - CloseHandle, GetLastError, ERROR_SUCCESS, ERROR_WMI_INSTANCE_NOT_FOUND, HANDLE, - INVALID_HANDLE_VALUE, WIN32_ERROR, -}; -use windows::Win32::Security::{ - AdjustTokenPrivileges, LookupPrivilegeValueW, SE_PRIVILEGE_ENABLED, TOKEN_ADJUST_PRIVILEGES, - TOKEN_PRIVILEGES, -}; -use windows::Win32::System::Diagnostics::Debug::{ - FormatMessageA, FORMAT_MESSAGE_FROM_SYSTEM, FORMAT_MESSAGE_IGNORE_INSERTS, + CloseHandle, ERROR_SUCCESS, ERROR_WMI_INSTANCE_NOT_FOUND, HANDLE, INVALID_HANDLE_VALUE, }; use windows::Win32::System::Diagnostics::Etw::{ CloseTrace, ControlTraceA, OpenTraceA, ProcessTrace, StartTraceA, SystemTraceControlGuid, @@ -41,12 +36,13 @@ use windows::Win32::System::Diagnostics::Etw::{ WNODE_FLAG_TRACED_GUID, }; use windows::Win32::System::SystemInformation::{GetVersionExA, OSVERSIONINFOA}; -use windows::Win32::System::SystemServices::SE_SYSTEM_PROFILE_NAME; use windows::Win32::System::Threading::{ - GetCurrentProcess, GetCurrentThread, OpenProcess, OpenProcessToken, SetThreadPriority, - WaitForSingleObject, CREATE_SUSPENDED, PROCESS_ALL_ACCESS, THREAD_PRIORITY_TIME_CRITICAL, + GetCurrentThread, SetThreadPriority, CREATE_SUSPENDED, THREAD_PRIORITY_TIME_CRITICAL, }; +use crate::error::get_last_error; +pub use crate::error::{Error, Result}; + /// Maximum stack depth/height of traces. // msdn says 192 but I got some that were bigger // const MAX_STACK_DEPTH: usize = 192; @@ -105,102 +101,6 @@ impl Drop for TraceContext { } } -/// The errors that may occur in blondie. -#[derive(Debug)] -pub enum Error { - /// Blondie requires administrator privileges - NotAnAdmin, - /// Error writing to the provided Writer - Write(std::io::Error), - /// Error spawning a suspended process - SpawnErr(std::io::Error), - /// Error waiting for child, abandoned - WaitOnChildErrAbandoned, - /// Error waiting for child, timed out - WaitOnChildErrTimeout, - /// A call to a windows API function returned an error and we didn't know how to handle it - Other(WIN32_ERROR, String, &'static str), - /// We require Windows 7 or greater - UnsupportedOsVersion, - /// This should never happen - UnknownError, -} -/// A [`std::result::Result`] alias where the `Err` case is [`blondie::Error`](Error). -pub type Result = std::result::Result; - -fn get_last_error(extra: &'static str) -> Error { - const BUF_LEN: usize = 1024; - let mut buf = [0u8; BUF_LEN]; - let code = unsafe { GetLastError() }; - let chars_written = unsafe { - FormatMessageA( - FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS, - None, - code.0, - 0, - PSTR(buf.as_mut_ptr()), - BUF_LEN as u32, - None, - ) - }; - assert!(chars_written != 0); - let code_str = std::ffi::CStr::from_bytes_until_nul(&buf) - .unwrap() - .to_str() - .unwrap_or("Invalid utf8 in error"); - Error::Other(code, code_str.to_string(), extra) -} - -/// A wrapper around `OpenProcess` that returns a handle with all access rights -unsafe fn handle_from_process_id(process_id: u32) -> Result { - match OpenProcess(PROCESS_ALL_ACCESS, false, process_id) { - Ok(handle) => Ok(handle), - Err(_) => Err(get_last_error("handle_from_process_id")), - } -} - -unsafe fn wait_for_process_by_handle(handle: HANDLE) -> Result<()> { - let ret = WaitForSingleObject(handle, 0xffffffff); - match ret.0 { - 0 => Ok(()), - 0x00000080 => Err(Error::WaitOnChildErrAbandoned), - 0x00000102 => Err(Error::WaitOnChildErrTimeout), - _ => Err(get_last_error("wait_for_process_by_handle")), - } -} - -fn acquire_privileges() -> Result<()> { - let mut privs = TOKEN_PRIVILEGES::default(); - privs.PrivilegeCount = 1; - privs.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED; - if unsafe { - LookupPrivilegeValueW(None, SE_SYSTEM_PROFILE_NAME, &mut privs.Privileges[0].Luid).0 == 0 - } { - return Err(get_last_error("acquire_privileges LookupPrivilegeValueA")); - } - let mut pt = HANDLE::default(); - if unsafe { OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &mut pt).0 == 0 } { - return Err(get_last_error("OpenProcessToken")); - } - let adjust = unsafe { AdjustTokenPrivileges(pt, false, Some(addr_of!(privs)), 0, None, None) }; - if adjust.0 == 0 { - let err = Err(get_last_error("AdjustTokenPrivileges")); - unsafe { - CloseHandle(pt); - } - return err; - } - let ret = unsafe { CloseHandle(pt) }; - if ret.0 == 0 { - return Err(get_last_error("acquire_privileges CloseHandle")); - } - let status = unsafe { GetLastError() }; - if status != ERROR_SUCCESS { - return Err(Error::NotAnAdmin); - } - Ok(()) -} - /// The main tracing logic. Traces the process with the given `target_process_id`. /// /// # Safety @@ -224,7 +124,7 @@ unsafe fn trace_from_process_id( { return Err(Error::UnsupportedOsVersion); } - acquire_privileges()?; + util::acquire_privileges()?; // Set the sampling interval // Only for Win8 or more @@ -246,12 +146,8 @@ unsafe fn trace_from_process_id( } } - let mut kernel_logger_name_with_nul = KERNEL_LOGGER_NAMEA - .as_bytes() - .iter() - .cloned() - .chain(Some(0)) - .collect::>(); + let mut kernel_logger_name_with_nul = KERNEL_LOGGER_NAMEA.as_bytes().to_vec(); + kernel_logger_name_with_nul.push(b'\0'); // Build the trace properties, we want EVENT_TRACE_FLAG_PROFILE for the "SampledProfile" event // https://docs.microsoft.com/en-us/windows/win32/etw/sampledprofile // In https://docs.microsoft.com/en-us/windows/win32/etw/event-tracing-mof-classes that event is listed as a "kernel event" @@ -353,7 +249,7 @@ unsafe fn trace_from_process_id( } } - let target_proc_handle = handle_from_process_id(target_process_id)?; + let target_proc_handle = util::handle_from_process_id(target_process_id)?; let mut context = TraceContext::new(target_proc_handle, target_process_id, kernel_stacks)?; // TODO: Do we need to Box the context? @@ -533,7 +429,7 @@ unsafe fn trace_from_process_id( } // Wait for it to end - wait_for_process_by_handle(target_proc_handle)?; + util::wait_for_process_by_handle(target_proc_handle)?; // This unblocks ProcessTrace let ret = ControlTraceA( ::default(), @@ -552,7 +448,7 @@ unsafe fn trace_from_process_id( .map_err(|_err_any| Error::UnknownError)??; if context.show_kernel_samples { - let kernel_module_paths = list_kernel_modules(); + let kernel_module_paths = util::list_kernel_modules(); context.image_paths.extend(kernel_module_paths); } @@ -869,80 +765,3 @@ struct ImageLoadEvent { Reserved3: u32, Reserved4: u32, } - -/// Returns a sequence of (image_file_path, image_base) -fn list_kernel_modules() -> Vec<(OsString, u64, u64)> { - // kernel module enumeration code based on http://www.rohitab.com/discuss/topic/40696-list-loaded-drivers-with-ntquerysysteminformation/ - #[link(name = "ntdll")] - extern "system" { - fn NtQuerySystemInformation( - SystemInformationClass: u32, - SystemInformation: *mut (), - SystemInformationLength: u32, - ReturnLength: *mut u32, - ) -> i32; - } - - const BUF_LEN: usize = 1024 * 1024; - let mut out_buf = vec![0u8; BUF_LEN]; - let mut out_size = 0u32; - // 11 = SystemModuleInformation - let retcode = unsafe { - NtQuerySystemInformation( - 11, - out_buf.as_mut_ptr().cast(), - BUF_LEN as u32, - &mut out_size, - ) - }; - if retcode < 0 { - // println!("Failed to load kernel modules"); - return vec![]; - } - let number_of_modules = unsafe { out_buf.as_ptr().cast::().read_unaligned() as usize }; - #[repr(C)] - #[derive(Debug)] - #[allow(non_snake_case)] - #[allow(non_camel_case_types)] - struct _RTL_PROCESS_MODULE_INFORMATION { - Section: *mut std::ffi::c_void, - MappedBase: *mut std::ffi::c_void, - ImageBase: *mut std::ffi::c_void, - ImageSize: u32, - Flags: u32, - LoadOrderIndex: u16, - InitOrderIndex: u16, - LoadCount: u16, - OffsetToFileName: u16, - FullPathName: [u8; 256], - } - let modules = unsafe { - let modules_ptr = out_buf - .as_ptr() - .cast::() - .offset(2) - .cast::<_RTL_PROCESS_MODULE_INFORMATION>(); - std::slice::from_raw_parts(modules_ptr, number_of_modules) - }; - - let kernel_module_paths = modules - .iter() - .filter_map(|module| { - std::ffi::CStr::from_bytes_until_nul(&module.FullPathName) - .unwrap() - .to_str() - .ok() - .map(|mod_str_filepath| { - let verbatim_path_osstring: OsString = mod_str_filepath - .replacen("\\SystemRoot\\", "\\\\?\\C:\\Windows\\", 1) - .into(); - ( - verbatim_path_osstring, - module.ImageBase as u64, - module.ImageSize as u64, - ) - }) - }) - .collect(); - kernel_module_paths -} diff --git a/src/util.rs b/src/util.rs new file mode 100644 index 0000000..40e3c42 --- /dev/null +++ b/src/util.rs @@ -0,0 +1,141 @@ +use crate::{get_last_error, Result, Error}; + +use std::ffi::OsString; +use std::ptr::addr_of; + +use windows::Win32::Foundation::{CloseHandle, GetLastError, ERROR_SUCCESS, HANDLE}; +use windows::Win32::Security::{ + AdjustTokenPrivileges, LookupPrivilegeValueW, SE_PRIVILEGE_ENABLED, TOKEN_ADJUST_PRIVILEGES, + TOKEN_PRIVILEGES, +}; +use windows::Win32::System::SystemServices::SE_SYSTEM_PROFILE_NAME; +use windows::Win32::System::Threading::{ + GetCurrentProcess, OpenProcess, OpenProcessToken, WaitForSingleObject, PROCESS_ALL_ACCESS, +}; + +/// A wrapper around `OpenProcess` that returns a handle with all access rights +pub(crate) unsafe fn handle_from_process_id(process_id: u32) -> Result { + match OpenProcess(PROCESS_ALL_ACCESS, false, process_id) { + Ok(handle) => Ok(handle), + Err(_) => Err(get_last_error("handle_from_process_id")), + } +} + +pub(crate) unsafe fn wait_for_process_by_handle(handle: HANDLE) -> Result<()> { + let ret = WaitForSingleObject(handle, 0xffffffff); + match ret.0 { + 0 => Ok(()), + 0x00000080 => Err(Error::WaitOnChildErrAbandoned), + 0x00000102 => Err(Error::WaitOnChildErrTimeout), + _ => Err(get_last_error("wait_for_process_by_handle")), + } +} + +pub(crate) fn acquire_privileges() -> Result<()> { + let mut privs = TOKEN_PRIVILEGES::default(); + privs.PrivilegeCount = 1; + privs.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED; + if unsafe { + LookupPrivilegeValueW(None, SE_SYSTEM_PROFILE_NAME, &mut privs.Privileges[0].Luid).0 == 0 + } { + return Err(get_last_error("acquire_privileges LookupPrivilegeValueA")); + } + let mut pt = HANDLE::default(); + if unsafe { OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &mut pt).0 == 0 } { + return Err(get_last_error("OpenProcessToken")); + } + let adjust = unsafe { AdjustTokenPrivileges(pt, false, Some(addr_of!(privs)), 0, None, None) }; + if adjust.0 == 0 { + let err = Err(get_last_error("AdjustTokenPrivileges")); + unsafe { + CloseHandle(pt); + } + return err; + } + let ret = unsafe { CloseHandle(pt) }; + if ret.0 == 0 { + return Err(get_last_error("acquire_privileges CloseHandle")); + } + let status = unsafe { GetLastError() }; + if status != ERROR_SUCCESS { + return Err(Error::NotAnAdmin); + } + Ok(()) +} + +/// Returns a sequence of (image_file_path, image_base, image_size) +pub(crate) fn list_kernel_modules() -> Vec<(OsString, u64, u64)> { + // kernel module enumeration code based on http://www.rohitab.com/discuss/topic/40696-list-loaded-drivers-with-ntquerysysteminformation/ + #[link(name = "ntdll")] + extern "system" { + fn NtQuerySystemInformation( + SystemInformationClass: u32, + SystemInformation: *mut (), + SystemInformationLength: u32, + ReturnLength: *mut u32, + ) -> i32; + } + + const BUF_LEN: usize = 1024 * 1024; + let mut out_buf = vec![0u8; BUF_LEN]; + let mut out_size = 0u32; + // 11 = SystemModuleInformation + let retcode = unsafe { + NtQuerySystemInformation( + 11, + out_buf.as_mut_ptr().cast(), + BUF_LEN as u32, + &mut out_size, + ) + }; + if retcode < 0 { + // println!("Failed to load kernel modules"); + return vec![]; + } + let number_of_modules = unsafe { out_buf.as_ptr().cast::().read_unaligned() as usize }; + #[repr(C)] + #[derive(Debug)] + #[allow(non_snake_case)] + #[allow(non_camel_case_types)] + struct _RTL_PROCESS_MODULE_INFORMATION { + Section: *mut std::ffi::c_void, + MappedBase: *mut std::ffi::c_void, + ImageBase: *mut std::ffi::c_void, + ImageSize: u32, + Flags: u32, + LoadOrderIndex: u16, + InitOrderIndex: u16, + LoadCount: u16, + OffsetToFileName: u16, + FullPathName: [u8; 256], + } + let modules = unsafe { + let modules_ptr = out_buf + .as_ptr() + .cast::() + .offset(2) + .cast::<_RTL_PROCESS_MODULE_INFORMATION>(); + std::slice::from_raw_parts(modules_ptr, number_of_modules) + }; + + let kernel_module_paths = modules + .iter() + .filter_map(|module| { + std::ffi::CStr::from_bytes_until_nul(&module.FullPathName) + .unwrap() + .to_str() + .ok() + .map(|mod_str_filepath| { + let verbatim_path_osstring: OsString = mod_str_filepath + .replacen("\\SystemRoot\\", "\\\\?\\C:\\Windows\\", 1) + .into(); + ( + verbatim_path_osstring, + module.ImageBase as u64, + module.ImageSize as u64, + ) + }) + }) + .collect(); + kernel_module_paths +}