From 4e975ed4499a37ff6257eb28e667e6f1902e6f32 Mon Sep 17 00:00:00 2001
From: kkent030315 <hrn832@protonmail.com>
Date: Mon, 11 Nov 2024 07:26:58 +0900
Subject: [PATCH 1/2] PE: TLS parser rework

---
 src/pe/mod.rs | 28 +++++++++-------------------
 src/pe/tls.rs | 30 +++++++++++++++++++-----------
 2 files changed, 28 insertions(+), 30 deletions(-)

diff --git a/src/pe/mod.rs b/src/pe/mod.rs
index e921f2e2..ea5f431f 100644
--- a/src/pe/mod.rs
+++ b/src/pe/mod.rs
@@ -221,25 +221,15 @@ impl<'a> PE<'a> {
             }
 
             if let Some(tls_table) = optional_header.data_directories.get_tls_table() {
-                tls_data = if is_64 {
-                    tls::TlsData::parse_with_opts::<u64>(
-                        bytes,
-                        image_base,
-                        tls_table,
-                        &sections,
-                        file_alignment,
-                        opts,
-                    )?
-                } else {
-                    tls::TlsData::parse_with_opts::<u32>(
-                        bytes,
-                        image_base,
-                        &tls_table,
-                        &sections,
-                        file_alignment,
-                        opts,
-                    )?
-                };
+                tls_data = tls::TlsData::parse_with_opts(
+                    bytes,
+                    image_base,
+                    tls_table,
+                    &sections,
+                    file_alignment,
+                    opts,
+                    is_64,
+                )?;
                 debug!("tls data: {:#?}", tls_data);
             }
 
diff --git a/src/pe/tls.rs b/src/pe/tls.rs
index 2369bb0e..001d7d4e 100644
--- a/src/pe/tls.rs
+++ b/src/pe/tls.rs
@@ -45,27 +45,30 @@ pub struct TlsData<'a> {
 }
 
 impl ImageTlsDirectory {
-    pub fn parse<T: Sized>(
+    pub fn parse(
         bytes: &[u8],
         dd: data_directories::DataDirectory,
         sections: &[section_table::SectionTable],
         file_alignment: u32,
+        is_64: bool,
     ) -> error::Result<Self> {
-        Self::parse_with_opts::<T>(
+        Self::parse_with_opts(
             bytes,
             dd,
             sections,
             file_alignment,
             &options::ParseOptions::default(),
+            is_64,
         )
     }
 
-    pub fn parse_with_opts<T: Sized>(
+    pub fn parse_with_opts(
         bytes: &[u8],
         dd: data_directories::DataDirectory,
         sections: &[section_table::SectionTable],
         file_alignment: u32,
         opts: &options::ParseOptions,
+        is_64: bool,
     ) -> error::Result<Self> {
         let rva = dd.virtual_address as usize;
         let mut offset =
@@ -76,8 +79,6 @@ impl ImageTlsDirectory {
                 ))
             })?;
 
-        let is_64 = core::mem::size_of::<T>() == 8;
-
         let start_address_of_raw_data = if is_64 {
             bytes.gread_with::<u64>(&mut offset, scroll::LE)?
         } else {
@@ -115,39 +116,40 @@ impl ImageTlsDirectory {
 }
 
 impl<'a> TlsData<'a> {
-    pub fn parse<T: Sized>(
+    pub fn parse(
         bytes: &'a [u8],
         image_base: usize,
         dd: &data_directories::DataDirectory,
         sections: &[section_table::SectionTable],
         file_alignment: u32,
+        is_64: bool,
     ) -> error::Result<Option<Self>> {
-        Self::parse_with_opts::<T>(
+        Self::parse_with_opts(
             bytes,
             image_base,
             dd,
             sections,
             file_alignment,
             &options::ParseOptions::default(),
+            is_64,
         )
     }
 
-    pub fn parse_with_opts<T: Sized>(
+    pub fn parse_with_opts(
         bytes: &'a [u8],
         image_base: usize,
         dd: &data_directories::DataDirectory,
         sections: &[section_table::SectionTable],
         file_alignment: u32,
         opts: &options::ParseOptions,
+        is_64: bool,
     ) -> error::Result<Option<Self>> {
         let mut raw_data = None;
         let mut slot = None;
         let mut callbacks = Vec::new();
 
-        let is_64 = core::mem::size_of::<T>() == 8;
-
         let itd =
-            ImageTlsDirectory::parse_with_opts::<T>(bytes, *dd, sections, file_alignment, opts)?;
+            ImageTlsDirectory::parse_with_opts(bytes, *dd, sections, file_alignment, opts, is_64)?;
 
         // Parse the raw data if any
         if itd.end_address_of_raw_data != 0 && itd.start_address_of_raw_data != 0 {
@@ -176,6 +178,12 @@ impl<'a> TlsData<'a> {
                         rva
                     ))
                 })?;
+            if offset + size as usize > bytes.len() {
+                return Err(error::Error::Malformed(format!(
+                    "tls raw data offset ({:#x}) and size ({:#x}) greater than byte slice len ({:#x})",
+                    offset,size,bytes.len()
+                )));
+            }
             raw_data = Some(&bytes[offset..offset + size as usize]);
         }
 

From 63f47dfba934dd86d9a99cc1089f74e5c123a5d7 Mon Sep 17 00:00:00 2001
From: kkent030315 <hrn832@protonmail.com>
Date: Mon, 11 Nov 2024 07:40:46 +0900
Subject: [PATCH 2/2] Format

---
 src/pe/tls.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/pe/tls.rs b/src/pe/tls.rs
index 001d7d4e..ae66bc24 100644
--- a/src/pe/tls.rs
+++ b/src/pe/tls.rs
@@ -181,7 +181,7 @@ impl<'a> TlsData<'a> {
             if offset + size as usize > bytes.len() {
                 return Err(error::Error::Malformed(format!(
                     "tls raw data offset ({:#x}) and size ({:#x}) greater than byte slice len ({:#x})",
-                    offset,size,bytes.len()
+                    offset, size, bytes.len()
                 )));
             }
             raw_data = Some(&bytes[offset..offset + size as usize]);