Public release of PDB parser plugin

1 files changed, 499 insertions, 0 deletions

diff --git a/rust/examples/pdb-ng/src/parser.rs b/rust/examples/pdb-ng/src/parser.rs
new file mode 100644
index 00000000..cbc4f07a
--- /dev/null
+++ b/rust/examples/pdb-ng/src/parser.rs
@@ -0,0 +1,499 @@
+// Copyright 2022-2024 Vector 35 Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::{HashMap, HashSet};
+use std::env;
+use std::fmt::Display;
+
+use anyhow::{anyhow, Result};
+use log::{debug, info};
+use pdb::*;
+
+use binaryninja::architecture::{Architecture, CoreArchitecture};
+use binaryninja::binaryview::{BinaryView, BinaryViewExt};
+use binaryninja::callingconvention::CallingConvention;
+use binaryninja::debuginfo::{DebugFunctionInfo, DebugInfo};
+use binaryninja::platform::Platform;
+use binaryninja::rc::Ref;
+use binaryninja::settings::Settings;
+use binaryninja::types::{
+    min_confidence, Conf, DataVariableAndName, EnumerationBuilder, NamedTypeReference,
+    NamedTypeReferenceClass, StructureBuilder, StructureType, Type, TypeClass,
+};
+
+use crate::symbol_parser::{ParsedDataSymbol, ParsedProcedure, ParsedSymbol};
+use crate::type_parser::ParsedType;
+
+/// Megastruct for all the parsing
+/// Certain fields are only used by specific files, as marked below.
+/// Why not make new structs for them? Because vvvv this garbage
+pub struct PDBParserInstance<'a, S: Source<'a> + 'a> {
+    /// DebugInfo where types/functions will be stored eventually
+    pub(crate) debug_info: &'a mut DebugInfo,
+    /// Parent binary view (usually during BinaryView::Finalize)
+    pub(crate) bv: &'a BinaryView,
+    /// Default arch of self.bv
+    pub(crate) arch: CoreArchitecture,
+    /// Default calling convention for self.arch
+    pub(crate) default_cc: Ref<CallingConvention<CoreArchitecture>>,
+    /// Thiscall calling convention for self.bv, or default_cc if we can't find one
+    pub(crate) thiscall_cc: Ref<CallingConvention<CoreArchitecture>>,
+    /// Cdecl calling convention for self.bv, or default_cc if we can't find one
+    pub(crate) cdecl_cc: Ref<CallingConvention<CoreArchitecture>>,
+    /// Default platform of self.bv
+    pub(crate) platform: Ref<Platform>,
+    /// pdb-rs structure for making lifetime hell a real place
+    pub(crate) pdb: PDB<'a, S>,
+    /// pdb-rs Mapping of modules to addresses for resolving RVAs
+    pub(crate) address_map: AddressMap<'a>,
+    /// Binja Settings instance (for optimization)
+    pub(crate) settings: Ref<Settings>,
+
+    /// type_parser.rs
+
+    /// TypeIndex -> ParsedType enum used during parsing
+    pub(crate) indexed_types: HashMap<TypeIndex, ParsedType>,
+    /// QName -> Binja Type for finished types
+    pub(crate) named_types: HashMap<String, Ref<Type>>,
+    /// Raw (mangled) name -> TypeIndex for resolving forward references
+    pub(crate) full_type_indices: HashMap<String, TypeIndex>,
+    /// Stack of types we're currently parsing
+    pub(crate) type_stack: Vec<TypeIndex>,
+    /// Stack of parent types we're parsing nested types inside of
+    pub(crate) namespace_stack: Vec<String>,
+    /// Type Index -> Does it return on the stack
+    pub(crate) type_default_returnable: HashMap<TypeIndex, bool>,
+
+    /// symbol_parser.rs
+
+    /// List of fully parsed symbols from all modules
+    pub(crate) parsed_symbols: Vec<ParsedSymbol>,
+    /// Raw name -> index in parsed_symbols
+    pub(crate) parsed_symbols_by_name: HashMap<String, usize>,
+    /// Raw name -> Symbol index for looking up symbols for the currently parsing module (mostly for thunks)
+    pub(crate) named_symbols: HashMap<String, SymbolIndex>,
+    /// Parent -> Children symbol index tree for the currently parsing module
+    pub(crate) symbol_tree: HashMap<SymbolIndex, Vec<SymbolIndex>>,
+    /// Child -> Parent symbol index mapping, inverse of symbol_tree
+    pub(crate) symbol_parents: HashMap<SymbolIndex, SymbolIndex>,
+    /// Stack of (start, end) indices for the current symbols being parsed while constructing the tree
+    pub(crate) symbol_stack: Vec<(SymbolIndex, SymbolIndex)>,
+    /// Index -> parsed symbol for the currently parsing module
+    pub(crate) indexed_symbols: HashMap<SymbolIndex, ParsedSymbol>,
+    /// Symbol address -> Symbol for looking up by address
+    pub(crate) addressed_symbols: HashMap<u64, Vec<ParsedSymbol>>,
+    /// CPU type of the currently parsing module
+    pub(crate) module_cpu_type: Option<CPUType>,
+}
+
+impl<'a, S: Source<'a> + 'a> PDBParserInstance<'a, S> {
+    /// Try to create a new parser instance from a given bv/pdb
+    pub fn new(
+        debug_info: &'a mut DebugInfo,
+        bv: &'a BinaryView,
+        mut pdb: PDB<'a, S>,
+    ) -> Result<Self> {
+        let arch = if let Some(arch) = bv.default_arch() {
+            arch
+        } else {
+            return Err(anyhow!("Cannot parse to view with no architecture"));
+        };
+
+        let platform = bv
+            .default_platform()
+            .expect("Expected bv to have a platform");
+
+        let address_map = pdb.address_map()?;
+
+        let default_cc = platform
+            .get_default_calling_convention()
+            .expect("Expected default calling convention");
+
+        let thiscall_cc = Self::find_calling_convention(platform.as_ref(), "thiscall")
+            .unwrap_or(default_cc.clone());
+
+        let cdecl_cc = platform
+            .get_cdecl_calling_convention()
+            .unwrap_or(default_cc.clone());
+
+        Ok(Self {
+            debug_info,
+            bv,
+            arch,
+            default_cc,
+            thiscall_cc,
+            cdecl_cc,
+            platform,
+            pdb,
+            address_map,
+            settings: Settings::new(""),
+            indexed_types: Default::default(),
+            named_types: Default::default(),
+            full_type_indices: Default::default(),
+            type_stack: Default::default(),
+            namespace_stack: Default::default(),
+            type_default_returnable: Default::default(),
+            parsed_symbols: Default::default(),
+            parsed_symbols_by_name: Default::default(),
+            named_symbols: Default::default(),
+            symbol_tree: Default::default(),
+            symbol_parents: Default::default(),
+            symbol_stack: Default::default(),
+            indexed_symbols: Default::default(),
+            addressed_symbols: Default::default(),
+            module_cpu_type: None,
+        })
+    }
+
+    /// Try to parse the pdb into the DebugInfo
+    pub fn try_parse_info(
+        &mut self,
+        progress: Box<dyn Fn(usize, usize) -> Result<()> + 'a>,
+    ) -> Result<()> {
+        self.parse_types(Self::split_progress(&progress, 0, &[1.0, 3.0, 0.5, 0.5]))?;
+        for (name, ty) in self.named_types.iter() {
+            self.debug_info.add_type(name, ty.as_ref(), &[]); // TODO : Components
+        }
+
+        info!("PDB found {} types", self.named_types.len());
+
+        let (symbols, functions) =
+            self.parse_symbols(Self::split_progress(&progress, 1, &[1.0, 3.0, 0.5, 0.5]))?;
+
+        if self
+            .settings
+            .get_bool("pdb.features.createMissingNamedTypes", Some(self.bv), None)
+        {
+            self.resolve_missing_ntrs(
+                &symbols,
+                Self::split_progress(&progress, 2, &[1.0, 3.0, 0.5, 0.5]),
+            )?;
+            self.resolve_missing_ntrs(
+                &functions,
+                Self::split_progress(&progress, 3, &[1.0, 3.0, 0.5, 0.5]),
+            )?;
+        }
+
+        info!("PDB found {} data variables", symbols.len());
+        info!("PDB found {} functions", functions.len());
+
+        let allow_void =
+            self.settings
+                .get_bool("pdb.features.allowVoidGlobals", Some(self.bv), None);
+
+        for sym in symbols {
+            match sym {
+                ParsedSymbol::Data(ParsedDataSymbol {
+                    address,
+                    name,
+                    type_,
+                    ..
+                }) => {
+                    let real_type =
+                        type_.unwrap_or_else(|| Conf::new(Type::void(), min_confidence()));
+
+                    if real_type.contents.type_class() == TypeClass::VoidTypeClass {
+                        if !allow_void {
+                            self.log(|| {
+                                format!("Not adding void-typed symbol {:?}@{:x}", name, address)
+                            });
+                            continue;
+                        }
+                    }
+
+                    self.log(|| {
+                        format!(
+                            "Adding data variable: 0x{:x}: {} {:?}",
+                            address, &name.raw_name, real_type
+                        )
+                    });
+                    self.debug_info
+                        .add_data_variable_info(DataVariableAndName::new(
+                            address,
+                            real_type,
+                            true,
+                            name.full_name.unwrap_or(name.raw_name),
+                        ));
+                }
+                s => {
+                    self.log(|| format!("Not adding non-data symbol {:?}", s));
+                }
+            }
+        }
+
+        for sym in functions {
+            match sym {
+                ParsedSymbol::Procedure(ParsedProcedure {
+                    address,
+                    name,
+                    type_,
+                    ..
+                }) => {
+                    self.log(|| {
+                        format!(
+                            "Adding function: 0x{:x}: {} {:?}",
+                            address, &name.raw_name, type_
+                        )
+                    });
+                    self.debug_info.add_function(DebugFunctionInfo::new(
+                        Some(name.short_name.unwrap_or(name.raw_name.clone())),
+                        Some(name.full_name.unwrap_or(name.raw_name.clone())),
+                        Some(name.raw_name),
+                        type_.clone().and_then(|conf| {
+                            // TODO: When DebugInfo support confidence on function types, remove this
+                            if conf.confidence == 0 {
+                                None
+                            } else {
+                                Some(conf.contents)
+                            }
+                        }),
+                        Some(address),
+                        Some(self.platform.clone()),
+                        vec![], // TODO : Components
+                    ));
+                }
+                _ => {}
+            }
+        }
+
+        Ok(())
+    }
+
+    fn collect_name(
+        &self,
+        name: &NamedTypeReference,
+        unknown_names: &mut HashMap<String, NamedTypeReferenceClass>,
+    ) {
+        let used_name = name.name().to_string();
+        if let Some(&found) =
+            unknown_names.iter().find_map(
+                |(key, value)| {
+                    if key == &used_name {
+                        Some(value)
+                    } else {
+                        None
+                    }
+                },
+            )
+        {
+            if found != name.class() {
+                // Interesting case, not sure we care
+                self.log(|| {
+                    format!(
+                        "Mismatch unknown NTR class for {}: {} ?",
+                        &used_name,
+                        name.class() as u32
+                    )
+                });
+            }
+        } else {
+            self.log(|| format!("Found new unused name: {}", &used_name));
+            unknown_names.insert(used_name, name.class());
+        }
+    }
+
+    fn collect_names(
+        &self,
+        ty: &Type,
+        unknown_names: &mut HashMap<String, NamedTypeReferenceClass>,
+    ) {
+        match ty.type_class() {
+            TypeClass::StructureTypeClass => {
+                if let Ok(structure) = ty.get_structure() {
+                    if let Ok(members) = structure.members() {
+                        for member in members {
+                            self.collect_names(member.ty.contents.as_ref(), unknown_names);
+                        }
+                    }
+                    if let Ok(bases) = structure.base_structures() {
+                        for base in bases {
+                            self.collect_name(base.ty.as_ref(), unknown_names);
+                        }
+                    }
+                }
+            }
+            TypeClass::PointerTypeClass => {
+                if let Ok(target) = ty.target() {
+                    self.collect_names(target.contents.as_ref(), unknown_names);
+                }
+            }
+            TypeClass::ArrayTypeClass => {
+                if let Ok(element_type) = ty.element_type() {
+                    self.collect_names(element_type.contents.as_ref(), unknown_names);
+                }
+            }
+            TypeClass::FunctionTypeClass => {
+                if let Ok(return_value) = ty.return_value() {
+                    self.collect_names(return_value.contents.as_ref(), unknown_names);
+                }
+                if let Ok(params) = ty.parameters() {
+                    for param in params {
+                        self.collect_names(param.t.contents.as_ref(), unknown_names);
+                    }
+                }
+            }
+            TypeClass::NamedTypeReferenceClass => {
+                if let Ok(ntr) = ty.get_named_type_reference() {
+                    self.collect_name(ntr.as_ref(), unknown_names);
+                }
+            }
+            _ => {}
+        }
+    }
+
+    fn resolve_missing_ntrs(
+        &mut self,
+        symbols: &Vec<ParsedSymbol>,
+        progress: Box<dyn Fn(usize, usize) -> Result<()> + '_>,
+    ) -> Result<()> {
+        let mut unknown_names = HashMap::new();
+        let mut known_names = self
+            .bv
+            .types()
+            .iter()
+            .map(|qnat| qnat.name().string())
+            .collect::<HashSet<_>>();
+
+        for ty in &self.named_types {
+            known_names.insert(ty.0.clone());
+        }
+
+        let count = symbols.len();
+        for (i, sym) in symbols.into_iter().enumerate() {
+            match sym {
+                ParsedSymbol::Data(ParsedDataSymbol {
+                    type_: Some(type_), ..
+                }) => {
+                    self.collect_names(type_.contents.as_ref(), &mut unknown_names);
+                }
+                ParsedSymbol::Procedure(ParsedProcedure {
+                    type_: Some(type_),
+                    locals,
+                    ..
+                }) => {
+                    self.collect_names(type_.contents.as_ref(), &mut unknown_names);
+                    for l in locals {
+                        if let Some(ltype) = &l.type_ {
+                            self.collect_names(ltype.contents.as_ref(), &mut unknown_names);
+                        }
+                    }
+                }
+                _ => {}
+            }
+            (progress)(i, count)?;
+        }
+
+        for (name, class) in unknown_names.into_iter() {
+            if known_names.iter().any(|known| known == &name) {
+                self.log(|| format!("Found referenced name and ignoring: {}", &name));
+                continue;
+            }
+            self.log(|| format!("Adding referenced but unknown type {} (likely due to demangled name and stripped type)", &name));
+            match class {
+                NamedTypeReferenceClass::UnknownNamedTypeClass
+                | NamedTypeReferenceClass::TypedefNamedTypeClass => {
+                    self.debug_info.add_type(name, Type::void().as_ref(), &[]); // TODO : Components
+                }
+                NamedTypeReferenceClass::ClassNamedTypeClass
+                | NamedTypeReferenceClass::StructNamedTypeClass
+                | NamedTypeReferenceClass::UnionNamedTypeClass => {
+                    let structure = StructureBuilder::new();
+                    match class {
+                        NamedTypeReferenceClass::ClassNamedTypeClass => {
+                            structure.set_structure_type(StructureType::ClassStructureType);
+                        }
+                        NamedTypeReferenceClass::StructNamedTypeClass => {
+                            structure.set_structure_type(StructureType::StructStructureType);
+                        }
+                        NamedTypeReferenceClass::UnionNamedTypeClass => {
+                            structure.set_structure_type(StructureType::UnionStructureType);
+                        }
+                        _ => {}
+                    }
+                    structure.set_width(1);
+                    structure.set_alignment(1);
+
+                    self.debug_info.add_type(
+                        name,
+                        Type::structure(structure.finalize().as_ref()).as_ref(),
+                        &[], // TODO : Components
+                    );
+                }
+                NamedTypeReferenceClass::EnumNamedTypeClass => {
+                    let enumeration = EnumerationBuilder::new();
+                    self.debug_info.add_type(
+                        name,
+                        Type::enumeration(
+                            enumeration.finalize().as_ref(),
+                            self.arch.default_integer_size(),
+                            false,
+                        )
+                        .as_ref(),
+                        &[], // TODO : Components
+                    );
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Lazy logging function that prints like 20MB of messages
+    pub(crate) fn log<F: FnOnce() -> D, D: Display>(&self, msg: F) {
+        if env::var("BN_DEBUG_PDB").is_ok() {
+            let space = "\t".repeat(self.type_stack.len()) + &"\t".repeat(self.symbol_stack.len());
+            let msg = format!("{}", msg());
+            debug!(
+                "{}{}",
+                space,
+                msg.replace("\n", &*("\n".to_string() + &space))
+            );
+        }
+    }
+
+    pub(crate) fn split_progress<'b, F: Fn(usize, usize) -> Result<()> + 'b>(
+        original_fn: F,
+        subpart: usize,
+        subpart_weights: &[f64],
+    ) -> Box<dyn Fn(usize, usize) -> Result<()> + 'b> {
+        // Normalize weights
+        let weight_sum: f64 = subpart_weights.iter().sum();
+        if weight_sum < 0.0001 {
+            return Box::new(|_, _| Ok(()));
+        }
+
+        // Keep a running count of weights for the start
+        let mut subpart_starts = vec![];
+        let mut start = 0f64;
+        for w in subpart_weights {
+            subpart_starts.push(start);
+            start += *w;
+        }
+
+        let subpart_start = subpart_starts[subpart] / weight_sum;
+        let weight = subpart_weights[subpart] / weight_sum;
+
+        Box::new(move |cur: usize, max: usize| {
+            // Just use a large number for easy divisibility
+            let steps = 1000000f64;
+            let subpart_size = steps * weight;
+            let subpart_progress = ((cur as f64) / (max as f64)) * subpart_size;
+
+            original_fn(
+                (subpart_start * steps + subpart_progress) as usize,
+                steps as usize,
+            )
+        })
+    }
+}