Skip to content

Commit

Permalink
Rearrange some code into modules
Browse files Browse the repository at this point in the history
  • Loading branch information
@MingweiSamuel
MingweiSamuel committed Aug 20, 2024
1 parent e1baeb4 commit 73f7667
Show file tree
Hide file tree
Showing 3 changed files with 206 additions and 195 deletions.
51 changes: 51 additions & 0 deletions src/error.rs
View file
Original file line number Diff line number Diff line change
@@ -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<T> = std::result::Result<T, Error>;

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)
}
209 changes: 14 additions & 195 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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,
Expand All @@ -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;
Expand Down Expand Up @@ -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<T> = std::result::Result<T, Error>;

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<HANDLE> {
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
Expand All @@ -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
Expand All @@ -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::<Vec<u8>>();
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"
Expand Down Expand Up @@ -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?

Expand Down Expand Up @@ -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(
<CONTROLTRACE_HANDLE as Default>::default(),
Expand All @@ -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);
}

Expand Down Expand Up @@ -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::<u32>().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::<u32>()
.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
}
Loading

0 comments on commit 73f7667

Please sign in to comment.