diff --git a/androscalpel/src/dex_writer.rs b/androscalpel/src/dex_writer.rs index 7f5b97a..822be15 100644 --- a/androscalpel/src/dex_writer.rs +++ b/androscalpel/src/dex_writer.rs @@ -1306,7 +1306,9 @@ impl DexWriter { let debug_info_off = if code.debug_info.is_empty() { 0 } else { - let debug_info_off = self.section_manager.get_size(Section::DebugInfoItem); + let debug_info_off = self + .section_manager + .get_aligned_size(Section::DebugInfoItem); let item = DebugInfoItem::deserialize_from_slice(&code.debug_info)?; self.section_manager .add_elt(Section::DebugInfoItem, Some(item.size())); @@ -1409,7 +1411,7 @@ impl DexWriter { Uleb128(class.direct_methods.get(id).unwrap().get_raw_access_flags()); // No if let because ownership gunfooterie let code_off = if class.direct_methods.get(id).unwrap().code.is_some() { - let code_off = self.section_manager.get_size(Section::CodeItem); + let code_off = self.section_manager.get_aligned_size(Section::CodeItem); self.insert_code_item(id.clone(), true)?; Uleb128(code_off + 1) } else { @@ -1444,7 +1446,7 @@ impl DexWriter { ); // No if let because ownership gunfooterie let code_off = if class.virtual_methods.get(id).unwrap().code.is_some() { - let code_off = self.section_manager.get_size(Section::CodeItem); + let code_off = self.section_manager.get_aligned_size(Section::CodeItem); self.insert_code_item(id.clone(), false)?; Uleb128(code_off + 1) } else { @@ -1645,7 +1647,9 @@ impl DexWriter { values.push(DexValue::MethodType(call_site.type_.clone())); values.extend(call_site.args.iter().cloned()); self.call_site_ids.push(CallSiteIdItem { - call_site_off: self.section_manager.get_size(Section::EncodedArrayItem), + call_site_off: self + .section_manager + .get_aligned_size(Section::EncodedArrayItem), }); // linked in link_call_site_ids() self.section_manager.add_elt(Section::CallSiteIdItem, None); self.insert_encoded_array_item(DexArray(values)) @@ -1722,7 +1726,9 @@ impl DexWriter { annotations.sort_by_key(|annot| annot.annotation.type_.clone()); for annot in annotations { annotation_set.entries.push(AnnotationOffItem { - annotation_off: self.section_manager.get_size(Section::AnnotationItem), + annotation_off: self + .section_manager + .get_aligned_size(Section::AnnotationItem), }); // linked in link_annotations() let item = AnnotationItem { @@ -1769,7 +1775,9 @@ impl DexWriter { annotations.sort_by_key(|annot| annot.annotation.type_.clone()); for annot in annotations { annotation_set.entries.push(AnnotationOffItem { - annotation_off: self.section_manager.get_size(Section::AnnotationItem), + annotation_off: self + .section_manager + .get_aligned_size(Section::AnnotationItem), }); // linked in link_annotations() let item = AnnotationItem { @@ -1816,7 +1824,9 @@ impl DexWriter { annotations.sort_by_key(|annot| annot.annotation.type_.clone()); for annot in annotations { annotation_set.entries.push(AnnotationOffItem { - annotation_off: self.section_manager.get_size(Section::AnnotationItem), + annotation_off: self + .section_manager + .get_aligned_size(Section::AnnotationItem), }); // linked in link_annotations() let item = AnnotationItem { @@ -1864,7 +1874,9 @@ impl DexWriter { annotations.sort_by_key(|annot| annot.annotation.type_.clone()); for annot in annotations { annotation_set.entries.push(AnnotationOffItem { - annotation_off: self.section_manager.get_size(Section::AnnotationItem), + annotation_off: self + .section_manager + .get_aligned_size(Section::AnnotationItem), }); // linked in link_annotations() let item = AnnotationItem { @@ -1914,7 +1926,9 @@ impl DexWriter { for (param_idx, has_annotation) in param_has_annotation.into_iter().enumerate() { list.list.push(AnnotationSetRefItem { annotations_off: if has_annotation { - let annotation_off = self.section_manager.get_size(Section::AnnotationSetItem); + let annotation_off = self + .section_manager + .get_aligned_size(Section::AnnotationSetItem); self.insert_parameters_annotation_set(method_id, is_direct_method, param_idx)?; annotation_off + 1 } else { @@ -1942,7 +1956,9 @@ impl DexWriter { fn insert_annotations(&mut self, class_id: &IdType) -> Result<()> { let (class, _) = self.class_defs.get(class_id).unwrap(); let class_annotations_off = if !class.annotations.is_empty() { - let class_annotations_off = self.section_manager.get_size(Section::AnnotationSetItem); + let class_annotations_off = self + .section_manager + .get_aligned_size(Section::AnnotationSetItem); self.insert_class_annotation_set(class_id) .with_context(|| { format!( @@ -1973,7 +1989,10 @@ impl DexWriter { field_id.__repr__(), class_id.__repr__()), }; if !field.annotations.is_empty() { - let annotations_off = self.section_manager.get_size(Section::AnnotationSetItem) + 1; + let annotations_off = self + .section_manager + .get_aligned_size(Section::AnnotationSetItem) + + 1; self.insert_field_annotation_set(&field_id, is_static)?; field_annotations.push(FieldAnnotation { field_idx: *self.field_ids.get(&field_id).ok_or(anyhow!( @@ -2004,7 +2023,10 @@ impl DexWriter { method_id.__repr__(), class_id.__repr__()), }; if !method.annotations.is_empty() { - let annotations_off = self.section_manager.get_size(Section::AnnotationSetItem) + 1; + let annotations_off = self + .section_manager + .get_aligned_size(Section::AnnotationSetItem) + + 1; self.insert_method_annotation_set(method_id, is_direct)?; method_annotations.push(MethodAnnotation { method_idx: *self.method_ids.get(method_id).ok_or(anyhow!( @@ -2030,8 +2052,10 @@ impl DexWriter { method_id.__repr__(), class_id.__repr__()), }; if !method.parameters_annotations.is_empty() { - let annotations_off = - self.section_manager.get_size(Section::AnnotationSetRefList) + 1; + let annotations_off = self + .section_manager + .get_aligned_size(Section::AnnotationSetRefList) + + 1; self.insert_parameters_annotation_set_list(&method_id, is_direct)?; parameter_annotations.push(ParameterAnnotation { method_idx: *self.method_ids.get(&method_id).ok_or(anyhow!( @@ -2081,7 +2105,9 @@ impl DexWriter { .and_modify(|(_, i)| *i = idx); let (class, _) = self.class_defs.get(class_id).unwrap(); let class_data_off = if class.has_data_item() { - let class_data_off = self.section_manager.get_size(Section::ClassDataItem); + let class_data_off = self + .section_manager + .get_aligned_size(Section::ClassDataItem); self.insert_class_data_item(class_id)?; class_data_off + 1 } else { @@ -2091,7 +2117,9 @@ impl DexWriter { // mutating self with `insert_class_data_item`, and get a new ref afterward let (class, _) = self.class_defs.get(class_id).unwrap(); let static_values_off = if class.has_static_values_array() { - let static_values_off = self.section_manager.get_size(Section::EncodedArrayItem); + let static_values_off = self + .section_manager + .get_aligned_size(Section::EncodedArrayItem); self.insert_class_static_values(class_id)?; static_values_off + 1 } else { @@ -2101,7 +2129,7 @@ impl DexWriter { let annotations_off = if class.has_annotations() { let annotations_off = self .section_manager - .get_size(Section::AnnotationsDirectoryItem); + .get_aligned_size(Section::AnnotationsDirectoryItem); self.insert_annotations(class_id)?; annotations_off + 1 } else { @@ -2255,7 +2283,7 @@ impl DexWriter { self.header.method_ids_off = self.section_manager.get_offset(Section::MethodIdItem); self.header.class_defs_size = self.section_manager.get_nb_elt(Section::ClassDefItem) as u32; self.header.class_defs_off = self.section_manager.get_offset(Section::ClassDefItem); - self.header.data_size = self.section_manager.get_size(Section::Data); + self.header.data_size = self.section_manager.get_unaligned_size(Section::Data); self.header.data_off = self.section_manager.get_offset(Section::Data); } @@ -2622,6 +2650,8 @@ impl DexWriter { let mut hasher = Sha1::new(); io::copy(&mut buffer, &mut hasher)?; self.header.signature = hasher.finalize().into(); + let size = buffer.seek(SeekFrom::End(0))? as u32; + self.header.file_size = size; buffer.rewind()?; self.header.serialize(&mut buffer)?; @@ -3009,7 +3039,7 @@ impl SectionManager { size } - fn get_size(&self, section: Section) -> u32 { + fn get_unaligned_size(&self, section: Section) -> u32 { if section.is_data() { self.sizes[section.get_index()..].iter().sum() } else { @@ -3017,6 +3047,15 @@ impl SectionManager { } } + /// The position of a potential new item in the section considering alignment. + fn get_aligned_size(&self, section: Section) -> u32 { + let mut size = self.get_unaligned_size(section); + while size % section.get_item_alignment() != 0 { + size += 1; + } + size + } + fn get_nb_elt(&self, section: Section) -> usize { self.nb_elt[section.get_index()] } @@ -3045,7 +3084,7 @@ impl SectionManager { fn show(&self) { let mut offset = 0; for section in Section::VARIANT_LIST { - let size = self.get_size(*section); + let size = self.get_unaligned_size(*section); let new_offset = offset + size; let nb_elt = self.get_nb_elt(*section); println!(