wip

androscalpel/src/dex_fragment.rs

@@ -45,12 +45,6 @@ pub struct DexFragment {
     class_def: ClassDefItem,
     /// The class_data.
     class_data: Option<ClassDataItem>,
-    // TODO: type list should be handle like other ids
-    // TODO: type list inside of proto ids are not handled here
-    // /// The type lists found in the classes associated to their index in the type_lists section.
-    // type_lists_index: HashMap<TypeList, usize>,
-    // /// The type_lists section and the offset of the lists inside the section.
-    // type_lists_with_offset: Vec<(TypeList, u32)>,
     /// The encoded_array_items section.
     encoded_array_items: Vec<EncodedArrayItem>,
     /// The annotations_directory_item.
@@ -69,9 +63,40 @@ pub struct DexFragment {
     debug_info_items: Vec<DebugInfoItem>,
     /// The list of interfaces of the class.
     interfaces: Vec<IdType>,
+    /// The current link state of the fragment.
+    link_state: FragLinkState,
 }
 
 impl DexFragment {
+    /// The strings in the dex fragment, sorted.
+    pub fn strings(&self) -> &[DexString] {
+        &self.strings
+    }
+    /// The types in the dex fragment, sorted.
+    pub fn type_ids(&self) -> &[IdType] {
+        &self.type_ids
+    }
+    /// The prototypes in the dex fragment, sorted.
+    pub fn proto_ids(&self) -> &[IdMethodType] {
+        &self.proto_ids
+    }
+    /// The field ids in the dex fragment, sorted.
+    pub fn field_ids(&self) -> &[IdField] {
+        &self.field_ids
+    }
+    /// The methods ids in the dex fragment, sorted.
+    pub fn method_ids(&self) -> &[IdMethod] {
+        &self.method_ids
+    }
+    /// The list of interfaces of the class.
+    pub fn interfaces(&self) -> &[IdType] {
+        &self.interfaces
+    }
+    /// The method_handles section.
+    pub fn method_handles(&self) -> &[MethodHandleItem] {
+        &self.method_handles
+    }
+
     pub fn new(class: &Class) -> Result<Self> {
         debug!(
             "Building dex fragment for class {}",
@@ -108,6 +133,7 @@ impl DexFragment {
             annotation_set_lists: vec![],
             debug_info_items: vec![],
             interfaces: class.interfaces.clone(),
+            link_state: FragLinkState::Unlinked,
         };
         frag.strings = class.get_all_strings().into_iter().collect();
         frag.strings.sort();
@@ -1852,6 +1878,168 @@ impl DexFragment {
         self.annotation_set_lists.push(list);
         Ok(())
     }
+
+    /// Some structure change size depending of the value of the index they refere to.
+    /// The new size needs to be known before the offset can be linked, so the idx are
+    /// linked before.
+    pub fn link_global_ids(
+        &mut self,
+        global_strings: &[DexString],
+        global_type_ids: &[IdType],
+        global_proto_ids: &[IdMethodType],
+        global_field_ids: &[IdField],
+        global_method_ids: &[IdMethod],
+        nb_method_handle_before_fragment: usize,
+    ) -> Result<()> {
+        self.link_state.start_linking_idx()?;
+        let string_reindex = Vec::with_capacity(self.strings.len());
+        let mut global_idx = 0;
+        for s in self.strings {
+            while global_idx < global_strings.len() && global_strings[global_idx] != s {
+                global_idx += 1;
+            }
+            if global_idx == global_strings.len() {
+                bail!("String {} not found in global index", s.__str__());
+            }
+            string_reindex.push(global_idx as u32);
+        }
+        let type_reindex = Vec::with_capacity(self.type_ids.len());
+        let mut global_idx = 0;
+        for ty in self.type_ids {
+            while global_idx < global_type_ids.len() && global_type_ids[global_idx] != ty {
+                global_idx += 1;
+            }
+            if global_idx == global_type_ids.len() {
+                bail!("Type {} not found in global index", ty.__str__());
+            }
+            type_reindex.push(global_idx as u32);
+        }
+        let proto_reindex = Vec::with_capacity(self.proto_ids.len());
+        let mut global_idx = 0;
+        for proto in self.proto_ids {
+            while global_idx < global_proto_ids.len() && global_proto_ids[global_idx] != proto {
+                global_idx += 1;
+            }
+            if global_idx == global_proto_ids.len() {
+                bail!("Prototype {} not found in global index", proto.__str__());
+            }
+            proto_reindex.push(global_idx as u32);
+        }
+        let field_reindex = Vec::with_capacity(self.field_ids.len());
+        let mut global_idx = 0;
+        for field in self.field_ids {
+            while global_idx < global_field_ids.len() && global_field_ids[global_idx] != field {
+                global_idx += 1;
+            }
+            if global_idx == global_field_ids.len() {
+                bail!("Field {} not found in global index", field.__str__());
+            }
+            field_reindex.push(global_idx as u16);
+        }
+        let method_reindex = Vec::with_capacity(self.method_ids.len());
+        let mut global_idx = 0;
+        for meth in self.method_ids {
+            while global_idx < global_method_ids.len() && global_method_ids[global_idx] != meth {
+                global_idx += 1;
+            }
+            if global_idx == global_method_ids.len() {
+                bail!("Method {} not found in global index", meth.__str__());
+            }
+            method_reindex.push(global_idx as u16);
+        }
+
+        self.link_id_class_def(&string_reindex, &type_reindex);
+        self.link_id_method_handle(&field_reindex, &method_reindex);
+
+        todo!()
+    }
+
+    fn link_id_class_def(&mut self, string_reindex: &[u32], type_reindex: &[u32]) {
+        self.class_def.class_idx = type_reindex[self.class_def.class_idx as usize];
+        self.class_def.superclass_idx = type_reindex[self.class_def.superclass_idx as usize];
+        self.class_def.source_file_idx = string_reindex[self.class_def.source_file_idx as usize];
+    }
+
+    fn link_id_method_handle(&mut self, field_reindex: &[u16], method_reindex: &[u16]) {
+        for handle in self.method_handles {
+            let reindexer = match handle.method_handle_type {
+                MethodHandleType::StaticPut => field_reindex,
+                MethodHandleType::StaticGet => field_reindex,
+                MethodHandleType::InstancePut => field_reindex,
+                MethodHandleType::InstanceGet => field_reindex,
+                MethodHandleType::InvokeStatic => method_reindex,
+                MethodHandleType::InvokeInstance => method_reindex,
+                MethodHandleType::InvokeConstructor => method_reindex,
+                MethodHandleType::InvokeDirect => method_reindex,
+                MethodHandleType::InvokeInterface => method_reindex,
+            };
+            handle.field_or_method_id = reindexer[handle.field_or_method_id as usize];
+        }
+    }
+
+    fn link_id_code(&mut self, string_reindex: &[u32], type_reindex: &[u32]) {
+        let mut total_size = 0;
+        let mut code_item_relocation = if let FragLinkState::LinkedIdx {
+            code_item_relocation,
+            ..
+        } = self.link_state
+        {
+            code_item_relocation
+        } else {
+            // link_global_ids() should prevent that
+            panic!("link_id_code should not be run outside of fn link_global_ids(..)");
+        };
+        for code in self.code_items {
+            let current_size = code.size();
+            for ins in &mut code.insns {
+                Self::link_id_ins(ins, string_reindex);
+            }
+            // TODO: TryItem recompute handler_off
+            if let Some(handlers) = code.handlers {
+                let mut handler_off_reindex = HashMap::new();
+                let mut current_offset = handlers.size_field().size();
+                let mut old_offset = handlers.size_field().size();
+                for handlers in handlers.list {
+                    handler_off_reindex.insert(old_offset as u16, current_offset as u16);
+                    old_offset += handlers.size();
+                    for handler in handlers.handlers {
+                        handler.type_idx.0 = type_reindex[handler.type_idx.0 as usize];
+                    }
+                    current_offset += handlers.size();
+                }
+                for try_ in code.tries {
+                    try_.handler_off = *handler_off_reindex
+                        .get(&try_.handler_off)
+                        .expect("Something whent wrong with the handle reindexing");
+                }
+            }
+        }
+    }
+    fn link_id_ins(ins: &mut InsFormat, string_reindex: &[u32]) {
+        match ins {
+            InsFormat::Format31C { op: 0x1b, b, .. } => *b = string_reindex[b as usize],
+            InsFormat::Format21C { op: 0x1a, b, .. } => todo!(), // TODO FUCK this
+            _ => todo!(),
+        }
+    }
+    fn link_id_string_data(&mut self) {
+        todo!()
+    }
+    fn link_id_debug_info(&mut self) {
+        todo!()
+    }
+    fn link_id_annotation(&mut self) {
+        todo!()
+    }
+    fn link_id_encoded_array(&mut self) {
+        todo!()
+    }
+    fn link_id_annotation_dir(&mut self) {
+        todo!()
+    }
+    fn link_id_class_data(&mut self) {
+        todo!()
+    }
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -1859,7 +2047,6 @@ enum FragSection {
     ClassDefItem,
     CallSiteIdItem,
     MethodHandleItem,
-    TypeList,
     AnnotationSetRefList,
     AnnotationSetItem,
     CodeItem,
@@ -1877,7 +2064,6 @@ impl FragSection {
         Self::ClassDefItem,
         Self::CallSiteIdItem,
         Self::MethodHandleItem,
-        Self::TypeList,
         Self::AnnotationSetRefList,
         Self::AnnotationSetItem,
         Self::CodeItem,
@@ -1886,8 +2072,7 @@ impl FragSection {
         Self::AnnotationItem,
         Self::EncodedArrayItem,
         Self::AnnotationsDirectoryItem,
-        // must be last because contains offsets in Uleb,
-        // so size change when linking !
+        // /!\ contains offsets in Uleb, so size change when linking !
         Self::ClassDataItem,
         // Self::HiddenapiClassDataItem,
     ];
@@ -1897,17 +2082,16 @@ impl FragSection {
             Self::ClassDefItem => 0,
             Self::CallSiteIdItem => 1,
             Self::MethodHandleItem => 2,
-            Self::TypeList => 3,
-            Self::AnnotationSetRefList => 4,
-            Self::AnnotationSetItem => 5,
-            Self::CodeItem => 6,
-            Self::StringDataItem => 7,
-            Self::DebugInfoItem => 8,
-            Self::AnnotationItem => 9,
-            Self::EncodedArrayItem => 10,
-            Self::AnnotationsDirectoryItem => 11,
-            Self::ClassDataItem => 12,
-            // Self::HiddenapiClassDataItem => 13,
+            Self::AnnotationSetRefList => 3,
+            Self::AnnotationSetItem => 4,
+            Self::CodeItem => 5,
+            Self::StringDataItem => 6,
+            Self::DebugInfoItem => 7,
+            Self::AnnotationItem => 8,
+            Self::EncodedArrayItem => 9,
+            Self::AnnotationsDirectoryItem => 10,
+            Self::ClassDataItem => 11,
+            // Self::HiddenapiClassDataItem => 12,
         }
     }
 
@@ -1916,7 +2100,6 @@ impl FragSection {
             Self::ClassDefItem => Some(0x20),
             Self::CallSiteIdItem => Some(4),
             Self::MethodHandleItem => Some(8),
-            Self::TypeList => None,
             Self::AnnotationSetRefList => None,
             Self::AnnotationSetItem => None,
             Self::CodeItem => None,
@@ -1953,7 +2136,6 @@ impl FragSection {
             Self::ClassDefItem => 4,
             Self::CallSiteIdItem => 1,
             Self::MethodHandleItem => 4,
-            Self::TypeList => 4,
             Self::AnnotationSetRefList => 4,
             Self::AnnotationSetItem => 4,
             Self::CodeItem => 4,
@@ -2124,6 +2306,7 @@ impl FragSectionManager {
 }
 
 /// Index that associate a type to its local id in a fragment.
+#[derive(Debug, Clone)]
 struct FragIndex {
     pub strings: HashMap<DexString, usize>,
     pub types: HashMap<IdType, usize>,
@@ -2131,3 +2314,38 @@ struct FragIndex {
     pub fields: HashMap<IdField, usize>,
     pub methods: HashMap<IdMethod, usize>,
 }
+
+/// The link state of a fragment. Describe the linking of a fragment and the related data.
+#[derive(Debug, Clone)]
+enum FragLinkState {
+    /// The fragment is not linked
+    Unlinked,
+    /// The index used in the fragment are linked of beeing linked.
+    /// The new local addresses of the element whose size depend on the value
+    /// of the index are stored here.
+    LinkedIdx {
+        code_item_relocation: HashMap<u32, u32>,
+        debug_info_item_relocation: HashMap<u32, u32>,
+        annotation_item_relocation: HashMap<u32, u32>,
+        encoded_array_relocation: HashMap<u32, u32>, // TODO only one?
+    },
+}
+
+impl FragLinkState {
+    fn start_linking_idx(&mut self) -> Result<()> {
+        match self {
+            Self::Unlinked => {
+                *self = Self::LinkedIdx {
+                    code_item_relocation: HashMap::<u32, u32>::new(),
+                    debug_info_item_relocation: HashMap::<u32, u32>::new(),
+                    annotation_item_relocation: HashMap::<u32, u32>::new(),
+                    encoded_array_relocation: HashMap::<u32, u32>::new(),
+                };
+                Ok(())
+            }
+            _ => {
+                bail!("Cannot link the idx in a fragment that is already linked");
+            }
+        }
+    }
+}