WIP link annotation offset

androscalpel/src/dex_fragment.rs

@@ -389,8 +389,7 @@ impl DexFragment {
             annotation_set.entries.push(AnnotationOffItem {
                 annotation_off: self
                     .section_manager
-                    .get_aligned_size(FragSection::AnnotationItem)
-                    + 1,
+                    .get_aligned_size(FragSection::AnnotationItem),
             });
 
             let item = AnnotationItem {
@@ -468,8 +467,7 @@ impl DexFragment {
             annotation_set.entries.push(AnnotationOffItem {
                 annotation_off: self
                     .section_manager
-                    .get_aligned_size(FragSection::AnnotationItem)
-                    + 1,
+                    .get_aligned_size(FragSection::AnnotationItem),
             });
 
             let item = AnnotationItem {
@@ -1843,8 +1841,7 @@ impl DexFragment {
             annotation_set.entries.push(AnnotationOffItem {
                 annotation_off: self
                     .section_manager
-                    .get_aligned_size(FragSection::AnnotationItem)
-                    + 1,
+                    .get_aligned_size(FragSection::AnnotationItem),
             });
 
             let item = AnnotationItem {
@@ -2433,15 +2430,36 @@ impl DexFragment {
         type_lists: &HashMap<Vec<IdType>, u32>,
     ) -> Result<()> {
         self.link_state = self.link_state.start_linking_offset()?;
-        self.link_off_call_site_ids(offsets);
+        self.link_off_call_site_ids(offsets)?;
         self.link_off_class_def(offsets, type_lists)?;
+        self.link_off_code(offsets);
+        self.link_off_annotations(offsets)?;
         Ok(())
     }
 
-    fn link_off_call_site_ids(&mut self, offsets: &SectionOffsets) {
+    fn link_off_call_site_ids(&mut self, offsets: &SectionOffsets) -> Result<()> {
+        let mut encoded_array_relocation = if let FragLinkState::LinkingOffset {
+            encoded_array_relocation,
+            ..
+        } = self.link_state
+        {
+            encoded_array_relocation
+        } else {
+            panic!(
+                "link_off_call_site_ids() should only be called by link_offsets_but_class_data()"
+            );
+        };
         for id in &mut self.call_site_ids {
-            id.call_site_off += offsets.get_offset(Section::EncodedArrayItem);
+            id.call_site_off = offsets.get_offset(Section::EncodedArrayItem)
+                + encoded_array_relocation
+                    .get(&id.call_site_off)
+                    .ok_or(anyhow!(
+                        "Could not found the new local offset of array \
+                    previously located at local offset 0x{:x}",
+                        id.call_site_off
+                    ))?;
         }
+        Ok(())
     }
 
     fn link_off_class_def(
@@ -2466,6 +2484,72 @@ impl DexFragment {
         // class data offset not set here
         Ok(())
     }
+
+    fn link_off_code(&mut self, offsets: &SectionOffsets) {
+        for code in &mut self.code_items {
+            if code.debug_info_off != 0 {
+                code.debug_info_off += offsets.get_offset(Section::DebugInfoItem) - 1;
+            }
+        }
+    }
+
+    fn link_off_annotations(&mut self, offsets: &SectionOffsets) -> Result<()> {
+        let mut annotation_item_relocation = if let FragLinkState::LinkingOffset {
+            annotation_item_relocation,
+            ..
+        } = self.link_state
+        {
+            annotation_item_relocation
+        } else {
+            panic!("link_off_annotations() should only be called by link_offsets_but_class_data()");
+        };
+        if let Some(annotation) = self.annotations_directory_item {
+            if annotation.class_annotations_off != 0 {
+                annotation.class_annotations_off +=
+                    offsets.get_offset(Section::AnnotationSetItem) - 1;
+            }
+            for field_annotation in annotation.field_annotations.iter_mut() {
+                if field_annotation.annotations_off != 0 {
+                    field_annotation.annotations_off +=
+                        offsets.get_offset(Section::AnnotationSetItem) - 1;
+                }
+            }
+            for method_annotation in annotation.method_annotations.iter_mut() {
+                if method_annotation.annotations_off != 0 {
+                    method_annotation.annotations_off +=
+                        offsets.get_offset(Section::AnnotationSetItem) - 1;
+                }
+            }
+            for parameter_annotation in annotation.parameter_annotations.iter_mut() {
+                if parameter_annotation.annotations_off != 0 {
+                    parameter_annotation.annotations_off +=
+                        offsets.get_offset(Section::AnnotationSetRefList) - 1;
+                }
+            }
+        }
+        for annotation_set in self.annotation_set_items.iter_mut() {
+            for entry in annotation_set.entries.iter_mut() {
+                entry.annotation_off = annotation_item_relocation
+                    .get(&entry.annotation_off)
+                    .ok_or(anyhow!(
+                        "Could not found new local offset for annotation with previous\
+                    local offset 0x{:x}",
+                        entry.annotation_off
+                    ))?
+                    + offsets.get_offset(Section::AnnotationItem);
+            }
+        }
+        for list in self.annotation_set_lists.iter_mut() {
+            for annotation in list.list.iter_mut() {
+                if annotation.annotations_off != 0 {
+                    annotation.annotations_off +=
+                        offsets.get_offset(Section::AnnotationSetItem) - 1;
+                }
+            }
+        }
+
+        Ok(())
+    }
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]

androscalpel/src/dex_writer_old.rs

@@ -2568,7 +2568,6 @@ impl DexWriter {
         for annotation_set in self.annotation_set_items.iter_mut() {
             for entry in annotation_set.entries.iter_mut() {
                 entry.annotation_off += self.section_manager.get_offset(Section::AnnotationItem);
-                // -1 now when linking frag
             }
         }
         for list in self.annotation_set_lists.iter_mut() {