From d8a4727d59e4aa8234fe17588480b06ebfbbd8ac Mon Sep 17 00:00:00 2001
From: Jean-Marie Mineau <jean-marie.mineau@inria.fr>
Date: Fri, 9 Feb 2024 18:01:19 +0100
Subject: [PATCH] WIP: fix data size

---
 androscalpel/src/dex_writer.rs             | 11 ++++++++++-
 androscalpel_serializer/src/file_reader.rs | 18 ++++++++++++++++--
 test.py                                    |  2 +-
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/androscalpel/src/dex_writer.rs b/androscalpel/src/dex_writer.rs
index 59c2e4c..19a36b5 100644
--- a/androscalpel/src/dex_writer.rs
+++ b/androscalpel/src/dex_writer.rs
@@ -2498,7 +2498,9 @@ impl DexWriter {
         }
         self.gen_type_list_section()?;
 
-        self.gen_map_list()?;
+        self.gen_map_list()?; // TODO TODO TODO not good, the values are not yes set because:
+                              // - alignment
+                              // - F***ing class_data that change size during linking
 
         // From now on, all section are generated and the value in section_manager do not change,
         // except for class data items, because F (see doc of self.corrected_class_data_offset).
@@ -2577,6 +2579,7 @@ impl DexWriter {
             handle.serialize(&mut buffer)?;
         }
         // MapList
+        self.check_section_offset(&buffer, Section::Data);
         self.check_section_offset(&buffer, Section::MapList);
         Self::fix_section_alignement(&mut buffer, Section::MapList)?;
         self.map_list.serialize(&mut buffer)?;
@@ -2646,6 +2649,12 @@ impl DexWriter {
         Self::fix_section_alignement(&mut buffer, Section::HiddenapiClassDataItem)?;
         */
 
+        let end_data = buffer.position();
+        assert_eq!(
+            end_data as u32,
+            self.header.data_off + self.header.data_size
+        );
+
         // compute signature
         buffer.seek(SeekFrom::Start(8 + 4 + 20))?;
         let mut hasher = Sha1::new();
diff --git a/androscalpel_serializer/src/file_reader.rs b/androscalpel_serializer/src/file_reader.rs
index 06e94a3..5be7138 100644
--- a/androscalpel_serializer/src/file_reader.rs
+++ b/androscalpel_serializer/src/file_reader.rs
@@ -377,12 +377,26 @@ impl<'a> DexFileReader<'a> {
     pub fn get_struct_at_offset<T: Serializable>(&self, offset: u32) -> Result<T> {
         let mut buffer = Cursor::new(self.data);
         buffer.seek(SeekFrom::Start(offset as u64)).unwrap();
-        T::deserialize(&mut buffer).map_err(|err| {
+        let r = T::deserialize(&mut buffer).map_err(|err| {
             Error::DeserializationError(format!(
                 "Failed to deserialize {} at 0x{:x}: {err}",
                 std::any::type_name::<T>(),
                 offset
             ))
-        })
+        });
+        if buffer.position() as u32 > self.header.data_off + self.header.data_size {
+            // Return error? Android won't run an apk that does not respect this condition for most
+            // (all?) struct, but a lot a tools don't care.
+            warn!(
+                "Deserialized {} at 0x{:x}, but ended up reading out of the data section \
+                (0x{:x} to 0x{:x}, current cursor at 0x{:x}",
+                std::any::type_name::<T>(),
+                offset,
+                self.header.data_off,
+                self.header.data_off + self.header.data_size,
+                buffer.position()
+            );
+        }
+        r
     }
 }
diff --git a/test.py b/test.py
index fa28b53..ab31c9f 100644
--- a/test.py
+++ b/test.py
@@ -30,7 +30,7 @@ clazz = apk.classes[clazz_id]
 method = clazz.virtual_methods[method_id]
 code = method.code
 
-logging.getLogger().setLevel(logging.ERROR)
+logging.getLogger().setLevel(logging.WARNING)
 
 print(f"[+] Code of {method_id} ")
 for i in code.insns: