use std::cmp::{self, Ordering};
use std::fmt;
use std::iter::FusedIterator;
use std::mem;
use std::ops::Range;
use crate::common::*;
use crate::msf::Stream;
use crate::pe::ImageSectionHeader;
#[repr(C)]
#[derive(Clone, Copy, Eq, PartialEq)]
pub(crate) struct OMAPRecord {
source_address: u32,
target_address: u32,
}
impl OMAPRecord {
pub fn new(source_address: u32, target_address: u32) -> Self {
Self {
source_address: source_address.to_le(),
target_address: target_address.to_le(),
}
}
#[inline]
pub fn source_address(self) -> u32 {
u32::from_le(self.source_address)
}
#[inline]
pub fn target_address(self) -> u32 {
u32::from_le(self.target_address)
}
#[inline]
fn translate(self, address: u32) -> u32 {
debug_assert!(self.source_address() <= address);
(address - self.source_address()) + self.target_address()
}
}
impl fmt::Debug for OMAPRecord {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("OMAPRecord")
.field(
"source_address",
&format_args!("{:#010x}", self.source_address()),
)
.field(
"target_address",
&format_args!("{:#010x}", self.target_address()),
)
.finish()
}
}
impl PartialOrd for OMAPRecord {
#[inline]
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
self.source_address().partial_cmp(&other.source_address())
}
}
impl Ord for OMAPRecord {
#[inline]
fn cmp(&self, other: &Self) -> Ordering {
self.source_address().cmp(&other.source_address())
}
}
pub(crate) struct OMAPTable<'s> {
stream: Stream<'s>,
}
impl<'s> OMAPTable<'s> {
pub(crate) fn parse(stream: Stream<'s>) -> Result<Self> {
match cast_aligned::<OMAPRecord>(stream.as_slice()) {
Some(_) => Ok(OMAPTable { stream }),
None => Err(Error::InvalidStreamLength("OMAP")),
}
}
#[inline]
pub fn records(&self) -> &[OMAPRecord] {
cast_aligned(self.stream.as_slice()).unwrap()
}
pub fn lookup(&self, source_address: u32) -> Option<u32> {
let records = self.records();
let index = match records.binary_search_by_key(&source_address, |r| r.source_address()) {
Ok(i) => i,
Err(0) => return None,
Err(i) => i - 1,
};
let record = records[index];
if record.target_address() == 0 {
return None;
}
Some(record.translate(source_address))
}
pub fn lookup_range(&self, range: Range<u32>) -> RangeIter<'_> {
let Range { start, end } = range;
if end <= start {
return RangeIter::empty();
}
let records = self.records();
let (record, next) = match records.binary_search_by_key(&start, |r| r.source_address()) {
Ok(i) => (records[i], &records[i + 1..]),
Err(0) => (OMAPRecord::new(0, 0), records),
Err(i) => (records[i - 1], &records[i..]),
};
RangeIter {
records: next.iter(),
record,
addr: start,
end,
}
}
}
impl fmt::Debug for OMAPTable<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_tuple("OMAPTable").field(&self.records()).finish()
}
}
pub(crate) struct RangeIter<'t> {
records: std::slice::Iter<'t, OMAPRecord>,
record: OMAPRecord,
addr: u32,
end: u32,
}
impl<'t> RangeIter<'t> {
pub fn empty() -> Self {
RangeIter {
records: [].iter(),
record: OMAPRecord::new(0, 0),
addr: 0,
end: 0,
}
}
pub fn identity(range: Range<u32>) -> Self {
RangeIter {
records: [].iter(),
record: OMAPRecord::new(range.start, range.start),
addr: range.start,
end: range.end,
}
}
}
impl Default for RangeIter<'_> {
fn default() -> Self {
Self::empty()
}
}
impl Iterator for RangeIter<'_> {
type Item = Range<u32>;
fn next(&mut self) -> Option<Self::Item> {
while self.addr < self.end {
let next_record = match self.records.next() {
Some(record) => *record,
None => OMAPRecord::new(self.end, 0),
};
let subrange_end = cmp::min(next_record.source_address(), self.end);
let subrange_start = mem::replace(&mut self.addr, subrange_end);
let last_record = mem::replace(&mut self.record, next_record);
if subrange_start >= subrange_end || last_record.target_address() == 0 {
continue;
}
let translated_start = last_record.translate(subrange_start);
let translated_end = last_record.translate(subrange_end);
return Some(translated_start..translated_end);
}
None
}
}
impl FusedIterator for RangeIter<'_> {}
pub struct RvaRangeIter<'t>(RangeIter<'t>);
impl Iterator for RvaRangeIter<'_> {
type Item = Range<Rva>;
fn next(&mut self) -> Option<Self::Item> {
self.0.next().map(|range| Rva(range.start)..Rva(range.end))
}
}
impl FusedIterator for RvaRangeIter<'_> {}
pub struct PdbInternalRvaRangeIter<'t>(RangeIter<'t>);
impl Iterator for PdbInternalRvaRangeIter<'_> {
type Item = Range<PdbInternalRva>;
fn next(&mut self) -> Option<Self::Item> {
self.0
.next()
.map(|range| PdbInternalRva(range.start)..PdbInternalRva(range.end))
}
}
impl FusedIterator for PdbInternalRvaRangeIter<'_> {}
#[derive(Debug, Default)]
pub struct AddressMap<'s> {
pub(crate) original_sections: Vec<ImageSectionHeader>,
pub(crate) transformed_sections: Option<Vec<ImageSectionHeader>>,
pub(crate) transformed_to_original: Option<OMAPTable<'s>>,
pub(crate) original_to_transformed: Option<OMAPTable<'s>>,
}
impl<'s> AddressMap<'s> {
pub fn rva_ranges(&self, range: Range<PdbInternalRva>) -> RvaRangeIter<'_> {
RvaRangeIter(match self.original_to_transformed {
Some(ref omap) => omap.lookup_range(range.start.0..range.end.0),
None => RangeIter::identity(range.start.0..range.end.0),
})
}
pub fn internal_rva_ranges(&self, range: Range<Rva>) -> PdbInternalRvaRangeIter<'_> {
PdbInternalRvaRangeIter(match self.transformed_to_original {
Some(ref omap) => omap.lookup_range(range.start.0..range.end.0),
None => RangeIter::identity(range.start.0..range.end.0),
})
}
}
fn get_section_offset(sections: &[ImageSectionHeader], address: u32) -> Option<(u16, u32)> {
let (index, section) = sections
.iter()
.take_while(|s| s.virtual_address <= address)
.enumerate()
.find(|(_, s)| address < s.virtual_address + s.size_of_raw_data)?;
Some((index as u16 + 1, address - section.virtual_address))
}
fn get_virtual_address(sections: &[ImageSectionHeader], section: u16, offset: u32) -> Option<u32> {
(section as usize)
.checked_sub(1)
.and_then(|i| sections.get(i))
.map(|section| section.virtual_address + offset)
}
impl Rva {
pub fn to_internal_rva(self, translator: &AddressMap<'_>) -> Option<PdbInternalRva> {
match translator.transformed_to_original {
Some(ref omap) => omap.lookup(self.0).map(PdbInternalRva),
None => Some(PdbInternalRva(self.0)),
}
}
pub fn to_section_offset(self, translator: &AddressMap<'_>) -> Option<SectionOffset> {
let (section, offset) = match translator.transformed_sections {
Some(ref sections) => get_section_offset(sections, self.0)?,
None => get_section_offset(&translator.original_sections, self.0)?,
};
Some(SectionOffset { section, offset })
}
pub fn to_internal_offset(
self,
translator: &AddressMap<'_>,
) -> Option<PdbInternalSectionOffset> {
self.to_internal_rva(translator)?
.to_internal_offset(translator)
}
}
impl PdbInternalRva {
pub fn to_rva(self, translator: &AddressMap<'_>) -> Option<Rva> {
match translator.original_to_transformed {
Some(ref omap) => omap.lookup(self.0).map(Rva),
None => Some(Rva(self.0)),
}
}
pub fn to_section_offset(self, translator: &AddressMap<'_>) -> Option<SectionOffset> {
self.to_rva(translator)?.to_section_offset(translator)
}
pub fn to_internal_offset(
self,
translator: &AddressMap<'_>,
) -> Option<PdbInternalSectionOffset> {
let (section, offset) = get_section_offset(&translator.original_sections, self.0)?;
Some(PdbInternalSectionOffset { section, offset })
}
}
impl SectionOffset {
pub fn to_rva(self, translator: &AddressMap<'_>) -> Option<Rva> {
let address = match translator.transformed_sections {
Some(ref sections) => get_virtual_address(sections, self.section, self.offset)?,
None => get_virtual_address(&translator.original_sections, self.section, self.offset)?,
};
Some(Rva(address))
}
pub fn to_internal_rva(self, translator: &AddressMap<'_>) -> Option<PdbInternalRva> {
self.to_rva(translator)?.to_internal_rva(translator)
}
pub fn to_internal_offset(
self,
translator: &AddressMap<'_>,
) -> Option<PdbInternalSectionOffset> {
if translator.transformed_sections.is_none() {
let Self { section, offset } = self;
return Some(PdbInternalSectionOffset { section, offset });
}
self.to_internal_rva(translator)?
.to_internal_offset(translator)
}
}
impl PdbInternalSectionOffset {
pub fn to_rva(self, translator: &AddressMap<'_>) -> Option<Rva> {
self.to_internal_rva(translator)?.to_rva(translator)
}
pub fn to_internal_rva(self, translator: &AddressMap<'_>) -> Option<PdbInternalRva> {
get_virtual_address(&translator.original_sections, self.section, self.offset)
.map(PdbInternalRva)
}
pub fn to_section_offset(self, translator: &AddressMap<'_>) -> Option<SectionOffset> {
if translator.transformed_sections.is_none() {
let Self { section, offset } = self;
return Some(SectionOffset { section, offset });
}
self.to_rva(translator)?.to_section_offset(translator)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::mem;
#[test]
fn test_omap_record() {
assert_eq!(mem::size_of::<OMAPRecord>(), 8);
assert_eq!(mem::align_of::<OMAPRecord>(), 4);
}
#[test]
fn test_get_virtual_address() {
let sections = vec![ImageSectionHeader {
virtual_address: 0x1000_0000,
..Default::default()
}];
assert_eq!(get_virtual_address(§ions, 1, 0x1234), Some(0x1000_1234));
assert_eq!(get_virtual_address(§ions, 2, 0x1234), None);
assert_eq!(get_virtual_address(§ions, 0, 0x1234), None);
}
}