use core::cmp::Ordering;
use core::fmt;
use core::iter::FusedIterator;
use scroll::ctx::TryFromCtx;
use scroll::{self, Pread, Pwrite};
use crate::error;
use crate::pe::data_directories;
use crate::pe::options;
use crate::pe::section_table;
use crate::pe::utils;
const UNW_FLAG_EHANDLER: u8 = 0x01;
const UNW_FLAG_UHANDLER: u8 = 0x02;
const UNW_FLAG_CHAININFO: u8 = 0x04;
const UWOP_PUSH_NONVOL: u8 = 0;
const UWOP_ALLOC_LARGE: u8 = 1;
const UWOP_ALLOC_SMALL: u8 = 2;
const UWOP_SET_FPREG: u8 = 3;
const UWOP_SAVE_NONVOL: u8 = 4;
const UWOP_SAVE_NONVOL_FAR: u8 = 5;
const UWOP_EPILOG: u8 = 6;
const UWOP_SPARE_CODE: u8 = 7;
const UWOP_SAVE_XMM128: u8 = 8;
const UWOP_SAVE_XMM128_FAR: u8 = 9;
const UWOP_PUSH_MACHFRAME: u8 = 10;
const RUNTIME_FUNCTION_SIZE: usize = 12;
const UNWIND_CODE_SIZE: usize = 2;
#[repr(C)]
#[derive(Copy, Clone, PartialEq, Default, Pread, Pwrite)]
pub struct RuntimeFunction {
pub begin_address: u32,
pub end_address: u32,
pub unwind_info_address: u32,
}
impl fmt::Debug for RuntimeFunction {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("RuntimeFunction")
.field("begin_address", &format_args!("{:#x}", self.begin_address))
.field("end_address", &format_args!("{:#x}", self.end_address))
.field(
"unwind_info_address",
&format_args!("{:#x}", self.unwind_info_address),
)
.finish()
}
}
#[derive(Debug)]
pub struct RuntimeFunctionIterator<'a> {
data: &'a [u8],
}
impl Iterator for RuntimeFunctionIterator<'_> {
type Item = error::Result<RuntimeFunction>;
fn next(&mut self) -> Option<Self::Item> {
if self.data.is_empty() {
return None;
}
Some(match self.data.pread_with(0, scroll::LE) {
Ok(func) => {
self.data = &self.data[RUNTIME_FUNCTION_SIZE..];
Ok(func)
}
Err(error) => {
self.data = &[];
Err(error.into())
}
})
}
fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.data.len() / RUNTIME_FUNCTION_SIZE;
(len, Some(len))
}
}
impl FusedIterator for RuntimeFunctionIterator<'_> {}
impl ExactSizeIterator for RuntimeFunctionIterator<'_> {}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)]
pub struct Register(pub u8);
impl Register {
fn xmm(number: u8) -> Self {
Register(number + 17)
}
pub fn name(self) -> &'static str {
match self.0 {
0 => "$rax",
1 => "$rcx",
2 => "$rdx",
3 => "$rbx",
4 => "$rsp",
5 => "$rbp",
6 => "$rsi",
7 => "$rdi",
8 => "$r8",
9 => "$r9",
10 => "$r10",
11 => "$r11",
12 => "$r12",
13 => "$r13",
14 => "$r14",
15 => "$r15",
16 => "$rip",
17 => "$xmm0",
18 => "$xmm1",
19 => "$xmm2",
20 => "$xmm3",
21 => "$xmm4",
22 => "$xmm5",
23 => "$xmm6",
24 => "$xmm7",
25 => "$xmm8",
26 => "$xmm9",
27 => "$xmm10",
28 => "$xmm11",
29 => "$xmm12",
30 => "$xmm13",
31 => "$xmm14",
32 => "$xmm15",
_ => "",
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum StackFrameOffset {
RSP(u32),
FP(u32),
}
impl StackFrameOffset {
fn with_ctx(offset: u32, ctx: UnwindOpContext) -> Self {
match ctx.frame_register {
Register(0) => StackFrameOffset::RSP(offset),
Register(_) => StackFrameOffset::FP(offset),
}
}
}
impl fmt::Display for Register {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.write_str(self.name())
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum UnwindOperation {
PushNonVolatile(Register),
Alloc(u32),
SetFPRegister,
SaveNonVolatile(Register, StackFrameOffset),
SaveXMM(Register, StackFrameOffset),
Epilog,
SaveXMM128(Register, StackFrameOffset),
PushMachineFrame(bool),
Noop,
}
#[derive(Clone, Copy, Debug, PartialEq)]
struct UnwindOpContext {
version: u8,
frame_register: Register,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct UnwindCode {
pub code_offset: u8,
pub operation: UnwindOperation,
}
impl<'a> TryFromCtx<'a, UnwindOpContext> for UnwindCode {
type Error = error::Error;
#[inline]
fn try_from_ctx(bytes: &'a [u8], ctx: UnwindOpContext) -> Result<(Self, usize), Self::Error> {
let mut read = 0;
let code_offset = bytes.gread_with::<u8>(&mut read, scroll::LE)?;
let operation = bytes.gread_with::<u8>(&mut read, scroll::LE)?;
let operation_code = operation & 0xf;
let operation_info = operation >> 4;
let operation = match operation_code {
self::UWOP_PUSH_NONVOL => {
let register = Register(operation_info);
UnwindOperation::PushNonVolatile(register)
}
self::UWOP_ALLOC_LARGE => {
let offset = match operation_info {
0 => u32::from(bytes.gread_with::<u16>(&mut read, scroll::LE)?) * 8,
1 => bytes.gread_with::<u32>(&mut read, scroll::LE)?,
i => {
let msg = format!("invalid op info ({}) for UWOP_ALLOC_LARGE", i);
return Err(error::Error::Malformed(msg));
}
};
UnwindOperation::Alloc(offset)
}
self::UWOP_ALLOC_SMALL => {
let offset = u32::from(operation_info) * 8 + 8;
UnwindOperation::Alloc(offset)
}
self::UWOP_SET_FPREG => UnwindOperation::SetFPRegister,
self::UWOP_SAVE_NONVOL => {
let register = Register(operation_info);
let offset = u32::from(bytes.gread_with::<u16>(&mut read, scroll::LE)?) * 8;
UnwindOperation::SaveNonVolatile(register, StackFrameOffset::with_ctx(offset, ctx))
}
self::UWOP_SAVE_NONVOL_FAR => {
let register = Register(operation_info);
let offset = bytes.gread_with::<u32>(&mut read, scroll::LE)?;
UnwindOperation::SaveNonVolatile(register, StackFrameOffset::with_ctx(offset, ctx))
}
self::UWOP_EPILOG => {
let data = u32::from(bytes.gread_with::<u16>(&mut read, scroll::LE)?) * 16;
if ctx.version == 1 {
let register = Register::xmm(operation_info);
UnwindOperation::SaveXMM(register, StackFrameOffset::with_ctx(data, ctx))
} else {
UnwindOperation::Epilog
}
}
self::UWOP_SPARE_CODE => {
let data = bytes.gread_with::<u32>(&mut read, scroll::LE)?;
if ctx.version == 1 {
let register = Register::xmm(operation_info);
UnwindOperation::SaveXMM128(register, StackFrameOffset::with_ctx(data, ctx))
} else {
UnwindOperation::Noop
}
}
self::UWOP_SAVE_XMM128 => {
let register = Register::xmm(operation_info);
let offset = u32::from(bytes.gread_with::<u16>(&mut read, scroll::LE)?) * 16;
UnwindOperation::SaveXMM128(register, StackFrameOffset::with_ctx(offset, ctx))
}
self::UWOP_SAVE_XMM128_FAR => {
let register = Register::xmm(operation_info);
let offset = bytes.gread_with::<u32>(&mut read, scroll::LE)?;
UnwindOperation::SaveXMM128(register, StackFrameOffset::with_ctx(offset, ctx))
}
self::UWOP_PUSH_MACHFRAME => {
let is_error = match operation_info {
0 => false,
1 => true,
i => {
let msg = format!("invalid op info ({}) for UWOP_PUSH_MACHFRAME", i);
return Err(error::Error::Malformed(msg));
}
};
UnwindOperation::PushMachineFrame(is_error)
}
op => {
let msg = format!("unknown unwind op code ({})", op);
return Err(error::Error::Malformed(msg));
}
};
let code = UnwindCode {
code_offset,
operation,
};
Ok((code, read))
}
}
#[derive(Clone, Debug)]
pub struct UnwindCodeIterator<'a> {
bytes: &'a [u8],
offset: usize,
context: UnwindOpContext,
}
impl Iterator for UnwindCodeIterator<'_> {
type Item = error::Result<UnwindCode>;
fn next(&mut self) -> Option<Self::Item> {
if self.offset >= self.bytes.len() {
return None;
}
Some(self.bytes.gread_with(&mut self.offset, self.context))
}
fn size_hint(&self) -> (usize, Option<usize>) {
let upper = (self.bytes.len() - self.offset) / UNWIND_CODE_SIZE;
let lower = (upper + 3 - (upper % 3)) / 3;
(lower, Some(upper))
}
}
impl FusedIterator for UnwindCodeIterator<'_> {}
#[derive(Copy, Clone, Debug, PartialEq)]
pub enum UnwindHandler<'a> {
ExceptionHandler(u32, &'a [u8]),
TerminationHandler(u32, &'a [u8]),
}
#[derive(Clone)]
pub struct UnwindInfo<'a> {
pub version: u8,
pub size_of_prolog: u8,
pub frame_register: Register,
pub frame_register_offset: u32,
pub chained_info: Option<RuntimeFunction>,
pub handler: Option<UnwindHandler<'a>>,
code_bytes: &'a [u8],
}
impl<'a> UnwindInfo<'a> {
pub fn parse(bytes: &'a [u8], mut offset: usize) -> error::Result<Self> {
let version_flags: u8 = bytes.gread_with(&mut offset, scroll::LE)?;
let version = version_flags & 0b111;
let flags = version_flags >> 3;
if version < 1 || version > 2 {
let msg = format!("unsupported unwind code version ({})", version);
return Err(error::Error::Malformed(msg));
}
let size_of_prolog = bytes.gread_with::<u8>(&mut offset, scroll::LE)?;
let count_of_codes = bytes.gread_with::<u8>(&mut offset, scroll::LE)?;
let frame_info = bytes.gread_with::<u8>(&mut offset, scroll::LE)?;
let frame_register = Register(frame_info & 0xf);
let frame_register_offset = u32::from((frame_info >> 4) * 16);
let codes_size = count_of_codes as usize * UNWIND_CODE_SIZE;
let code_bytes = bytes.gread_with(&mut offset, codes_size)?;
if count_of_codes % 2 != 0 {
offset += 2;
}
debug_assert!(offset % 4 == 0);
let mut chained_info = None;
let mut handler = None;
if flags & UNW_FLAG_CHAININFO != 0 {
chained_info = Some(bytes.gread_with(&mut offset, scroll::LE)?);
} else if flags & (UNW_FLAG_EHANDLER | UNW_FLAG_UHANDLER) != 0 {
let address = bytes.gread_with::<u32>(&mut offset, scroll::LE)?;
let data = &bytes[offset..];
handler = Some(if flags & UNW_FLAG_EHANDLER != 0 {
UnwindHandler::ExceptionHandler(address, data)
} else {
UnwindHandler::TerminationHandler(address, data)
});
}
Ok(UnwindInfo {
version,
size_of_prolog,
frame_register,
frame_register_offset,
chained_info,
handler,
code_bytes,
})
}
pub fn unwind_codes(&self) -> UnwindCodeIterator<'a> {
UnwindCodeIterator {
bytes: self.code_bytes,
offset: 0,
context: UnwindOpContext {
version: self.version,
frame_register: self.frame_register,
},
}
}
}
impl fmt::Debug for UnwindInfo<'_> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let count_of_codes = self.code_bytes.len() / UNWIND_CODE_SIZE;
f.debug_struct("UnwindInfo")
.field("version", &self.version)
.field("size_of_prolog", &self.size_of_prolog)
.field("frame_register", &self.frame_register)
.field("frame_register_offset", &self.frame_register_offset)
.field("count_of_codes", &count_of_codes)
.field("chained_info", &self.chained_info)
.field("handler", &self.handler)
.finish()
}
}
impl<'a> IntoIterator for &'_ UnwindInfo<'a> {
type Item = error::Result<UnwindCode>;
type IntoIter = UnwindCodeIterator<'a>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
self.unwind_codes()
}
}
pub struct ExceptionData<'a> {
bytes: &'a [u8],
offset: usize,
size: usize,
file_alignment: u32,
}
impl<'a> ExceptionData<'a> {
pub fn parse(
bytes: &'a [u8],
directory: data_directories::DataDirectory,
sections: &[section_table::SectionTable],
file_alignment: u32,
) -> error::Result<Self> {
Self::parse_with_opts(
bytes,
directory,
sections,
file_alignment,
&options::ParseOptions::default(),
)
}
pub fn parse_with_opts(
bytes: &'a [u8],
directory: data_directories::DataDirectory,
sections: &[section_table::SectionTable],
file_alignment: u32,
opts: &options::ParseOptions,
) -> error::Result<Self> {
let size = directory.size as usize;
if size % RUNTIME_FUNCTION_SIZE != 0 {
return Err(error::Error::from(scroll::Error::BadInput {
size,
msg: "invalid exception directory table size",
}));
}
let rva = directory.virtual_address as usize;
let offset = utils::find_offset(rva, sections, file_alignment, opts).ok_or_else(|| {
error::Error::Malformed(format!("cannot map exception_rva ({:#x}) into offset", rva))
})?;
if offset % 4 != 0 {
return Err(error::Error::from(scroll::Error::BadOffset(offset)));
}
Ok(ExceptionData {
bytes,
offset,
size,
file_alignment,
})
}
pub fn len(&self) -> usize {
self.size / RUNTIME_FUNCTION_SIZE
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn functions(&self) -> RuntimeFunctionIterator<'a> {
RuntimeFunctionIterator {
data: &self.bytes[self.offset..self.offset + self.size],
}
}
pub fn get_function(&self, index: usize) -> error::Result<RuntimeFunction> {
self.get_function_by_offset(self.offset + index * RUNTIME_FUNCTION_SIZE)
}
pub fn find_function(&self, rva: u32) -> error::Result<Option<RuntimeFunction>> {
let mut size = self.len();
if size == 0 {
return Ok(None);
}
let mut base = 0;
while size > 1 {
let half = size / 2;
let mid = base + half;
let offset = self.offset + mid * RUNTIME_FUNCTION_SIZE;
let addr = self.bytes.pread_with::<u32>(offset, scroll::LE)?;
base = if addr > rva { base } else { mid };
size -= half;
}
let offset = self.offset + base * RUNTIME_FUNCTION_SIZE;
let addr = self.bytes.pread_with::<u32>(offset, scroll::LE)?;
let function = match addr.cmp(&rva) {
Ordering::Less | Ordering::Equal => self.get_function(base)?,
Ordering::Greater if base == 0 => return Ok(None),
Ordering::Greater => self.get_function(base - 1)?,
};
if function.end_address > rva {
Ok(Some(function))
} else {
Ok(None)
}
}
pub fn get_unwind_info(
&self,
function: RuntimeFunction,
sections: &[section_table::SectionTable],
) -> error::Result<UnwindInfo<'a>> {
self.get_unwind_info_with_opts(function, sections, &options::ParseOptions::default())
}
pub fn get_unwind_info_with_opts(
&self,
mut function: RuntimeFunction,
sections: &[section_table::SectionTable],
opts: &options::ParseOptions,
) -> error::Result<UnwindInfo<'a>> {
while function.unwind_info_address % 2 != 0 {
let rva = (function.unwind_info_address & !1) as usize;
function = self.get_function_by_rva_with_opts(rva, sections, opts)?;
}
let rva = function.unwind_info_address as usize;
let offset =
utils::find_offset(rva, sections, self.file_alignment, opts).ok_or_else(|| {
error::Error::Malformed(format!("cannot map unwind rva ({:#x}) into offset", rva))
})?;
UnwindInfo::parse(self.bytes, offset)
}
#[allow(dead_code)]
fn get_function_by_rva(
&self,
rva: usize,
sections: &[section_table::SectionTable],
) -> error::Result<RuntimeFunction> {
self.get_function_by_rva_with_opts(rva, sections, &options::ParseOptions::default())
}
fn get_function_by_rva_with_opts(
&self,
rva: usize,
sections: &[section_table::SectionTable],
opts: &options::ParseOptions,
) -> error::Result<RuntimeFunction> {
let offset =
utils::find_offset(rva, sections, self.file_alignment, opts).ok_or_else(|| {
error::Error::Malformed(format!(
"cannot map exception rva ({:#x}) into offset",
rva
))
})?;
self.get_function_by_offset(offset)
}
#[inline]
fn get_function_by_offset(&self, offset: usize) -> error::Result<RuntimeFunction> {
debug_assert!((offset - self.offset) % RUNTIME_FUNCTION_SIZE == 0);
debug_assert!(offset < self.offset + self.size);
Ok(self.bytes.pread_with(offset, scroll::LE)?)
}
}
impl fmt::Debug for ExceptionData<'_> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("ExceptionData")
.field("file_alignment", &self.file_alignment)
.field("offset", &format_args!("{:#x}", self.offset))
.field("size", &format_args!("{:#x}", self.size))
.field("len", &self.len())
.finish()
}
}
impl<'a> IntoIterator for &'_ ExceptionData<'a> {
type Item = error::Result<RuntimeFunction>;
type IntoIter = RuntimeFunctionIterator<'a>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
self.functions()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_size_of_runtime_function() {
assert_eq!(
std::mem::size_of::<RuntimeFunction>(),
RUNTIME_FUNCTION_SIZE
);
}
#[test]
fn test_iter_unwind_codes() {
let unwind_info = UnwindInfo {
version: 1,
size_of_prolog: 4,
frame_register: Register(0),
frame_register_offset: 0,
chained_info: None,
handler: None,
code_bytes: &[4, 98],
};
let unwind_codes: Vec<UnwindCode> = unwind_info
.unwind_codes()
.map(|result| result.expect("parse unwind code"))
.collect();
assert_eq!(unwind_codes.len(), 1);
let expected = UnwindCode {
code_offset: 4,
operation: UnwindOperation::Alloc(56),
};
assert_eq!(unwind_codes[0], expected);
}
}