diff --git a/src/agent/debugger/src/debugger.rs b/src/agent/debugger/src/debugger.rs index 87605898e..38b1da603 100644 --- a/src/agent/debugger/src/debugger.rs +++ b/src/agent/debugger/src/debugger.rs @@ -8,9 +8,8 @@ #![allow(clippy::redundant_closure)] #![allow(clippy::redundant_clone)] use std::{ - collections::{hash_map, HashMap}, + collections::HashMap, ffi::OsString, - fs, mem::MaybeUninit, os::windows::process::CommandExt, path::{Path, PathBuf}, @@ -19,32 +18,29 @@ use std::{ use anyhow::{Context, Result}; use log::{error, trace}; -use win_util::{check_winapi, file, last_os_error, process}; +use win_util::{check_winapi, last_os_error, process}; use winapi::{ shared::{ - minwindef::{DWORD, FALSE, LPCVOID, LPVOID, TRUE}, + minwindef::{DWORD, FALSE, LPCVOID, TRUE}, winerror::ERROR_SEM_TIMEOUT, }, um::{ dbghelp::ADDRESS64, debugapi::{ContinueDebugEvent, WaitForDebugEvent}, errhandlingapi::GetLastError, - handleapi::CloseHandle, minwinbase::{ CREATE_PROCESS_DEBUG_INFO, CREATE_THREAD_DEBUG_INFO, EXCEPTION_BREAKPOINT, EXCEPTION_DEBUG_INFO, EXCEPTION_SINGLE_STEP, EXIT_PROCESS_DEBUG_INFO, EXIT_THREAD_DEBUG_INFO, LOAD_DLL_DEBUG_INFO, RIP_INFO, UNLOAD_DLL_DEBUG_INFO, }, winbase::{DebugSetProcessKillOnExit, DEBUG_ONLY_THIS_PROCESS, INFINITE}, - winnt::{ - DBG_CONTINUE, DBG_EXCEPTION_NOT_HANDLED, HANDLE, IMAGE_FILE_MACHINE_AMD64, - IMAGE_FILE_MACHINE_I386, - }, + winnt::{DBG_CONTINUE, DBG_EXCEPTION_NOT_HANDLED, HANDLE}, }, }; +use crate::target::Target; use crate::{ - dbghelp::{self, FrameContext, SymInfo}, + dbghelp::{self, SymInfo}, debug_event::{DebugEvent, DebugEventInfo}, stack, }; @@ -57,7 +53,7 @@ const STATUS_WX86_BREAKPOINT: u32 = ::winapi::shared::ntstatus::STATUS_WX86_BREA pub struct BreakpointId(pub u32); #[derive(Copy, Clone)] -enum StepState { +pub(crate) enum StepState { Breakpoint { pc: u64 }, SingleStep, } @@ -69,12 +65,26 @@ pub enum BreakpointType { StepOut { rsp: u64 }, } -struct ModuleBreakpoint { +pub(crate) struct ModuleBreakpoint { rva: u64, kind: BreakpointType, id: BreakpointId, } +impl ModuleBreakpoint { + pub fn rva(&self) -> u64 { + self.rva + } + + pub fn kind(&self) -> BreakpointType { + self.kind + } + + pub fn id(&self) -> BreakpointId { + self.id + } +} + #[allow(unused)] struct UnresolvedBreakpoint { sym: String, @@ -85,7 +95,7 @@ struct UnresolvedBreakpoint { /// A breakpoint for a specific target. We can say it is bound because we know exactly /// where to set it, but it might disabled. #[derive(Clone)] -struct Breakpoint { +pub struct Breakpoint { /// The address of the breakpoint. ip: u64, @@ -102,6 +112,70 @@ struct Breakpoint { id: BreakpointId, } +impl Breakpoint { + pub fn new( + ip: u64, + kind: BreakpointType, + enabled: bool, + original_byte: Option, + hit_count: usize, + id: BreakpointId, + ) -> Self { + Breakpoint { + ip, + kind, + enabled, + original_byte, + hit_count, + id, + } + } + + pub fn ip(&self) -> u64 { + self.ip + } + + pub fn kind(&self) -> BreakpointType { + self.kind + } + + pub(crate) fn set_kind(&mut self, kind: BreakpointType) { + self.kind = kind; + } + + pub fn enabled(&self) -> bool { + self.enabled + } + + pub fn set_enabled(&mut self, enabled: bool) { + self.enabled = enabled; + } + + pub fn original_byte(&self) -> Option { + self.original_byte + } + + pub(crate) fn set_original_byte(&mut self, original_byte: Option) { + self.original_byte = original_byte; + } + + pub fn hit_count(&self) -> usize { + self.hit_count + } + + pub(crate) fn increment_hit_count(&mut self) { + self.hit_count += 1; + } + + pub fn id(&self) -> BreakpointId { + self.id + } + + pub(crate) fn set_id(&mut self, id: BreakpointId) { + self.id = id; + } +} + pub struct StackFrame { return_address: u64, stack_pointer: u64, @@ -124,458 +198,6 @@ impl StackFrame { } } -struct Module { - path: PathBuf, - file_handle: HANDLE, - base_address: u64, - image_size: u32, - machine: Machine, - - // Track if we need to call SymLoadModule for the dll. - sym_module_loaded: bool, -} - -impl Module { - fn new(module_handle: HANDLE, base_address: u64) -> Result { - let path = file::get_path_from_handle(module_handle).unwrap_or_else(|e| { - error!("Error getting path from file handle: {}", e); - "???".into() - }); - - let image_details = get_image_details(&path)?; - - Ok(Self { - path, - file_handle: module_handle, - base_address, - image_size: image_details.image_size, - machine: image_details.machine, - sym_module_loaded: false, - }) - } - - fn sym_load_module(&mut self, process_handle: HANDLE) -> Result<()> { - if !self.sym_module_loaded { - let dbghelp = dbghelp::lock()?; - - dbghelp.sym_load_module( - process_handle, - self.file_handle, - &self.path, - self.base_address, - self.image_size, - )?; - - self.sym_module_loaded = true; - } - - Ok(()) - } - - fn name(&self) -> &Path { - // Unwrap guaranteed by construction, we always have a filename. - self.path.file_stem().unwrap().as_ref() - } -} - -impl Drop for Module { - fn drop(&mut self) { - unsafe { CloseHandle(self.file_handle) }; - } -} - -pub struct Target { - process_id: DWORD, - process_handle: HANDLE, - current_thread_handle: HANDLE, - saw_initial_bp: bool, - saw_initial_wow64_bp: bool, - - // Track if we need to call SymInitialize for the process and if we need to notify - // dbghelp about loaded/unloaded dlls. - sym_initialized: bool, - exited: bool, - - thread_handles: fnv::FnvHashMap, - - // We cache the current thread context for possible repeated queries and modifications. - // We want to call GetThreadContext once, then call SetThreadContext (if necessary) before - // resuming. Calling Get/Set/Get/Set doesn't seem to work because the second Get doesn't - // see any the changes made in the Set call. - current_context: Option, - - // True if we need to set the thread context before resuming. - context_is_modified: bool, - - // Key is base address (which also happens to be the HANDLE). - modules: fnv::FnvHashMap, - - breakpoints: fnv::FnvHashMap, - - // Map of thread to stepping state (e.g. breakpoint address to restore breakpoints) - single_step: fnv::FnvHashMap, -} - -impl Target { - fn new( - process_id: DWORD, - thread_id: DWORD, - process_handle: HANDLE, - thread_handle: HANDLE, - ) -> Self { - let mut thread_handles = fnv::FnvHashMap::default(); - thread_handles.insert(thread_id, thread_handle); - - Self { - process_id, - current_thread_handle: thread_handle, - process_handle, - saw_initial_bp: false, - saw_initial_wow64_bp: false, - sym_initialized: false, - exited: false, - thread_handles, - current_context: None, - context_is_modified: false, - modules: fnv::FnvHashMap::default(), - breakpoints: fnv::FnvHashMap::default(), - single_step: fnv::FnvHashMap::default(), - } - } - - fn modules(&self) -> hash_map::Iter { - self.modules.iter() - } - - fn initial_bp(&mut self, load_symbols: bool) -> Result<()> { - self.saw_initial_bp = true; - - if load_symbols || !self.breakpoints.is_empty() { - self.sym_initialize()?; - - for (_, module) in self.modules.iter_mut() { - if let Err(e) = module.sym_load_module(self.process_handle) { - error!("Error loading symbols: {}", e); - } - } - } - - Ok(()) - } - - fn sym_initialize(&mut self) -> Result<()> { - if !self.sym_initialized { - let dbghelp = dbghelp::lock()?; - if let Err(e) = dbghelp.sym_initialize(self.process_handle) { - error!("Error in SymInitializeW: {}", e); - - if let Err(e) = dbghelp.sym_cleanup(self.process_handle) { - error!("Error in SymCleanup: {}", e); - } - - return Err(e); - } - - for (_, module) in self.modules.iter_mut() { - if let Err(e) = module.sym_load_module(self.process_handle) { - error!( - "Error loading symbols for module {}: {}", - module.path.display(), - e - ); - } - } - - self.sym_initialized = true; - } - - Ok(()) - } - - /// Register the module loaded at `base_address`, returning the module name. - fn load_module(&mut self, module_handle: HANDLE, base_address: u64) -> Result> { - let mut module = Module::new(module_handle, base_address)?; - - trace!( - "Loading module {} at {:x}", - module.name().display(), - base_address - ); - - if module.machine == Machine::X64 && process::is_wow64_process(self.process_handle) { - // We ignore native dlls in wow64 processes. - return Ok(None); - } - - let module_name = module.name().to_owned(); - if self.sym_initialized { - if let Err(e) = module.sym_load_module(self.process_handle) { - error!("Error loading symbols: {}", e); - } - } - - let base_address = module.base_address; - if let Some(old_value) = self.modules.insert(base_address, module) { - error!( - "Existing module {} replace at base_address {}", - old_value.path.display(), - base_address - ); - } - - Ok(Some(module_name)) - } - - fn unload_module(&mut self, base_address: u64) { - // Drop the module and remove any breakpoints. - if let Some(module) = self.modules.remove(&base_address) { - let image_size = module.image_size as u64; - self.breakpoints - .retain(|&ip, _| ip < base_address || ip >= base_address + image_size); - } - } - - fn apply_absolute_breakpoint( - &mut self, - address: u64, - kind: BreakpointType, - id: BreakpointId, - ) -> Result<()> { - let original_byte: u8 = process::read_memory(self.process_handle, address as LPVOID)?; - - self.breakpoints - .entry(address) - .and_modify(|e| { - e.kind = kind; - e.enabled = true; - e.original_byte = Some(original_byte); - e.id = id; - }) - .or_insert(Breakpoint { - ip: address, - kind, - enabled: true, - original_byte: Some(original_byte), - hit_count: 0, - id, - }); - - write_instruction_byte(self.process_handle, address, 0xcc)?; - - Ok(()) - } - - fn apply_module_breakpoints( - &mut self, - base_address: u64, - breakpoints: &[ModuleBreakpoint], - ) -> Result<()> { - if breakpoints.is_empty() { - return Ok(()); - } - - // We want to set every breakpoint for the module at once. We'll read the just the - // memory we need to do that, so find the min and max rva to compute how much memory - // to read and update in the remote process. - let (min, max) = breakpoints - .iter() - .fold((u64::max_value(), u64::min_value()), |acc, bp| { - (acc.0.min(bp.rva), acc.1.max(bp.rva)) - }); - - // Add 1 to include the final byte. - let region_size = (max - min) - .checked_add(1) - .ok_or_else(|| anyhow::anyhow!("overflow in region size trying to set breakpoints"))? - as usize; - let remote_address = base_address.checked_add(min).ok_or_else(|| { - anyhow::anyhow!("overflow in remote address calculation trying to set breakpoints") - })? as LPVOID; - - let mut buffer: Vec = Vec::with_capacity(region_size); - unsafe { - buffer.set_len(region_size); - } - process::read_memory_array(self.process_handle, remote_address, &mut buffer[..])?; - - for mbp in breakpoints { - let ip = base_address + mbp.rva; - let offset = (mbp.rva - min) as usize; - - trace!("Setting breakpoint at {:x}", ip); - - let bp = Breakpoint { - ip, - kind: mbp.kind, - enabled: true, - original_byte: Some(buffer[offset]), - hit_count: 0, - id: mbp.id, - }; - - buffer[offset] = 0xcc; - - self.breakpoints.insert(ip, bp); - } - - process::write_memory_slice(self.process_handle, remote_address, &buffer[..])?; - process::flush_instruction_cache(self.process_handle, remote_address, region_size)?; - - Ok(()) - } - - fn prepare_to_resume(&mut self) -> Result<()> { - if let Some(context) = self.current_context.take() { - if self.context_is_modified { - context.set_thread_context(self.current_thread_handle)?; - } - } - - self.context_is_modified = false; - - Ok(()) - } - - fn ensure_current_context(&mut self) -> Result<()> { - if self.current_context.is_none() { - self.current_context = Some(dbghelp::get_thread_frame( - self.process_handle, - self.current_thread_handle, - )?); - } - - Ok(()) - } - - fn get_current_context(&mut self) -> Result<&FrameContext> { - self.ensure_current_context()?; - Ok(self.current_context.as_ref().unwrap()) - } - - fn get_current_context_mut(&mut self) -> Result<&mut FrameContext> { - self.ensure_current_context()?; - - // Assume the caller will modify the context. When it is modified, - // we must set it before resuming the target. - self.context_is_modified = true; - Ok(self.current_context.as_mut().unwrap()) - } - - fn read_register_u64(&mut self, reg: iced_x86::Register) -> Result { - let current_context = self.get_current_context()?; - Ok(current_context.get_register_u64(reg)) - } - - fn read_flags_register(&mut self) -> Result { - let current_context = self.get_current_context()?; - Ok(current_context.get_flags()) - } - - /// Handle a breakpoint that we set (as opposed to a breakpoint in user code, e.g. - /// assertion.) - /// - /// Return the breakpoint id if it should be reported to the client. - fn handle_breakpoint(&mut self, pc: u64) -> Result> { - enum HandleBreakpoint { - User(BreakpointId, bool), - StepOut(u64), - } - - let handle_breakpoint = { - let bp = self.breakpoints.get_mut(&pc).unwrap(); - - bp.hit_count += 1; - - write_instruction_byte(self.process_handle, bp.ip, bp.original_byte.unwrap())?; - - match bp.kind { - BreakpointType::OneTime => { - bp.enabled = false; - bp.original_byte = None; - - // We are clearing the breakpoint after hitting it, so we do not need - // to single step. - HandleBreakpoint::User(bp.id, false) - } - - BreakpointType::Counter => { - // Single step so we can restore the breakpoint after stepping. - HandleBreakpoint::User(bp.id, true) - } - - BreakpointType::StepOut { rsp } => HandleBreakpoint::StepOut(rsp), - } - }; - - let context = self.get_current_context_mut()?; - context.set_program_counter(pc); - - // We need to single step if we need to restore the breakpoint. - let single_step = match handle_breakpoint { - HandleBreakpoint::User(_, single_step) => single_step, - - // Single step only when in a recursive call, which is inferred when the current - // stack pointer (from context) is less then at the target to step out from (rsp). - // Note this only works if the stack grows down. - HandleBreakpoint::StepOut(rsp) => rsp > context.stack_pointer(), - }; - - if single_step { - context.set_single_step(true); - self.single_step - .insert(self.current_thread_handle, StepState::Breakpoint { pc }); - } - - Ok(match handle_breakpoint { - HandleBreakpoint::User(id, _) => Some(id), - HandleBreakpoint::StepOut(_) => None, - }) - } - - fn handle_single_step(&mut self, step_state: StepState) -> Result<()> { - match step_state { - StepState::Breakpoint { pc } => { - write_instruction_byte(self.process_handle, pc, 0xcc)?; - } - _ => {} - } - - self.single_step.remove(&self.current_thread_handle); - Ok(()) - } - - fn prepare_to_step(&mut self) -> Result { - // Don't change the reason we're single stepping on this thread if - // we previously set the reason (e.g. so we would restore a breakpoint). - self.single_step - .entry(self.current_thread_handle) - .or_insert(StepState::SingleStep); - - let context = self.get_current_context_mut()?; - context.set_single_step(true); - - Ok(true) - } - - fn set_exited(&mut self) -> Result<()> { - self.exited = true; - - if self.sym_initialized { - let dbghelp = dbghelp::lock()?; - dbghelp.sym_cleanup(self.process_handle)?; - } - Ok(()) - } -} - -fn write_instruction_byte(process_handle: HANDLE, ip: u64, b: u8) -> Result<()> { - let orig_byte = [b; 1]; - let remote_address = ip as LPVOID; - process::write_memory_slice(process_handle, remote_address, &orig_byte)?; - process::flush_instruction_cache(process_handle, remote_address, orig_byte.len())?; - Ok(()) -} - #[rustfmt::skip] #[allow(clippy::trivially_copy_pass_by_ref)] pub trait DebugEventHandler { @@ -782,7 +404,7 @@ impl Debugger { } pub fn run(&mut self, callbacks: &mut impl DebugEventHandler) -> Result<()> { - while !self.target.exited { + while !self.target.exited() { // Poll between every event so a client can add generic logic instead needing to // handle every possible event. callbacks.on_poll(self); @@ -798,9 +420,9 @@ impl Debugger { } pub fn quit_debugging(&self) { - if !self.target.exited { - trace!("timeout - terminating pid: {}", self.target.process_id); - process::terminate(self.target.process_handle); + if !self.target.exited() { + trace!("timeout - terminating pid: {}", self.target.process_id()); + process::terminate(self.target.process_handle()); } } @@ -808,13 +430,9 @@ impl Debugger { let mut continue_status = DBG_CONTINUE; if let DebugEventInfo::CreateThread(info) = de.info() { - self.target.current_thread_handle = info.hThread; - self.target - .thread_handles - .insert(de.thread_id(), info.hThread); + self.target.create_new_thread(info.hThread, de.thread_id()); } else { - self.target.current_thread_handle = - *self.target.thread_handles.get(&de.thread_id()).unwrap(); + self.target.set_current_thread(de.thread_id()); } match de.info() { @@ -832,7 +450,7 @@ impl Debugger { // breakpoint notification from the OS. Otherwise we may set // breakpoints in startup code before the debugger is properly // initialized. - if self.target.saw_initial_bp { + if self.target.saw_initial_bp() { self.apply_module_breakpoints(module_name, base_address) } } @@ -874,7 +492,7 @@ impl Debugger { DebugEventInfo::ExitThread(info) => { callbacks.on_exit_thread(self, *info); - self.target.thread_handles.remove(&de.thread_id()); + self.target.exit_thread(de.thread_id()); } DebugEventInfo::OutputDebugString(info) => { @@ -882,7 +500,7 @@ impl Debugger { let length = info.nDebugStringLength.saturating_sub(1) as usize; if info.fUnicode != 0 { if let Ok(message) = process::read_wide_string( - self.target.process_handle, + self.target.process_handle(), info.lpDebugStringData as LPCVOID, length, ) { @@ -890,7 +508,7 @@ impl Debugger { } } else { if let Ok(message) = process::read_narrow_string( - self.target.process_handle, + self.target.process_handle(), info.lpDebugStringData as LPCVOID, length, ) { @@ -921,7 +539,7 @@ impl Debugger { ) { Some(DebuggerNotification::InitialBreak) => { let modules = { - self.target.saw_initial_bp = true; + self.target.set_saw_initial_bp(); let load_symbols = !self.symbolic_breakpoints.is_empty(); self.target.initial_bp(load_symbols)?; self.target @@ -935,7 +553,7 @@ impl Debugger { Ok(DBG_CONTINUE) } Some(DebuggerNotification::InitialWow64Break) => { - self.target.saw_initial_wow64_bp = true; + self.target.set_saw_initial_wow64_bp(); Ok(DBG_CONTINUE) } Some(DebuggerNotification::Clr) => Ok(DBG_CONTINUE), @@ -950,7 +568,7 @@ impl Debugger { Ok(DBG_CONTINUE) } None => { - let process_handle = self.target.process_handle; + let process_handle = self.target.process_handle(); Ok(callbacks.on_exception(self, info, process_handle)) } } @@ -972,7 +590,7 @@ impl Debugger { Ok(dbghelp) => { for bp in unresolved_breakpoints { match dbghelp.sym_from_name( - self.target.process_handle, + self.target.process_handle(), module_name.as_ref(), &bp.sym, ) { @@ -1014,19 +632,19 @@ impl Debugger { // be too aggressive in dealing with failures. let resolve_symbols = self.target.sym_initialize().is_ok(); return stack::get_stack( - self.target.process_handle, - self.target.current_thread_handle, + self.target.process_handle(), + self.target.current_thread_handle(), resolve_symbols, ); } pub fn get_symbol(&self, pc: u64) -> Result { let dbghelp = dbghelp::lock()?; - dbghelp.sym_from_inline_context(self.target.process_handle, pc, 0) + dbghelp.sym_from_inline_context(self.target.process_handle(), pc, 0) } pub fn get_current_thread_id(&self) -> u64 { - self.target.current_thread_handle as u64 + self.target.current_thread_handle() as u64 } pub fn read_register_u64(&mut self, reg: iced_x86::Register) -> Result { @@ -1042,7 +660,7 @@ impl Debugger { remote_address: LPCVOID, buf: &mut [T], ) -> Result<()> { - process::read_memory_array(self.target.process_handle, remote_address, buf)?; + process::read_memory_array(self.target.process_handle(), remote_address, buf)?; Ok(()) } @@ -1052,8 +670,8 @@ impl Debugger { let mut return_address = ADDRESS64::default(); let mut stack_pointer = ADDRESS64::default(); dbghlp.stackwalk_ex( - self.target.process_handle, - self.target.current_thread_handle, + self.target.process_handle(), + self.target.current_thread_handle(), |_frame_context, frame| { return_address = frame.AddrReturn; stack_pointer = frame.AddrStack; @@ -1116,8 +734,8 @@ fn is_debugger_notification( // The first EXCEPTION_BREAKPOINT is sent to debuggers like us. EXCEPTION_BREAKPOINT => { - if target.saw_initial_bp { - if target.breakpoints.contains_key(&exception_address) { + if target.saw_initial_bp() { + if target.breakpoint_set_at_addr(exception_address) { Some(DebuggerNotification::Breakpoint(exception_address)) } else { None @@ -1130,7 +748,7 @@ fn is_debugger_notification( // We may see a second breakpoint (STATUS_WX86_BREAKPOINT) when debugging a // WoW64 process, this is also a debugger notification, not a real breakpoint. STATUS_WX86_BREAKPOINT => { - if target.saw_initial_wow64_bp { + if target.saw_initial_wow64_bp() { None } else { Some(DebuggerNotification::InitialWow64Break) @@ -1138,7 +756,7 @@ fn is_debugger_notification( } EXCEPTION_SINGLE_STEP => { - if let Some(&step_state) = target.single_step.get(&target.current_thread_handle) { + if let Some(step_state) = target.single_step(target.current_thread_handle()) { Some(DebuggerNotification::SingleStep(step_state)) } else { // Unexpected single step - could be a logic bug in the debugger or less @@ -1152,37 +770,3 @@ fn is_debugger_notification( _ => None, } } - -#[derive(Copy, Clone, Debug, PartialEq)] -enum Machine { - Unknown, - X64, - X86, -} - -struct ImageDetails { - image_size: u32, - machine: Machine, -} - -fn get_image_details(path: &Path) -> Result { - let file = fs::File::open(path)?; - let map = unsafe { memmap::Mmap::map(&file)? }; - - let header = goblin::pe::header::Header::parse(&map)?; - let image_size = header - .optional_header - .map(|h| h.windows_fields.size_of_image) - .ok_or_else(|| anyhow::anyhow!("Missing optional header in PE image"))?; - - let machine = match header.coff_header.machine { - IMAGE_FILE_MACHINE_AMD64 => Machine::X64, - IMAGE_FILE_MACHINE_I386 => Machine::X86, - _ => Machine::Unknown, - }; - - Ok(ImageDetails { - image_size, - machine, - }) -} diff --git a/src/agent/debugger/src/lib.rs b/src/agent/debugger/src/lib.rs index 21992b96d..cd38a5ae5 100644 --- a/src/agent/debugger/src/lib.rs +++ b/src/agent/debugger/src/lib.rs @@ -7,3 +7,4 @@ pub mod dbghelp; pub mod debug_event; pub mod debugger; pub mod stack; +mod target; diff --git a/src/agent/debugger/src/target.rs b/src/agent/debugger/src/target.rs new file mode 100644 index 000000000..7fb1b14fa --- /dev/null +++ b/src/agent/debugger/src/target.rs @@ -0,0 +1,567 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use std::{ + collections::hash_map, + fs, + path::{Path, PathBuf}, +}; + +use anyhow::Result; +use log::{error, trace}; +use win_util::{file, process}; +use winapi::{ + shared::minwindef::{DWORD, LPVOID}, + um::{ + handleapi::CloseHandle, + winnt::{HANDLE, IMAGE_FILE_MACHINE_AMD64, IMAGE_FILE_MACHINE_I386}, + }, +}; + +use crate::{ + dbghelp::{self, FrameContext}, + debugger::{Breakpoint, BreakpointId, BreakpointType, ModuleBreakpoint, StepState}, +}; + +pub struct Module { + path: PathBuf, + file_handle: HANDLE, + base_address: u64, + image_size: u32, + machine: Machine, + + // Track if we need to call SymLoadModule for the dll. + sym_module_loaded: bool, +} + +impl Module { + fn new(module_handle: HANDLE, base_address: u64) -> Result { + let path = file::get_path_from_handle(module_handle).unwrap_or_else(|e| { + error!("Error getting path from file handle: {}", e); + "???".into() + }); + + let image_details = get_image_details(&path)?; + + Ok(Self { + path, + file_handle: module_handle, + base_address, + image_size: image_details.image_size, + machine: image_details.machine, + sym_module_loaded: false, + }) + } + + fn sym_load_module(&mut self, process_handle: HANDLE) -> Result<()> { + if !self.sym_module_loaded { + let dbghelp = dbghelp::lock()?; + + dbghelp.sym_load_module( + process_handle, + self.file_handle, + &self.path, + self.base_address, + self.image_size, + )?; + + self.sym_module_loaded = true; + } + + Ok(()) + } + + pub fn name(&self) -> &Path { + // Unwrap guaranteed by construction, we always have a filename. + self.path.file_stem().unwrap().as_ref() + } +} + +impl Drop for Module { + fn drop(&mut self) { + unsafe { CloseHandle(self.file_handle) }; + } +} + +pub struct Target { + process_id: DWORD, + process_handle: HANDLE, + current_thread_handle: HANDLE, + saw_initial_bp: bool, + saw_initial_wow64_bp: bool, + + // Track if we need to call SymInitialize for the process and if we need to notify + // dbghelp about loaded/unloaded dlls. + sym_initialized: bool, + exited: bool, + + thread_handles: fnv::FnvHashMap, + + // We cache the current thread context for possible repeated queries and modifications. + // We want to call GetThreadContext once, then call SetThreadContext (if necessary) before + // resuming. Calling Get/Set/Get/Set doesn't seem to work because the second Get doesn't + // see any the changes made in the Set call. + current_context: Option, + + // True if we need to set the thread context before resuming. + context_is_modified: bool, + + // Key is base address (which also happens to be the HANDLE). + modules: fnv::FnvHashMap, + + breakpoints: fnv::FnvHashMap, + + // Map of thread to stepping state (e.g. breakpoint address to restore breakpoints) + single_step: fnv::FnvHashMap, +} + +impl Target { + pub fn new( + process_id: DWORD, + thread_id: DWORD, + process_handle: HANDLE, + thread_handle: HANDLE, + ) -> Self { + let mut thread_handles = fnv::FnvHashMap::default(); + thread_handles.insert(thread_id, thread_handle); + + Self { + process_id, + current_thread_handle: thread_handle, + process_handle, + saw_initial_bp: false, + saw_initial_wow64_bp: false, + sym_initialized: false, + exited: false, + thread_handles, + current_context: None, + context_is_modified: false, + modules: fnv::FnvHashMap::default(), + breakpoints: fnv::FnvHashMap::default(), + single_step: fnv::FnvHashMap::default(), + } + } + + pub fn current_thread_handle(&self) -> HANDLE { + self.current_thread_handle + } + + pub fn create_new_thread(&mut self, thread_handle: HANDLE, thread_id: DWORD) { + self.current_thread_handle = thread_handle; + self.thread_handles.insert(thread_id, thread_handle); + } + + pub fn set_current_thread(&mut self, thread_id: DWORD) { + self.current_thread_handle = *self.thread_handles.get(&thread_id).unwrap(); + } + + pub fn exit_thread(&mut self, thread_id: DWORD) { + self.thread_handles.remove(&thread_id); + } + + pub fn process_handle(&self) -> HANDLE { + self.process_handle + } + + pub fn process_id(&self) -> DWORD { + self.process_id + } + + pub fn saw_initial_wow64_bp(&self) -> bool { + self.saw_initial_wow64_bp + } + + pub fn set_saw_initial_wow64_bp(&mut self) { + self.saw_initial_wow64_bp = true; + } + + pub fn saw_initial_bp(&self) -> bool { + self.saw_initial_bp + } + + pub fn set_saw_initial_bp(&mut self) { + self.saw_initial_bp = true; + } + + pub fn exited(&self) -> bool { + self.exited + } + + pub fn modules(&self) -> hash_map::Iter { + self.modules.iter() + } + + pub fn initial_bp(&mut self, load_symbols: bool) -> Result<()> { + self.saw_initial_bp = true; + + if load_symbols || !self.breakpoints.is_empty() { + self.sym_initialize()?; + + for (_, module) in self.modules.iter_mut() { + if let Err(e) = module.sym_load_module(self.process_handle) { + error!("Error loading symbols: {}", e); + } + } + } + + Ok(()) + } + + pub fn breakpoint_set_at_addr(&self, address: u64) -> bool { + self.breakpoints.contains_key(&address) + } + + pub(crate) fn single_step(&self, thread_handle: HANDLE) -> Option { + self.single_step.get(&thread_handle).cloned() + } + + pub fn sym_initialize(&mut self) -> Result<()> { + if !self.sym_initialized { + let dbghelp = dbghelp::lock()?; + if let Err(e) = dbghelp.sym_initialize(self.process_handle) { + error!("Error in SymInitializeW: {}", e); + + if let Err(e) = dbghelp.sym_cleanup(self.process_handle) { + error!("Error in SymCleanup: {}", e); + } + + return Err(e); + } + + for (_, module) in self.modules.iter_mut() { + if let Err(e) = module.sym_load_module(self.process_handle) { + error!( + "Error loading symbols for module {}: {}", + module.path.display(), + e + ); + } + } + + self.sym_initialized = true; + } + + Ok(()) + } + + /// Register the module loaded at `base_address`, returning the module name. + pub fn load_module( + &mut self, + module_handle: HANDLE, + base_address: u64, + ) -> Result> { + let mut module = Module::new(module_handle, base_address)?; + + trace!( + "Loading module {} at {:x}", + module.name().display(), + base_address + ); + + if module.machine == Machine::X64 && process::is_wow64_process(self.process_handle) { + // We ignore native dlls in wow64 processes. + return Ok(None); + } + + let module_name = module.name().to_owned(); + if self.sym_initialized { + if let Err(e) = module.sym_load_module(self.process_handle) { + error!("Error loading symbols: {}", e); + } + } + + let base_address = module.base_address; + if let Some(old_value) = self.modules.insert(base_address, module) { + error!( + "Existing module {} replace at base_address {}", + old_value.path.display(), + base_address + ); + } + + Ok(Some(module_name)) + } + + pub fn unload_module(&mut self, base_address: u64) { + // Drop the module and remove any breakpoints. + if let Some(module) = self.modules.remove(&base_address) { + let image_size = module.image_size as u64; + self.breakpoints + .retain(|&ip, _| ip < base_address || ip >= base_address + image_size); + } + } + + pub fn apply_absolute_breakpoint( + &mut self, + address: u64, + kind: BreakpointType, + id: BreakpointId, + ) -> Result<()> { + let original_byte: u8 = process::read_memory(self.process_handle, address as LPVOID)?; + + self.breakpoints + .entry(address) + .and_modify(|bp| { + bp.set_kind(kind); + bp.set_enabled(true); + bp.set_original_byte(Some(original_byte)); + bp.set_id(id); + }) + .or_insert(Breakpoint::new( + address, + kind, + /*enabled*/ true, + /*original_byte*/ Some(original_byte), + /*hit_count*/ 0, + id, + )); + + write_instruction_byte(self.process_handle, address, 0xcc)?; + + Ok(()) + } + + pub(crate) fn apply_module_breakpoints( + &mut self, + base_address: u64, + breakpoints: &[ModuleBreakpoint], + ) -> Result<()> { + if breakpoints.is_empty() { + return Ok(()); + } + + // We want to set every breakpoint for the module at once. We'll read the just the + // memory we need to do that, so find the min and max rva to compute how much memory + // to read and update in the remote process. + let (min, max) = breakpoints + .iter() + .fold((u64::max_value(), u64::min_value()), |acc, bp| { + (acc.0.min(bp.rva()), acc.1.max(bp.rva())) + }); + + // Add 1 to include the final byte. + let region_size = (max - min) + .checked_add(1) + .ok_or_else(|| anyhow::anyhow!("overflow in region size trying to set breakpoints"))? + as usize; + let remote_address = base_address.checked_add(min).ok_or_else(|| { + anyhow::anyhow!("overflow in remote address calculation trying to set breakpoints") + })? as LPVOID; + + let mut buffer: Vec = Vec::with_capacity(region_size); + unsafe { + buffer.set_len(region_size); + } + process::read_memory_array(self.process_handle, remote_address, &mut buffer[..])?; + + for mbp in breakpoints { + let ip = base_address + mbp.rva(); + let offset = (mbp.rva() - min) as usize; + + trace!("Setting breakpoint at {:x}", ip); + + let bp = Breakpoint::new( + ip, + mbp.kind(), + /*enabled*/ true, + Some(buffer[offset]), + /*hit_count*/ 0, + mbp.id(), + ); + + buffer[offset] = 0xcc; + + self.breakpoints.insert(ip, bp); + } + + process::write_memory_slice(self.process_handle, remote_address, &buffer[..])?; + process::flush_instruction_cache(self.process_handle, remote_address, region_size)?; + + Ok(()) + } + + pub fn prepare_to_resume(&mut self) -> Result<()> { + if let Some(context) = self.current_context.take() { + if self.context_is_modified { + context.set_thread_context(self.current_thread_handle)?; + } + } + + self.context_is_modified = false; + + Ok(()) + } + + fn ensure_current_context(&mut self) -> Result<()> { + if self.current_context.is_none() { + self.current_context = Some(dbghelp::get_thread_frame( + self.process_handle, + self.current_thread_handle, + )?); + } + + Ok(()) + } + + fn get_current_context(&mut self) -> Result<&FrameContext> { + self.ensure_current_context()?; + Ok(self.current_context.as_ref().unwrap()) + } + + fn get_current_context_mut(&mut self) -> Result<&mut FrameContext> { + self.ensure_current_context()?; + + // Assume the caller will modify the context. When it is modified, + // we must set it before resuming the target. + self.context_is_modified = true; + Ok(self.current_context.as_mut().unwrap()) + } + + pub fn read_register_u64(&mut self, reg: iced_x86::Register) -> Result { + let current_context = self.get_current_context()?; + Ok(current_context.get_register_u64(reg)) + } + + pub fn read_flags_register(&mut self) -> Result { + let current_context = self.get_current_context()?; + Ok(current_context.get_flags()) + } + + /// Handle a breakpoint that we set (as opposed to a breakpoint in user code, e.g. + /// assertion.) + /// + /// Return the breakpoint id if it should be reported to the client. + pub fn handle_breakpoint(&mut self, pc: u64) -> Result> { + enum HandleBreakpoint { + User(BreakpointId, bool), + StepOut(u64), + } + + let handle_breakpoint = { + let bp = self.breakpoints.get_mut(&pc).unwrap(); + + bp.increment_hit_count(); + + write_instruction_byte(self.process_handle, bp.ip(), bp.original_byte().unwrap())?; + + match bp.kind() { + BreakpointType::OneTime => { + bp.set_enabled(false); + bp.set_original_byte(None); + + // We are clearing the breakpoint after hitting it, so we do not need + // to single step. + HandleBreakpoint::User(bp.id(), false) + } + + BreakpointType::Counter => { + // Single step so we can restore the breakpoint after stepping. + HandleBreakpoint::User(bp.id(), true) + } + + BreakpointType::StepOut { rsp } => HandleBreakpoint::StepOut(rsp), + } + }; + + let context = self.get_current_context_mut()?; + context.set_program_counter(pc); + + // We need to single step if we need to restore the breakpoint. + let single_step = match handle_breakpoint { + HandleBreakpoint::User(_, single_step) => single_step, + + // Single step only when in a recursive call, which is inferred when the current + // stack pointer (from context) is less then at the target to step out from (rsp). + // Note this only works if the stack grows down. + HandleBreakpoint::StepOut(rsp) => rsp > context.stack_pointer(), + }; + + if single_step { + context.set_single_step(true); + self.single_step + .insert(self.current_thread_handle, StepState::Breakpoint { pc }); + } + + Ok(match handle_breakpoint { + HandleBreakpoint::User(id, _) => Some(id), + HandleBreakpoint::StepOut(_) => None, + }) + } + + pub(crate) fn handle_single_step(&mut self, step_state: StepState) -> Result<()> { + match step_state { + StepState::Breakpoint { pc } => { + write_instruction_byte(self.process_handle, pc, 0xcc)?; + } + _ => {} + } + + self.single_step.remove(&self.current_thread_handle); + Ok(()) + } + + pub fn prepare_to_step(&mut self) -> Result { + // Don't change the reason we're single stepping on this thread if + // we previously set the reason (e.g. so we would restore a breakpoint). + self.single_step + .entry(self.current_thread_handle) + .or_insert(StepState::SingleStep); + + let context = self.get_current_context_mut()?; + context.set_single_step(true); + + Ok(true) + } + + pub fn set_exited(&mut self) -> Result<()> { + self.exited = true; + + if self.sym_initialized { + let dbghelp = dbghelp::lock()?; + dbghelp.sym_cleanup(self.process_handle)?; + } + Ok(()) + } +} + +#[derive(Copy, Clone, Debug, PartialEq)] +enum Machine { + Unknown, + X64, + X86, +} + +struct ImageDetails { + image_size: u32, + machine: Machine, +} + +fn get_image_details(path: &Path) -> Result { + let file = fs::File::open(path)?; + let map = unsafe { memmap::Mmap::map(&file)? }; + + let header = goblin::pe::header::Header::parse(&map)?; + let image_size = header + .optional_header + .map(|h| h.windows_fields.size_of_image) + .ok_or_else(|| anyhow::anyhow!("Missing optional header in PE image"))?; + + let machine = match header.coff_header.machine { + IMAGE_FILE_MACHINE_AMD64 => Machine::X64, + IMAGE_FILE_MACHINE_I386 => Machine::X86, + _ => Machine::Unknown, + }; + + Ok(ImageDetails { + image_size, + machine, + }) +} + +fn write_instruction_byte(process_handle: HANDLE, ip: u64, b: u8) -> Result<()> { + let orig_byte = [b; 1]; + let remote_address = ip as LPVOID; + process::write_memory_slice(process_handle, remote_address, &orig_byte)?; + process::flush_instruction_cache(process_handle, remote_address, orig_byte.len())?; + Ok(()) +}