From 01c496aaacdcdd88008b75d7fb49dbef35bf1dfb Mon Sep 17 00:00:00 2001 From: Jean-Marie Mineau Date: Thu, 28 Mar 2024 17:16:56 +0100 Subject: [PATCH] WIP section manager utils --- androscalpel/src/dex_fragment.rs | 200 ++++++------ androscalpel/src/dex_writer.rs | 509 ++++++++++++++++++++++++++++++- 2 files changed, 598 insertions(+), 111 deletions(-) diff --git a/androscalpel/src/dex_fragment.rs b/androscalpel/src/dex_fragment.rs index fe70f28..4a0b0a9 100644 --- a/androscalpel/src/dex_fragment.rs +++ b/androscalpel/src/dex_fragment.rs @@ -1905,6 +1905,65 @@ impl DexFragment { Ok(()) } + /// Return `(nb_call_site, nb_method_handle)` the numbers of call sites and method handles in + /// the fragment. This is a non-trivial task that because some are generated during + /// [`Self::link_global_ids()`]. + pub fn get_nb_call_site_and_method_handle(&self, class: &Class) -> (usize, usize) { + fn get_nb_call_site_and_method_handle_in_method(method: &Method) -> (usize, usize) { + let mut nb_call_site = 0; + let mut nb_method_handle = 0; + if let Some(Code { insns, .. }) = method.code { + for ins in insns { + match ins { + Instruction::ConstMethodHandle(_) => nb_method_handle += 1, + Instruction::InvokeCustom(ins) => { + nb_call_site += 1; + let mut val_to_process: Vec<&DexValue> = + ins.call_site.args.iter().collect(); + while let Some(val) = val_to_process.pop() { + match val { + DexValue::MethodHandle(_) => nb_method_handle += 1, + DexValue::Array(DexArray(array)) => { + val_to_process.extend(array) + } + DexValue::Annotation(DexAnnotation { elements, .. }) => { + val_to_process.extend(elements.values()) + } + + _ => (), + } + } + } + _ => (), + } + } + }; + (nb_call_site, nb_method_handle) + } + + match self.link_state { + FragLinkState::Unlinked => { + let mut nb_call_site = self.call_site_ids.len(); + let mut nb_method_handle = self.method_handles.len(); + for method in class.direct_methods.values() { + let (nnb_call_site, nnb_method_handle) = + get_nb_call_site_and_method_handle_in_method(method); + nb_call_site += nnb_call_site; + nb_method_handle += nnb_method_handle + } + for method in class.virtual_methods.values() { + let (nnb_call_site, nnb_method_handle) = + get_nb_call_site_and_method_handle_in_method(method); + nb_call_site += nnb_call_site; + nb_method_handle += nnb_method_handle + } + (nb_call_site, nb_method_handle) + } + // after Self::link_global_ids(), the lists are complete + _ => (self.call_site_ids.len(), self.method_handles.len()), + } + } + /// 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. @@ -1916,10 +1975,9 @@ impl DexFragment { /// ## 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 + /// by [`Self::get_nb_call_site_and_method_handle()`] because some of those values are not + /// generated before linking id (because code items are generated when linking id, not good but + /// all we have right now). pub fn link_global_ids( &mut self, class: &Class, @@ -2070,6 +2128,9 @@ impl DexFragment { panic!("link_id_annotation() should only be called by link_global_ids()"); }; // TODO: can we update the new value directly in the direstoryies/set? + // The size of the elements change, so the section size must be recomputed + self.section_manager + .reset_section(FragSection::AnnotationItem); let mut initial_offset = 0; let mut new_offset = 0; for item in self.annotation_items { @@ -2091,6 +2152,8 @@ impl DexFragment { annotation_item_relocation.insert(initial_offset, new_offset); initial_offset += old_size; new_offset += new_size; + self.section_manager + .add_elt(FragSection::AnnotationItem, Some(new_size as usize)); } } @@ -2197,6 +2260,9 @@ impl DexFragment { } else { panic!("link_id_encoded_arrays() should only be called by link_global_ids()"); }; + // The size of the elements change, so the section size must be recomputed + self.section_manager + .reset_section(FragSection::EncodedArrayItem); let mut initial_offset = 0; let mut relocated_offset = 0; for array in self.encoded_array_items { @@ -2214,6 +2280,8 @@ impl DexFragment { encoded_array_relocation.insert(old_size, new_size); initial_offset += old_size; relocated_offset += new_size; + self.section_manager + .add_elt(FragSection::EncodedArrayItem, Some(new_size as usize)); } } @@ -2343,10 +2411,18 @@ impl DexFragment { } Ok(()) } + + pub fn get_section_manager(&self) -> &FragSectionManager { + &self.section_manager + } + + pub fn get_interfaces(&self) -> &Vec { + &self.interfaces + } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum FragSection { +pub(crate) enum FragSection { ClassDefItem, CallSiteIdItem, MethodHandleItem, @@ -2363,7 +2439,7 @@ enum FragSection { } impl FragSection { - const VARIANT_LIST: &'static [Self] = &[ + pub(crate) const VARIANT_LIST: &'static [Self] = &[ Self::ClassDefItem, Self::CallSiteIdItem, Self::MethodHandleItem, @@ -2454,27 +2530,15 @@ impl FragSection { } #[derive(Debug, Default, Clone)] -struct FragSectionManager { +pub(crate) struct FragSectionManager { sizes: [u32; Self::NB_SECTION], nb_elt: [usize; Self::NB_SECTION], - //offsets: [u32; Self::NB_SECTION], - editable: bool, } impl FragSectionManager { const NB_SECTION: usize = 12; - fn reset(&mut self) { - self.sizes = [0; Self::NB_SECTION]; - self.nb_elt = [0; Self::NB_SECTION]; - //self.offsets = [0; Self::NB_SECTION]; - self.editable = true; - } - fn add_elt(&mut self, section: FragSection, size: Option) { - if !self.editable { - panic!("Try to modify a section when the sections are set to read only"); - } if (section == FragSection::ClassDefItem || section == FragSection::ClassDataItem || section == FragSection::AnnotationsDirectoryItem) @@ -2488,35 +2552,18 @@ impl FragSectionManager { self.sizes[section.get_index()] += section.get_elt_size(size) as u32; self.nb_elt[section.get_index()] += 1; } - /* - fn incr_section_size(&mut self, section: FragSection, size: usize) { - if !self.editable { - panic!("Try to modify a section when the sections are set to read only"); - } - self.sizes[section.get_index()] += size as u32; + + fn reset_section(&mut self, section: FragSection) { + self.sizes[section.get_index()] = 0; + self.nb_elt[section.get_index()] = 0; } - fn get_offset(&self, section: FragSection) -> u32 { - if self.editable { - panic!("Try to get section offset before sections are finilized"); - } - let size = self.offsets[section.get_index()]; - let alignment = section.get_item_alignment(); - if size % alignment != 0 { - panic!( - "section {section:?} should be aligned on {alignment} bytes, \ - found section offset 0x{size:x}" - ); // avoid by finilized - } - size - } - */ fn get_unaligned_size(&self, section: FragSection) -> u32 { self.sizes[section.get_index()] } /// The position of a potential new item in the section considering alignment. - fn get_aligned_size(&self, section: FragSection) -> u32 { + pub(crate) fn get_aligned_size(&self, section: FragSection) -> u32 { let mut size = self.get_unaligned_size(section); while size % section.get_item_alignment() != 0 { size += 1; @@ -2524,73 +2571,22 @@ impl FragSectionManager { size } - fn get_nb_elt(&self, section: FragSection) -> usize { + pub(crate) fn get_nb_elt(&self, section: FragSection) -> usize { self.nb_elt[section.get_index()] } - /* - /// Finialize the sections: switch to read only and fix the section alignment. - fn finalize_sections(&mut self) { - for section in Section::VARIANT_LIST { - while self.sizes[..section.get_index()].iter().sum::() - % section.get_item_alignment() - != 0 - { - self.incr_section_size( - section.prev().expect( - "First section (Header) should alway be aligned but \ - found unaligned section without predecessor", - ), - 1, - ); - } + /// Generate a new section manager with the content of an added fragment. + pub(crate) fn merge_fragment(&self, frag_sections: &FragSectionManager) -> Self { + let mut new_section = self.clone(); + for section in FragSection::VARIANT_LIST { + new_section.nb_elt[section.get_index()] += frag_sections.get_nb_elt(*section); + // Noticed the "aligned" + new_section.sizes[section.get_index()] = + new_section.get_aligned_size(*section) + frag_sections.get_aligned_size(*section); } - let mut offset = 0; - for section in Section::VARIANT_LIST { - self.offsets[section.get_index()] = offset; - offset += self.sizes[section.get_index()]; - } - - self.editable = false; + new_section } - /// This method exist for the only purpose of linking the method code offset inside - /// the class data items. This linking needs to be done before finilizing because it change the - /// size of the class data item section. - /// - /// Seriously, avoid using this. - fn get_code_item_offset_prefinalized(&mut self) -> u32 { - if !self.editable || self.get_nb_elt(Section::MapList) != 0 { - panic!("Don't use this method for other purpose than linking class_data_items"); - } - let mut map_list_size = 4; - let map_item_size = 12; /* = MapItem { - type_: MapItemType::HeaderItem, - unused: 0, - size: 0, - offset: 0, - } - .size(); */ - for section in Section::VARIANT_LIST { - if !section.is_data() - && (self.get_nb_elt(*section) != 0 || section == &Section::MapList) - { - map_list_size += map_item_size; - } - } - let mut offset = map_list_size; // This is aligned so it wont affect alignment - for section in &Section::VARIANT_LIST[..Section::CodeItem.get_index()] { - // size Section::Data and size Section::MapList are 0 - while offset % section.get_item_alignment() != 0 { - offset += 1; - } - offset += self.sizes[section.get_index()]; - } - - offset - } - */ - /// Display the sections informations. #[allow(dead_code)] fn show(&self) { diff --git a/androscalpel/src/dex_writer.rs b/androscalpel/src/dex_writer.rs index 854bb73..862a8eb 100644 --- a/androscalpel/src/dex_writer.rs +++ b/androscalpel/src/dex_writer.rs @@ -1,8 +1,10 @@ //! The structure that generate a .dex from classes. use std::collections::{HashMap, HashSet, VecDeque}; -use crate::dex_fragment::DexFragment; +use crate::dex_fragment::{DexFragment, FragSection, FragSectionManager}; use crate::{Class, DexString, IdField, IdMethod, IdMethodType, IdType, Result}; +use androscalpel_serializer::{MapItem, MapItemType, MapList, Serializable, TypeItem, TypeList}; +use anyhow::{anyhow, bail}; #[derive(Debug, Clone)] pub struct DexWriter<'a> { @@ -39,8 +41,6 @@ impl<'a> DexWriter<'a> { let mut field_set: HashSet = HashSet::new(); let mut method_set: HashSet = HashSet::new(); let mut type_list_set: HashSet> = HashSet::new(); - let mut nb_method_handle = 0; - let mut nb_method_handle_before = vec![]; let fragments: VecDeque<(&'a Class, DexFragment)> = self .classes .into_iter() @@ -48,7 +48,7 @@ impl<'a> DexWriter<'a> { Ok(frag) => Ok((class, frag)), Err(err) => Err(err), }) - .collect()?; + .collect::>()?; loop { let (class, new_fragment) = if let Some(new_fragment) = fragments.pop_front() { new_fragment @@ -109,9 +109,9 @@ impl<'a> DexWriter<'a> { field_set.extend(new_fragment.field_ids().iter().cloned()); method_set.extend(new_fragment.method_ids().iter().cloned()); type_list_set.insert(new_fragment.interfaces().to_vec()); - nb_method_handle_before.push(nb_method_handle); - nb_method_handle += new_fragment.method_handles().len(); - fragments_in_file.push((class, new_fragment)); + let (nb_call_site, nb_method_handle) = + new_fragment.get_nb_call_site_and_method_handle(&class); + fragments_in_file.push((class, new_fragment, nb_call_site, nb_method_handle)); } type_list_set.extend(proto_set.iter().map(|proto| proto.parameters.clone())); let mut strings: Vec = string_set.iter().cloned().collect(); @@ -127,10 +127,135 @@ impl<'a> DexWriter<'a> { let mut type_lists: Vec> = type_list_set.iter().cloned().collect(); let index = DexIndex::new(&strings, &type_ids, &proto_ids, &field_ids, &method_ids); - for (i, (class, fragment)) in fragments_in_file.iter().enumerate() { - fragment.link_global_ids(&index, nb_method_handle_before[i]); + // TODO: Sort fragments_in_file in topological order + // fragments_in_file.sort() magic stuff + let mut nb_call_site_total = 0; + let mut nb_method_handle_total = 0; + for (_, _, nb_call_site, nb_method_handle) in &mut fragments_in_file { + let nb_call_site_before_fragment = nb_call_site_total; + let nb_method_handle_before_fragment = nb_method_handle_total; + nb_call_site_total += *nb_call_site; + nb_method_handle_total += *nb_method_handle; + *nb_call_site = nb_call_site_before_fragment; + *nb_method_handle = nb_method_handle_before_fragment; } + // Link descriptor ids (and generate code items) + let fragments_in_file: Vec = fragments_in_file + .into_iter() + .map( + |( + class, + fragment, + nb_call_site_before_fragment, + nb_method_handle_before_fragment, + )| { + match fragment.link_global_ids( + class, + &index, + nb_call_site_before_fragment, + nb_method_handle_before_fragment, + ) { + Err(err) => Err(err), + Ok(()) => Ok(fragment), + } + }, + ) + .collect::>()?; + + let mut type_lists = HashMap::, u32>::new(); + // collect section offsets + // The type are quite uggly but I'm behind on time + let mut fragment_section_manager = FragSectionManager::default(); + let fragments_in_file: Vec<(FragSectionManager, DexFragment)> = fragments_in_file + .into_iter() + .map(|fragment| { + let old_section_manager = fragment_section_manager; + let new_section_manager = + fragment_section_manager.merge_fragment(fragment.get_section_manager()); + fragment_section_manager = new_section_manager; + if !fragment.get_interfaces().is_empty() { + type_lists.insert(fragment.get_interfaces().clone(), 0); + } + (old_section_manager, fragment) + }) + .collect(); + + let mut section_manager: SectionManager = fragment_section_manager.into(); + + section_manager.add_elt(Section::HeaderItem, None); + section_manager.add_multiple_elt(Section::StringIdItem, strings.len()); + section_manager.add_multiple_elt(Section::TypeIdItem, type_ids.len()); + section_manager.add_multiple_elt(Section::ProtoIdItem, proto_ids.len()); + section_manager.add_multiple_elt(Section::FieldIdItem, field_ids.len()); + section_manager.add_multiple_elt(Section::MethodIdItem, method_ids.len()); + for string in &strings { + section_manager.add_elt(Section::StringDataItem, Some(string.0.size())); + } + for proto in &proto_ids { + if !proto.parameters.is_empty() { + type_lists.insert(proto.parameters.clone(), 0); + } + } + let type_list_list: Vec = type_lists + .keys() + .into_iter() + .map(|ty_list| { + if let Some(local_off) = type_lists.get_mut(ty_list) { + *local_off = section_manager.get_aligned_size(Section::TypeList) + } else { + bail!("type list {ty_list:?} not needed but not found in dex_writter") + } + let ty_list = TypeList { + list: ty_list + .into_iter() + .map(|ty| { + index + .types + .get(ty) + .ok_or(anyhow!( + "type {} needed but not found in dex writer", + ty.__str__() + )) + .map(|idx| TypeItem { + type_idx: *idx as u16, + }) + }) + .collect::>()?, + }; + + section_manager.add_elt(Section::TypeList, Some(ty_list.size())); + Ok(ty_list) + }) + .collect::>()?; + section_manager.add_elt(Section::MapList, Some(4)); + let offsets = section_manager.get_offsets(); + let map_item_size = 12; /* = MapItem { + type_: MapItemType::HeaderItem, + unused: 0, + size: 0, + offset: 0, + } + .size(); */ + let mut map_list = MapList::default(); + for section in Section::VARIANT_LIST { + if !section.is_data() && section_manager.get_nb_elt(*section) != 0 { + section_manager.incr_section_size(Section::MapList, map_item_size); + map_list.list.push(MapItem { + type_: section.get_map_item_type(), + unused: 0, + size: section_manager.get_nb_elt(*section) as u32, + offset: offsets.get_offset(*section), + }); + } + } + + // TODO: link en correct data section + + // TODO: link everything else + + // TODO: serialize + Ok(vec![]) } } @@ -197,3 +322,369 @@ impl<'a> DexIndex<'a> { } } } + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Section { + HeaderItem, + StringIdItem, + TypeIdItem, + ProtoIdItem, + FieldIdItem, + MethodIdItem, + ClassDefItem, + CallSiteIdItem, + MethodHandleItem, + Data, + MapList, + TypeList, + AnnotationSetRefList, + AnnotationSetItem, + CodeItem, + StringDataItem, + DebugInfoItem, + AnnotationItem, + EncodedArrayItem, + AnnotationsDirectoryItem, + HiddenapiClassDataItem, + ClassDataItem, +} + +impl From<&FragSection> for Section { + fn from(section: &FragSection) -> Self { + match section { + FragSection::ClassDefItem => Section::ClassDefItem, + FragSection::CallSiteIdItem => Section::CallSiteIdItem, + FragSection::MethodHandleItem => Section::MethodHandleItem, + FragSection::AnnotationSetRefList => Section::AnnotationSetRefList, + FragSection::AnnotationSetItem => Section::AnnotationSetItem, + FragSection::CodeItem => Section::CodeItem, + FragSection::StringDataItem => Section::StringDataItem, + FragSection::DebugInfoItem => Section::DebugInfoItem, + FragSection::AnnotationItem => Section::AnnotationItem, + FragSection::EncodedArrayItem => Section::EncodedArrayItem, + FragSection::AnnotationsDirectoryItem => Section::AnnotationsDirectoryItem, + FragSection::ClassDataItem => Section::ClassDataItem, + // FragSection::HiddenapiClassDataItem => Section::HiddenapiClassDataItem, + } + } +} + +impl Section { + const VARIANT_LIST: &'static [Self] = &[ + Self::HeaderItem, + Self::StringIdItem, + Self::TypeIdItem, + Self::ProtoIdItem, + Self::FieldIdItem, + Self::MethodIdItem, + Self::ClassDefItem, + Self::CallSiteIdItem, + Self::MethodHandleItem, + Self::Data, + Self::MapList, + Self::TypeList, + Self::AnnotationSetRefList, + Self::AnnotationSetItem, + Self::CodeItem, + Self::StringDataItem, + Self::DebugInfoItem, + Self::AnnotationItem, + Self::EncodedArrayItem, + Self::AnnotationsDirectoryItem, + Self::ClassDataItem, // must be last because contains offsets in Uleb, + // so size change when linking ! + Self::HiddenapiClassDataItem, + ]; + + fn get_index(&self) -> usize { + match self { + Self::HeaderItem => 0, + Self::StringIdItem => 1, + Self::TypeIdItem => 2, + Self::ProtoIdItem => 3, + Self::FieldIdItem => 4, + Self::MethodIdItem => 5, + Self::ClassDefItem => 6, + Self::CallSiteIdItem => 7, + Self::MethodHandleItem => 8, + Self::Data => 9, + Self::MapList => 10, + Self::TypeList => 11, + Self::AnnotationSetRefList => 12, + Self::AnnotationSetItem => 13, + Self::CodeItem => 14, + Self::StringDataItem => 15, + Self::DebugInfoItem => 16, + Self::AnnotationItem => 17, + Self::EncodedArrayItem => 18, + Self::AnnotationsDirectoryItem => 19, + Self::ClassDataItem => 20, + Self::HiddenapiClassDataItem => 21, + } + } + + fn get_default_elt_size(&self) -> Option { + match self { + Self::HeaderItem => Some(0x70), + Self::StringIdItem => Some(4), + Self::TypeIdItem => Some(4), + Self::ProtoIdItem => Some(0xc), + Self::FieldIdItem => Some(8), + Self::MethodIdItem => Some(8), + Self::ClassDefItem => Some(0x20), + Self::CallSiteIdItem => Some(4), + Self::MethodHandleItem => Some(8), + Self::Data => panic!("element cannot be inserted in data dirctly"), + Self::MapList => None, + Self::TypeList => None, + Self::AnnotationSetRefList => None, + Self::AnnotationSetItem => None, + Self::ClassDataItem => None, + Self::CodeItem => None, + Self::StringDataItem => None, + Self::DebugInfoItem => None, + Self::AnnotationItem => None, + Self::EncodedArrayItem => None, + Self::AnnotationsDirectoryItem => None, + Self::HiddenapiClassDataItem => None, + } + } + + fn get_elt_size(&self, default_size: Option) -> usize { + let fixed_size = self.get_default_elt_size(); + if let (Some(fixed_size), Some(default_size)) = (fixed_size, default_size) { + if fixed_size == default_size { + default_size + } else { + panic!( + "Element in {:?} have a size of {}, not {}", + self, fixed_size, default_size + ) + } + } else { + fixed_size.or(default_size).unwrap_or_else(|| { + panic!( + "Element of {:?} don't have a fixed size, you need to provide one", + self + ) + }) + } + } + + /// Return the previous section if it exist. + fn prev(&self) -> Option { + match self { + Self::HeaderItem => None, + Self::StringIdItem => Some(Self::HeaderItem), + Self::TypeIdItem => Some(Self::StringIdItem), + Self::ProtoIdItem => Some(Self::TypeIdItem), + Self::FieldIdItem => Some(Self::ProtoIdItem), + Self::MethodIdItem => Some(Self::FieldIdItem), + Self::ClassDefItem => Some(Self::MethodIdItem), + Self::CallSiteIdItem => Some(Self::ClassDefItem), + Self::MethodHandleItem => Some(Self::CallSiteIdItem), + Self::Data => Some(Self::MethodHandleItem), + Self::MapList => Some(Self::MethodHandleItem), // Data is just an indicator + Self::TypeList => Some(Self::MapList), + Self::AnnotationSetRefList => Some(Self::TypeList), + Self::AnnotationSetItem => Some(Self::AnnotationSetRefList), + Self::CodeItem => Some(Self::AnnotationSetItem), + Self::StringDataItem => Some(Self::CodeItem), + Self::DebugInfoItem => Some(Self::StringDataItem), + Self::AnnotationItem => Some(Self::DebugInfoItem), + Self::EncodedArrayItem => Some(Self::AnnotationItem), + Self::AnnotationsDirectoryItem => Some(Self::EncodedArrayItem), + Self::ClassDataItem => Some(Self::AnnotationsDirectoryItem), + Self::HiddenapiClassDataItem => Some(Self::ClassDataItem), + } + } + + fn get_map_item_type(&self) -> MapItemType { + match self { + Self::HeaderItem => MapItemType::HeaderItem, + Self::StringIdItem => MapItemType::StringIdItem, + Self::TypeIdItem => MapItemType::TypeIdItem, + Self::ProtoIdItem => MapItemType::ProtoIdItem, + Self::FieldIdItem => MapItemType::FieldIdItem, + Self::MethodIdItem => MapItemType::MethodIdItem, + Self::ClassDefItem => MapItemType::ClassDefItem, + Self::CallSiteIdItem => MapItemType::CallSiteIdItem, + Self::MethodHandleItem => MapItemType::MethodHandleItem, + Self::Data => panic!("Data is not a MatItemType"), + Self::MapList => MapItemType::MapList, + Self::TypeList => MapItemType::TypeList, + Self::AnnotationSetRefList => MapItemType::AnnotationSetRefList, + Self::AnnotationSetItem => MapItemType::AnnotationSetItem, + Self::CodeItem => MapItemType::CodeItem, + Self::StringDataItem => MapItemType::StringDataItem, + Self::DebugInfoItem => MapItemType::DebugInfoItem, + Self::AnnotationItem => MapItemType::AnnotationItem, + Self::EncodedArrayItem => MapItemType::EncodedArrayItem, + Self::AnnotationsDirectoryItem => MapItemType::AnnotationsDirectoryItem, + Self::ClassDataItem => MapItemType::ClassDataItem, + Self::HiddenapiClassDataItem => MapItemType::HiddenapiClassDataItem, + } + } + + /// Return the alignment of the item in byte. + fn get_item_alignment(&self) -> u32 { + match self { + Self::HeaderItem => 4, + Self::StringIdItem => 4, + Self::TypeIdItem => 4, + Self::ProtoIdItem => 4, + Self::FieldIdItem => 4, + Self::MethodIdItem => 4, + Self::ClassDefItem => 4, + Self::CallSiteIdItem => 1, + Self::MethodHandleItem => 4, + Self::Data => 1, + Self::MapList => 4, + Self::TypeList => 4, + Self::AnnotationSetRefList => 4, + Self::AnnotationSetItem => 4, + Self::CodeItem => 4, + Self::StringDataItem => 1, + Self::DebugInfoItem => 1, + Self::AnnotationItem => 1, + Self::EncodedArrayItem => 1, + Self::AnnotationsDirectoryItem => 4, + Self::ClassDataItem => 1, + Self::HiddenapiClassDataItem => 1, + } + } + + fn is_data(&self) -> bool { + matches!(self, Self::Data) + } +} + +#[derive(Debug, Default, Clone)] +struct SectionManager { + sizes: [u32; Self::NB_SECTION], + nb_elt: [usize; Self::NB_SECTION], +} + +impl From for SectionManager { + fn from(frag_sections: FragSectionManager) -> Self { + let mut section_manager = Self::default(); + for frag_section in FragSection::VARIANT_LIST { + let section: Section = frag_section.into(); + section_manager.nb_elt[section.get_index()] += frag_sections.get_nb_elt(*frag_section); + // Noticed the "aligned" + section_manager.sizes[section.get_index()] = section_manager.get_aligned_size(section) + + frag_sections.get_aligned_size(*frag_section); + } + section_manager + } +} + +impl SectionManager { + const NB_SECTION: usize = 22; + + fn add_elt(&mut self, section: Section, size: Option) { + if section.is_data() { + panic!("Cannot add element directly in section data"); + } + while self.sizes[section.get_index()] % section.get_item_alignment() != 0 { + self.sizes[section.get_index()] += 1; + } + self.sizes[section.get_index()] += section.get_elt_size(size) as u32; + self.nb_elt[section.get_index()] += 1; + } + + /// Add `nb` element in `section`. The element must be of fixed size (the layout of the file + /// ensure the element of fixed size are always aligned). + fn add_multiple_elt(&mut self, section: Section, nb: usize) { + if section.is_data() { + panic!("Cannot add element directly in section data"); + } + let size = if let Some(size) = section.get_default_elt_size() { + size + } else { + panic!( + "Cannot add multiple elements at the same time for element that don't \ + have a fixed size." + ); + }; + self.sizes[section.get_index()] += (size * nb) as u32; + self.nb_elt[section.get_index()] += nb; + } + + fn incr_section_size(&mut self, section: Section, size: usize) { + self.sizes[section.get_index()] += size as u32; + } + + fn get_unaligned_size(&self, section: Section) -> u32 { + if section.is_data() { + self.sizes[section.get_index()..].iter().sum() + } else { + self.sizes[section.get_index()] + } + } + + /// 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()] + } + + /// Display the sections informations. + #[allow(dead_code)] + fn show(&self) { + for section in Section::VARIANT_LIST { + let size = self.get_unaligned_size(*section); + let nb_elt = self.get_nb_elt(*section); + println!("{section:?}: size: 0x{size:x}, nb elt: {nb_elt})"); + } + } + + /// Return the offset of the sections with the sections defined in by this manager. + fn get_offsets(&self) -> SectionOffsets { + let mut offsets = SectionOffsets::default(); + let mut offset = 0; + for section in Section::VARIANT_LIST { + while offset % section.get_item_alignment() != 0 { + offset += 1; + } + offsets.offsets[section.get_index()]; + offset += self.sizes[section.get_index()]; + } + offsets + } +} + +#[derive(Debug, Default, Clone)] +struct SectionOffsets { + offsets: [u32; SectionManager::NB_SECTION], +} + +impl SectionOffsets { + fn get_offset(&self, section: Section) -> u32 { + self.offsets[section.get_index()] + } + + /// Return the offset of the sections **of a fragment**, from the **start of the dex file**. + /// `previous_fragment_sections` is are the merged `FragSectionManager` of the previous + /// fragment. + fn get_fragment_offset( + &self, + previous_fragment_sections: &FragSectionManager, + ) -> SectionOffsets { + let mut offsets = self.clone(); + for frag_section in FragSection::VARIANT_LIST { + let section: Section = frag_section.into(); + // Noticed the "aligned" + offsets.offsets[section.get_index()] += + previous_fragment_sections.get_aligned_size(*frag_section); + } + offsets + } +}