finish-ish id linking

TODO.md

@@ -3,18 +3,6 @@
 - https://source.android.com/docs/core/runtime/dex-format#system-annotation
 - goto size computation
 - no nop when no payload
+- code items need a "fragment" repr
 - option to get label at every code addresses
 - name register / parameters
-
-# frag meth
-- insert_class_static_values
-
-# relink
-
-data.static_fields: field_idx_diff
-data.direct_methods: method_idx_diff
-data.code_off
-code_item.try_item
-instr
-
-EncodedValue

androscalpel/src/dex_fragment.rs

@@ -66,6 +66,14 @@ pub struct DexFragment {
     interfaces: Vec<IdType>,
     /// The current link state of the fragment.
     link_state: FragLinkState,
+
+    /// This should not be stored, but currently code items
+    /// need to be generated when linking ids, and to avoid having
+    /// fragmented data representations, the data used by the code
+    /// must be generated "pre-linking". Idealy, code items should be
+    /// precompiled in an intermediate state when creating the fragment,
+    /// and the `FragIndex` would not be needed anymore.
+    frag_index: Option<FragIndex>,
 }
 
 impl DexFragment {
@@ -135,6 +143,7 @@ impl DexFragment {
             debug_info_items: vec![],
             interfaces: class.interfaces.clone(),
             link_state: FragLinkState::Unlinked,
+            frag_index: None,
         };
         frag.strings = class.get_all_strings().into_iter().collect();
         frag.strings.sort();
@@ -234,6 +243,7 @@ impl DexFragment {
         } else {
             0
         };
+        frag.frag_index = Some(index);
 
         Ok(frag)
     }
@@ -579,7 +589,14 @@ impl DexFragment {
     /// This item cannot be cached, because the jump instructions depend on the size of
     /// instructions that depend on the size of the descriptor ids that depend on the
     /// list of all descriptors in the dex file.
-    fn insert_code_item(&mut self, code: &Code, index: &DexIndex) -> Result<()> {
+    fn insert_code_item(
+        &mut self,
+        code: &Code,
+        index: &DexIndex,
+        nb_call_site_before_fragment: usize,
+        nb_method_handle_before_fragment: usize,
+        frag_index: &FragIndex,
+    ) -> Result<()> {
         // Estimate instructions addresses
         let mut min_addr = 0;
         let mut max_addr = 0;
@@ -1350,15 +1367,16 @@ impl DexFragment {
                     insns.push(ins);
                 }
                 Instruction::InvokeCustom(ins) => {
-                    let call_site_idx = self.call_site_ids.len();
+                    let call_site_idx = self.call_site_ids.len() + nb_call_site_before_fragment;
-                    self.insert_call_site_item(&ins.call_site, index)?;
+                    self.insert_call_site_item(&ins.call_site, frag_index)?;
                     let ins = ins.get_raw_ins(call_site_idx);
                     addr += ins.size() / 2;
                     insns.push(ins);
                 }
                 Instruction::ConstMethodHandle(ins) => {
-                    let method_handle_idx = self.method_handles.len();
+                    let method_handle_idx =
-                    self.insert_method_handle(&ins.handle, index)?;
+                        self.method_handles.len() + nb_method_handle_before_fragment;
+                    self.insert_method_handle(&ins.handle, frag_index)?;
                     let ins = ins.get_raw_ins(method_handle_idx);
                     addr += ins.size() / 2;
                     insns.push(ins);
@@ -1890,12 +1908,31 @@ impl DexFragment {
     /// 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.
+    ///
+    /// `call_site_id_item`s are sorted by offset of the associated array, so they are naturally
+    /// sorted, and `method_handle_item`s are not sorted. This means linking those id is done
+    /// simply by adding the number of ids before the fragement.
+    ///
+    /// ## Warning
+    ///
+    /// The `nb_call_site_before_fragment` and `nb_method_handle_before_fragment` must be computed
+    /// **after** linking the id of the previous fragment as the code item is generated at this
+    /// step, doing so generate `call_site_id_item`s and `method_handle_item`s.
+    ///
+    /// This is bad as it prevent parallelistation of this step. TODO
     pub fn link_global_ids(
         &mut self,
         class: &Class,
         index: &DexIndex,
+        nb_call_site_before_fragment: usize,
         nb_method_handle_before_fragment: usize,
     ) -> Result<()> {
+        let frag_index = if let Some(frag_index) = self.frag_index {
+            self.frag_index = None;
+            frag_index
+        } else {
+            bail!("Fragment cannot be linked without the frag index (see DexFragment::frag_index doc)")
+        };
         self.link_state.start_linking_idx()?;
         let string_reindex = Vec::with_capacity(self.strings.len());
         // TODO: considering we have the map, this can be simplified a lot
@@ -1960,7 +1997,15 @@ impl DexFragment {
             method_reindex.push(global_idx as u16);
         }
 
-        self.link_id_class_data_and_gen_code(class, &field_reindex, &method_reindex, index)?;
+        self.link_id_class_data_and_gen_code(
+            class,
+            &field_reindex,
+            &method_reindex,
+            index,
+            nb_call_site_before_fragment,
+            nb_method_handle_before_fragment,
+            &frag_index,
+        )?;
         self.link_id_class_def(&string_reindex, &type_reindex);
         self.link_id_method_handle(&field_reindex, &method_reindex);
         self.link_id_annotation(
@@ -1979,9 +2024,8 @@ impl DexFragment {
             &proto_reindex,
             nb_method_handle_before_fragment,
         );
-        self.link_id_annotation_dir(field_reindex, method_reindex)
+        self.link_id_annotation_dir(&field_reindex, &method_reindex);
-
+        Ok(())
-        todo!()
     }
 
     fn link_id_class_def(&mut self, string_reindex: &[u32], type_reindex: &[u32]) {
@@ -2194,6 +2238,9 @@ impl DexFragment {
         field_reindex: &[u16],
         method_reindex: &[u16],
         index: &DexIndex,
+        nb_call_site_before_fragment: usize,
+        nb_method_handle_before_fragment: usize,
+        frag_index: &FragIndex,
     ) -> Result<()> {
         if let Some(data) = self.class_data {
             let mut last_local_id = 0;
@@ -2234,7 +2281,13 @@ impl DexFragment {
                         .code;
                     if let Some(code) = code {
                         let code_off = self.section_manager.get_aligned_size(FragSection::CodeItem);
-                        self.insert_code_item(&code, index)?;
+                        self.insert_code_item(
+                            &code,
+                            index,
+                            nb_call_site_before_fragment,
+                            nb_method_handle_before_fragment,
+                            frag_index,
+                        )?;
                         meth.code_off.0 = code_off + 1;
                     } else {
                         bail!(
@@ -2267,7 +2320,13 @@ impl DexFragment {
                         .code;
                     if let Some(code) = code {
                         let code_off = self.section_manager.get_aligned_size(FragSection::CodeItem);
-                        self.insert_code_item(&code, index)?;
+                        self.insert_code_item(
+                            &code,
+                            index,
+                            nb_call_site_before_fragment,
+                            nb_method_handle_before_fragment,
+                            frag_index,
+                        )?;
                         meth.code_off.0 = code_off + 1;
                     } else {
                         bail!(