perf: optimize parsing of imports and exports

Pre-allocate vector capacities in `parse_imports` and `parse_exports` to reduce reallocations and improve performance. Introduce an `ImportDll` enum to clarify and simplify logic for mapping import ordinals to names, avoiding repeated string lowercasing.

Author: plusvic

lib/src/modules/pe/parser.rs

@@ -4,7 +4,7 @@ use std::collections::{HashMap, VecDeque};
 use std::default::Default;
 use std::iter::zip;
 use std::mem;
-use std::str::{from_utf8, FromStr};
+use std::str::{FromStr, from_utf8};
 use std::sync::OnceLock;
 
 use bstr::{BStr, ByteSlice};
@@ -14,11 +14,11 @@ use memchr::memmem;
 use nom::branch::{alt, permutation};
 use nom::bytes::complete::{take, take_till, take_while_m_n};
 use nom::combinator::{
-    cond, consumed, iterator, map, opt, success, verify, Success,
+    Success, cond, consumed, iterator, map, opt, success, verify,
 };
 use nom::error::ErrorKind;
 use nom::multi::{
-    count, fold_many0, fold_many1, length_data, many0, many1, many_m_n,
+    count, fold_many0, fold_many1, length_data, many_m_n, many0, many1,
 };
 use nom::number::complete::{le_u16, le_u32, le_u64, u8};
 use nom::{Err, IResult, Parser, ToUsize};
@@ -905,12 +905,11 @@ impl<'a> PE<'a> {
                     .ok()
                     .and_then(|name| name.strip_prefix('/'))
                     .and_then(|offset| u32::from_str(offset).ok())
-                    && let Some(s) = string_table.get(offset as usize..)
-                        && let Ok((_, s)) =
-                            take_till::<_, &[u8], Error>(|c| c == 0)(s)
-                        {
-                            section.full_name = Some(BStr::new(s));
-                        }
+                && let Some(s) = string_table.get(offset as usize..)
+                && let Ok((_, s)) = take_till::<_, &[u8], Error>(|c| c == 0)(s)
+            {
+                section.full_name = Some(BStr::new(s));
+            }
 
             Ok((remainder, section))
         }
@@ -1473,23 +1472,23 @@ impl<'a> PE<'a> {
                         // value as the upper bound and avoid some completely
                         // corrupt entries with random values.
                         && (rsrc_entry.size as usize) < 0x3FFFFFFF
-                        {
-                            resources.push(Resource {
-                                type_id: ids.0,
-                                rsrc_id: ids.1,
-                                lang_id: ids.2,
-                                // `rsrc_entry.offset` is relative to the start of
-                                // the resource section, so it's actually an RVA.
-                                // Here we convert it to a file offset.
-                                offset: self.rva_to_offset(rsrc_entry.offset),
-                                rva: rsrc_entry.offset,
-                                length: rsrc_entry.size,
-                            });
-
-                            if resources.len() == Self::MAX_PE_RESOURCES {
-                                return Some((resources_info, resources));
-                            }
+                    {
+                        resources.push(Resource {
+                            type_id: ids.0,
+                            rsrc_id: ids.1,
+                            lang_id: ids.2,
+                            // `rsrc_entry.offset` is relative to the start of
+                            // the resource section, so it's actually an RVA.
+                            // Here we convert it to a file offset.
+                            offset: self.rva_to_offset(rsrc_entry.offset),
+                            rva: rsrc_entry.offset,
+                            length: rsrc_entry.size,
+                        });
+
+                        if resources.len() == Self::MAX_PE_RESOURCES {
+                            return Some((resources_info, resources));
                         }
+                    }
                 }
             }
         }
@@ -1525,8 +1524,10 @@ impl<'a> PE<'a> {
     /// Returns a parser that parses a WIN_CERTIFICATE structure.
     fn win_cert_parser(
         &self,
-    ) -> impl FnMut(&'a [u8]) -> IResult<&'a [u8], Vec<AuthenticodeSignature<'a>>>
-           + '_ {
+    ) -> impl FnMut(
+        &'a [u8],
+    ) -> IResult<&'a [u8], Vec<AuthenticodeSignature<'a>>>
+    + '_ {
         move |input: &'a [u8]| {
             // Parse the WIN_CERTIFICATE structure.
             let (remainder, (length, _revision, _cert_type)) = (
@@ -1556,8 +1557,10 @@ impl<'a> PE<'a> {
     /// Authenticode signature.
     fn signature_parser(
         &self,
-    ) -> impl FnMut(&'a [u8]) -> IResult<&'a [u8], Vec<AuthenticodeSignature<'a>>>
-           + '_ {
+    ) -> impl FnMut(
+        &'a [u8],
+    ) -> IResult<&'a [u8], Vec<AuthenticodeSignature<'a>>>
+    + '_ {
         move |input: &'a [u8]| {
             let signatures = AuthenticodeParser::parse(input, self)
                 .map_err(|_| Err::Error(Error::new(input, ErrorKind::Fail)))?;
@@ -1672,7 +1675,7 @@ impl<'a> PE<'a> {
             .parse(cv_info)
             {
                 Ok((_, (_signature, _padding, pdb_path))) => {
-                    return Some(pdb_path)
+                    return Some(pdb_path);
                 }
                 Err(_) => continue,
             };
@@ -1793,7 +1796,9 @@ impl<'a> PE<'a> {
         let is_32_bits =
             self.optional_hdr.magic != Self::IMAGE_NT_OPTIONAL_HDR64_MAGIC;
 
-        let mut imported_funcs = Vec::new();
+        let estimated_descriptors =
+            min(input.len() / Self::SIZE_OF_DIR_ENTRY, Self::MAX_PE_IMPORTS);
+        let mut imported_funcs = Vec::with_capacity(estimated_descriptors);
 
         // Parse import descriptors until finding one that is empty (filled
         // with null values), which indicates the end of the directory table;
@@ -1840,6 +1845,16 @@ impl<'a> PE<'a> {
                     continue;
                 };
 
+            let import_dll = if dll_name.eq_ignore_ascii_case("ws2_32.dll")
+                || dll_name.eq_ignore_ascii_case("wsock32.dll")
+            {
+                ImportDll::Wsock32
+            } else if dll_name.eq_ignore_ascii_case("oleaut32.dll") {
+                ImportDll::Oleaut32
+            } else {
+                ImportDll::Other
+            };
+
             // Use the INT (a.k.a: OriginalFirstThunk) if it is non-zero, but
             // fallback to using the IAT (a.k.a: FirstThunk).
             let thunks = if descriptor.import_name_table > 0 {
@@ -1849,20 +1864,25 @@ impl<'a> PE<'a> {
             }
             .or_else(|| self.data_at_rva(descriptor.import_address_table));
 
-            let thunks = match thunks {
+            let thunks_slice = match thunks {
                 Some(thunk) => thunk,
                 None => continue,
             };
 
+            let estimated_funcs = min(
+                thunks_slice.len() / if is_32_bits { 4 } else { 8 },
+                Self::MAX_PE_IMPORTS,
+            );
+
             // Parse the thunks, which are an array of 64-bits or 32-bits
             // values, depending on whether this is 64-bits PE file. The
             // array is terminated by a null thunk.
             let thunks = iterator(
-                thunks,
+                thunks_slice,
                 verify(uint(is_32_bits), |thunk| *thunk != 0),
             );
 
-            let mut funcs = Vec::new();
+            let mut funcs = Vec::with_capacity(estimated_funcs);
 
             for (i, mut thunk) in
                 &mut thunks.take(Self::MAX_PE_IMPORTS).enumerate()
@@ -1897,7 +1917,7 @@ impl<'a> PE<'a> {
                 if import_by_ordinal {
                     let ordinal = (thunk & 0xffff) as u16;
                     func.ordinal = Some(ordinal);
-                    func.name = ord_to_name(dll_name, ordinal);
+                    func.name = ord_to_name(import_dll, ordinal);
                 } else {
                     // When descriptor values are virtual addresses, thunks are
                     // virtual addresses too and need to be converted to RVAs.
@@ -2109,17 +2129,16 @@ impl<'a> PE<'a> {
 
         // Create a vector with one item per exported function. Items in the
         // array initially have function RVA and ordinal only.
-        let mut exported_funcs: Vec<_> = func_rvas
-            .take(num_exports)
-            .enumerate()
-            .filter_map(|(i, rva)| {
-                Some(ExportedFunc {
+        let mut exported_funcs = Vec::with_capacity(num_exports);
+        for (i, rva) in func_rvas.take(num_exports).enumerate() {
+            if let Some(ordinal) = exports.base.checked_add(i as u32) {
+                exported_funcs.push(ExportedFunc {
                     rva,
-                    ordinal: exports.base.checked_add(i as u32)?,
+                    ordinal,
                     ..Default::default()
-                })
-            })
-            .collect();
+                });
+            }
+        }
 
         let names = self
             .parse_at_rva(exports.address_of_names, count(le_u32, num_names))
@@ -2139,10 +2158,11 @@ impl<'a> PE<'a> {
                     .find_position(|ordinal| {
                         *ordinal as u32 == f.ordinal - exports.base
                     })
-                && let Some(name_rva) = names.get(idx) {
-                    f.name =
-                        self.str_at_rva(*name_rva, Self::MAX_FUNC_NAME_LENGTH);
-                }
+                && let Some(name_rva) = names.get(idx)
+            {
+                f.name =
+                    self.str_at_rva(*name_rva, Self::MAX_FUNC_NAME_LENGTH);
+            }
 
             // If the function's RVA is within the exports section (as given
             // by the RVA and size fields in the directory entry), this is a
@@ -2822,16 +2842,23 @@ fn utf16_le_string() -> impl FnMut(&[u8]) -> IResult<&[u8], String> {
     }
 }
 
+#[derive(Copy, Clone, PartialEq, Eq)]
+enum ImportDll {
+    Wsock32,
+    Oleaut32,
+    Other,
+}
+
 /// Convert ordinal number to function name.
 ///
 /// For some well-known DLLs the returned name is the one that that corresponds
 /// to the given ordinal. For the remaining DLLs the returned name has the form
 /// "ordN" where N is the ordinal (e.g: "ord1", "ord23").
-fn ord_to_name(dll_name: &str, ordinal: u16) -> Option<String> {
-    let func_name = match dll_name.to_ascii_lowercase().as_str() {
-        "ws2_32.dll" | "wsock32.dll" => wsock32_ord_to_name(ordinal),
-        "oleaut32.dll" => oleaut32_ord_to_name(ordinal),
-        _ => None,
+fn ord_to_name(well_known: ImportDll, ordinal: u16) -> Option<String> {
+    let func_name = match well_known {
+        ImportDll::Wsock32 => wsock32_ord_to_name(ordinal),
+        ImportDll::Oleaut32 => oleaut32_ord_to_name(ordinal),
+        ImportDll::Other => None,
     };
 
     func_name.map(|n| n.to_owned()).or_else(|| Some(format!("ord{ordinal}")))