use crate::{
parsed::*,
traits::{Decode, View},
File, TableIndex,
};
use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
pub struct TypeReader {
pub files: Vec<File>,
pub types: BTreeMap<String, BTreeMap<String, TypeDef>>,
}
impl TypeReader {
pub fn from_os() -> Self {
let dir = std::env::var("windir").expect("No `windir` environent variable set");
let mut dir = PathBuf::from(dir);
dir.push(SYSTEM32);
dir.push("winmetadata");
let files = std::fs::read_dir(dir)
.expect("Could not read Windowws metadata path")
.filter_map(|value| value.ok())
.map(|value| value.path());
Self::from_iter(files)
}
pub fn from_win32<P: AsRef<Path>>(file: P) -> Self {
let mut reader = Self {
files: Vec::default(),
types: BTreeMap::default(),
};
let file = File::new(file);
reader.insert_file_at_index(file, 0, InsertMode::All);
reader
}
pub fn from_iter<I: IntoIterator<Item = PathBuf>>(files: I) -> Self {
let mut reader = Self {
files: Vec::default(),
types: BTreeMap::default(),
};
for (file_index, file) in files.into_iter().enumerate() {
let file = File::new(file);
reader.insert_file_at_index(file, file_index, InsertMode::WinrtOnly);
}
reader
}
fn insert_file_at_index(&mut self, file: File, file_index: usize, insert_mode: InsertMode) {
let row_count = file.type_def_table().row_count;
self.files.push(file);
for row in 0..row_count {
let def = TypeDef(Row::new(row, TableIndex::TypeDef, file_index as u16));
if insert_mode == InsertMode::WinrtOnly && !def.is_winrt(&self) {
continue;
}
let (namespace, name) = def.name(&self);
let namespace = namespace.to_string();
let name = name.to_string();
self.types
.entry(namespace)
.or_default()
.entry(name)
.or_insert(def);
}
}
pub fn namespaces(&self) -> impl Iterator<Item = &String> {
self.types.keys()
}
pub fn namespace_types(&self, namespace: &str) -> impl Iterator<Item = (&str, &TypeDef)> {
self.types[namespace].iter().map(|(n, t)| (n.as_str(), t))
}
pub fn resolve_type_def(&self, (namespace, type_name): (&str, &str)) -> TypeDef {
if let Some(types) = self.types.get(namespace) {
if let Some(def) = types.get(type_name) {
return *def;
}
}
panic!("Could not find type `{}.{}`", namespace, type_name);
}
pub fn u32(&self, row: Row, column: u32) -> u32 {
let file = &self.files[row.file_index as usize];
let table = &file.tables[row.table_index as usize];
let offset = table.data + row.index * table.row_size + table.columns[column as usize].0;
match table.columns[column as usize].1 {
1 => file.bytes.copy_as::<u8>(offset) as u32,
2 => file.bytes.copy_as::<u16>(offset) as u32,
4 => file.bytes.copy_as::<u32>(offset) as u32,
_ => file.bytes.copy_as::<u64>(offset) as u32,
}
}
pub fn str(&self, row: Row, column: u32) -> &str {
let file = &self.files[row.file_index as usize];
let offset = (file.strings + self.u32(row, column)) as usize;
let last = file.bytes[offset..]
.iter()
.position(|c| *c == b'\0')
.unwrap();
std::str::from_utf8(&file.bytes[offset..offset + last]).unwrap()
}
pub(crate) fn decode<T: Decode>(&self, row: Row, column: u32) -> T {
T::decode(self.u32(row, column), row.file_index)
}
pub(crate) fn list(
&self,
row: Row,
table: TableIndex,
column: u32,
) -> impl Iterator<Item = Row> {
let file = &self.files[row.file_index as usize];
let first = self.u32(row, column) - 1;
let last = if row.index + 1 < file.tables[row.table_index as usize].row_count {
self.u32(row.next(), column) - 1
} else {
file.tables[table as usize].row_count
};
(first..last).map(move |value| Row::new(value, table, row.file_index))
}
pub fn blob(&self, row: Row, column: u32) -> Blob {
let file = &self.files[row.file_index as usize];
let offset = (file.blobs + self.u32(row, column)) as usize;
let initial_byte = file.bytes[offset];
let (mut blob_size, blob_size_bytes) = match initial_byte >> 5 {
0..=3 => (initial_byte & 0x7f, 1),
4..=5 => (initial_byte & 0x3f, 2),
6 => (initial_byte & 0x1f, 4),
_ => panic!("Invalid blob size"),
};
for byte in &file.bytes[offset + 1..offset + blob_size_bytes] {
blob_size = blob_size.checked_shl(8).unwrap_or(0) + byte;
}
Blob::new(self, row.file_index, offset + blob_size_bytes)
}
pub(crate) fn equal_range(
&self,
file: u16,
table: TableIndex,
column: u32,
value: u32,
) -> impl Iterator<Item = Row> {
let (first, last) = self.equal_range_of(
table,
file,
0,
self.files[file as usize].tables[table as usize].row_count,
column,
value,
);
(first..last).map(move |row| Row::new(row, table, file))
}
fn lower_bound_of(
&self,
table: TableIndex,
file: u16,
mut first: u32,
last: u32,
column: u32,
value: u32,
) -> u32 {
let mut count = last - first;
while count > 0 {
let count2 = count / 2;
let middle = first + count2;
if self.u32(Row::new(middle, table, file), column) < value {
first = middle + 1;
count -= count2 + 1;
} else {
count = count2;
}
}
first
}
pub(crate) fn upper_bound(&self, file: u16, table: TableIndex, column: u32, value: u32) -> Row {
Row::new(
self.upper_bound_of(
table,
file,
0,
self.files[file as usize].tables[table as usize].row_count,
column,
value,
),
table,
file,
)
}
fn upper_bound_of(
&self,
table: TableIndex,
file: u16,
mut first: u32,
last: u32,
column: u32,
value: u32,
) -> u32 {
let mut count = last - first;
while count > 0 {
let count2 = count / 2;
let middle = first + count2;
if value < self.u32(Row::new(middle, table, file), column) {
count = count2
} else {
first = middle + 1;
count -= count2 + 1;
}
}
first
}
fn equal_range_of(
&self,
table: TableIndex,
file: u16,
mut first: u32,
mut last: u32,
column: u32,
value: u32,
) -> (u32, u32) {
let mut count = last - first;
loop {
if count == 0 {
last = first;
break;
}
let count2 = count / 2;
let middle = first + count2;
let middle_value = self.u32(Row::new(middle, table, file), column);
match middle_value.cmp(&value) {
Ordering::Less => {
first = middle + 1;
count -= count2 + 1;
}
Ordering::Greater => count = count2,
Ordering::Equal => {
let first2 = self.lower_bound_of(table, file, first, middle, column, value);
first += count;
last = self.upper_bound_of(table, file, middle + 1, first, column, value);
first = first2;
break;
}
}
}
(first, last)
}
}
#[derive(Debug, PartialEq)]
enum InsertMode {
All,
WinrtOnly,
}
#[cfg(target_pointer_width = "64")]
const SYSTEM32: &str = "System32";
#[cfg(target_pointer_width = "32")]
const SYSTEM32: &str = "SysNative";