diff --git a/androscalpel/src/dex_writer.rs b/androscalpel/src/dex_writer.rs index 287995b..fc81392 100644 --- a/androscalpel/src/dex_writer.rs +++ b/androscalpel/src/dex_writer.rs @@ -45,6 +45,33 @@ pub struct DexWriter { // values // annotations // + /// A struct that keep track of sections size, nb elements and offsets. + section_manager: SectionManager, + /// The string_data section. Once populated, the elements must be sorted according to the spec. + string_data_list: Vec, + /// The type_ids sections. Once populated, the elements must be sorted according to the spec. + type_ids_list: Vec, + /// The proto_ids section. Once populated, the elements must be sorted according to the spec. + /// the `parameter_off` field of the prototypes must be linked before serializing the section. + proto_ids_list: Vec, + /// The field_ids section. Once populated, the elements must be sorted according to the spec. + field_ids_list: Vec, + /// The method_ids section. Once populated, the elements must be sorted according to the spec. + method_ids_list: Vec, + /// The class_defs section. Once populated, the elements must be sorted according to the spec. + /// The `interfaces_off`, `annotations_off`, `class_data_off` and `static_values_off` fields + /// must be linked before serializing. Before liking, the `annotations_off`, `class_data_off` + /// and `static_values_off` field take the value of the offset from the start of theire + /// respective section (or zero if the class does not have a value associated to the field). + class_defs_list: Vec, + /// The class_data section. + class_data_list: Vec, + /// The type lists found in the classes associated to their index in the type_lists section. + type_lists_index: HashMap, + /// The type_lists section and the offset of the lists inside the section. + type_lists_with_offset: Vec<(TypeList, u32)>, + /// The map_list + map_list: MapList, } impl Default for DexWriter { @@ -83,6 +110,17 @@ impl Default for DexWriter { field_ids: HashMap::new(), method_ids: HashMap::new(), class_defs: HashMap::new(), + section_manager: SectionManager::default(), + string_data_list: vec![], + type_ids_list: vec![], + proto_ids_list: vec![], + field_ids_list: vec![], + method_ids_list: vec![], + class_defs_list: vec![], + class_data_list: vec![], + type_lists_index: HashMap::new(), + type_lists_with_offset: vec![], + map_list: MapList::default(), } } } @@ -144,16 +182,7 @@ impl DexWriter { Ok(()) } - pub fn gen_dex_file_to_vec(&mut self) -> Result> { - let mut output = Cursor::new(Vec::::new()); - self.write_dex_file(&mut output)?; - Ok(output.into_inner()) - } - - fn write_dex_file(&mut self, writer: &mut dyn Write) -> Result<()> { - let mut section_manager = SectionManager::default(); - section_manager.incr_section_size(Section::HeaderItem, 0x70); - + fn gen_string_data_section(&mut self) -> Result<()> { debug!("Sort string and generate string_data_item and string_ids sections"); let mut string_ids_list: Vec = self.strings.keys().cloned().collect(); string_ids_list.sort(); @@ -161,35 +190,40 @@ impl DexWriter { self.strings .entry(string.clone()) .and_modify(|val| *val = idx); - section_manager.add_elt(Section::StringIdItem, None); - section_manager.add_elt(Section::StringDataItem, Some(string.0.size())); + self.section_manager.add_elt(Section::StringIdItem, None); + self.section_manager + .add_elt(Section::StringDataItem, Some(string.0.size())); } - let string_ids_list: Vec = string_ids_list + self.string_data_list = string_ids_list .into_iter() .map(|string| string.into()) .collect(); + Ok(()) + } + fn gen_type_ids_section(&mut self) -> Result<()> { debug!("Sort types and generate type_id_item section"); let mut type_ids_list: Vec = self.type_ids.keys().cloned().collect(); type_ids_list.sort(); for (idx, ty) in type_ids_list.iter().enumerate() { self.type_ids.entry(ty.clone()).and_modify(|val| *val = idx); - section_manager.add_elt(Section::TypeIdItem, None); + self.section_manager.add_elt(Section::TypeIdItem, None); } - let type_ids_list: Vec = { - let mut type_ids_list_aux = vec![]; - for ty in type_ids_list.into_iter() { - type_ids_list_aux.push(TypeIdItem { - descriptor_idx: *self.strings.get(&ty.0).ok_or(anyhow!( - "String {} (name of type {}) not found in dex builder", - ty.0.__repr__(), - ty.__repr__() - ))? as u32, - }); - } - type_ids_list_aux - }; + let mut type_ids_list_aux = vec![]; + for ty in type_ids_list.into_iter() { + type_ids_list_aux.push(TypeIdItem { + descriptor_idx: *self.strings.get(&ty.0).ok_or(anyhow!( + "String {} (name of type {}) not found in dex builder", + ty.0.__repr__(), + ty.__repr__() + ))? as u32, + }); + } + self.type_ids_list = type_ids_list_aux; + Ok(()) + } + fn gen_proto_ids_section(&mut self) -> Result<()> { debug!("Sort prototypes and generate proto_id_item section"); let mut proto_ids_list: Vec = self.proto_ids.keys().cloned().collect(); proto_ids_list.sort(); @@ -197,28 +231,29 @@ impl DexWriter { self.proto_ids .entry(proto.clone()) .and_modify(|val| *val = idx); - section_manager.add_elt(Section::ProtoIdItem, None); + self.section_manager.add_elt(Section::ProtoIdItem, None); } - let mut proto_ids_list = { - let mut proto_ids_list_aux = vec![]; - for proto in proto_ids_list { - proto_ids_list_aux.push(ProtoIdItem { - shorty_idx: *self.strings.get(&proto.shorty).ok_or(anyhow!( - "String {}, (shorty of prototype {}) not found in dex builder", - proto.shorty.__repr__(), - proto.__repr__() - ))? as u32, - return_type_idx: *self.type_ids.get(&proto.return_type).ok_or(anyhow!( - "Type {}, (return type of prototype {}) not found in dex builder", - proto.shorty.__repr__(), - proto.__repr__() - ))? as u32, - parameters_off: 0, // TO BE LINKED LATTER - }); - } - proto_ids_list_aux - }; + let mut proto_ids_list_aux = vec![]; + for proto in proto_ids_list { + proto_ids_list_aux.push(ProtoIdItem { + shorty_idx: *self.strings.get(&proto.shorty).ok_or(anyhow!( + "String {}, (shorty of prototype {}) not found in dex builder", + proto.shorty.__repr__(), + proto.__repr__() + ))? as u32, + return_type_idx: *self.type_ids.get(&proto.return_type).ok_or(anyhow!( + "Type {}, (return type of prototype {}) not found in dex builder", + proto.shorty.__repr__(), + proto.__repr__() + ))? as u32, + parameters_off: 0, // Will be linked later + }); + } + self.proto_ids_list = proto_ids_list_aux; + Ok(()) + } + fn gen_field_ids_section(&mut self) -> Result<()> { debug!("Sort field ids and generate field_ids_item"); let mut field_ids_list: Vec = self.field_ids.keys().cloned().collect(); field_ids_list.sort(); @@ -226,32 +261,34 @@ impl DexWriter { self.field_ids .entry(field_id.clone()) .and_modify(|val| *val = idx); - section_manager.add_elt(Section::FieldIdItem, None); + self.section_manager.add_elt(Section::FieldIdItem, None); } - let field_ids_list: Vec = { - let mut field_ids_list_aux = vec![]; - for field in field_ids_list.into_iter() { - field_ids_list_aux.push(FieldIdItem { - class_idx: *self.type_ids.get(&field.class_).ok_or(anyhow!( - "Type {} (class of field {}) not found in dex builder", - field.class_.__repr__(), - field.__repr__() - ))? as u16, - type_idx: *self.type_ids.get(&field.type_).ok_or(anyhow!( - "Type {} (type of field {}) not found in dex builder", - field.type_.__repr__(), - field.__repr__() - ))? as u16, - name_idx: *self.strings.get(&field.name).ok_or(anyhow!( - "String {} (name of field {}) not found in dex builder", - field.name.__repr__(), - field.__repr__() - ))? as u32, - }); - } - field_ids_list_aux - }; + let mut field_ids_list_aux = vec![]; + for field in field_ids_list.into_iter() { + field_ids_list_aux.push(FieldIdItem { + class_idx: *self.type_ids.get(&field.class_).ok_or(anyhow!( + "Type {} (class of field {}) not found in dex builder", + field.class_.__repr__(), + field.__repr__() + ))? as u16, + type_idx: *self.type_ids.get(&field.type_).ok_or(anyhow!( + "Type {} (type of field {}) not found in dex builder", + field.type_.__repr__(), + field.__repr__() + ))? as u16, + name_idx: *self.strings.get(&field.name).ok_or(anyhow!( + "String {} (name of field {}) not found in dex builder", + field.name.__repr__(), + field.__repr__() + ))? as u32, + }); + } + self.field_ids_list = field_ids_list_aux; + Ok(()) + } + + fn gen_method_ids_section(&mut self) -> Result<()> { debug!("Sort method ids and generate method_id_item section"); let mut method_ids_list: Vec = self.method_ids.keys().cloned().collect(); method_ids_list.sort(); @@ -259,190 +296,197 @@ impl DexWriter { self.method_ids .entry(method_id.clone()) .and_modify(|val| *val = idx); - section_manager.add_elt(Section::MethodIdItem, None); + self.section_manager.add_elt(Section::MethodIdItem, None); } - let method_ids_list: Vec = { - let mut method_ids_list_aux = vec![]; - for method in method_ids_list.into_iter() { - method_ids_list_aux.push(MethodIdItem { - class_idx: *self.type_ids.get(&method.class_).ok_or(anyhow!( - "Type {} (class of method {}) not found in dex builder", - method.class_.__repr__(), - method.__repr__() - ))? as u16, - proto_idx: *self.proto_ids.get(&method.proto).ok_or(anyhow!( - "Prototype {} (signature of method {}) not found in dex builder", - method.proto.__repr__(), - method.__repr__() - ))? as u16, - name_idx: *self.strings.get(&method.name).ok_or(anyhow!( - "String {} (name of method {}) not found in dex builder", - method.name.__repr__(), - method.__repr__() - ))? as u32, - }); - } - method_ids_list_aux - }; - - debug!("Sort classes and generate the class_defs and class_data section"); - let mut class_defs_list = vec![]; - let mut class_data_list = vec![]; - for (idx, class_id) in self.get_sorted_class_def()?.into_iter().enumerate() { - self.class_defs - .entry(class_id.clone()) - .and_modify(|(_, i)| *i = idx); - let (class, _) = self.class_defs.get(&class_id).unwrap(); - class_defs_list.push(ClassDefItem { - class_idx: *self.type_ids.get(&class.descriptor).ok_or(anyhow!( - "Type {} (type of class {}) not found in dex builder", - class.descriptor.__repr__(), - class.__repr__() + let mut method_ids_list_aux = vec![]; + for method in method_ids_list.into_iter() { + method_ids_list_aux.push(MethodIdItem { + class_idx: *self.type_ids.get(&method.class_).ok_or(anyhow!( + "Type {} (class of method {}) not found in dex builder", + method.class_.__repr__(), + method.__repr__() + ))? as u16, + proto_idx: *self.proto_ids.get(&method.proto).ok_or(anyhow!( + "Prototype {} (signature of method {}) not found in dex builder", + method.proto.__repr__(), + method.__repr__() + ))? as u16, + name_idx: *self.strings.get(&method.name).ok_or(anyhow!( + "String {} (name of method {}) not found in dex builder", + method.name.__repr__(), + method.__repr__() ))? as u32, - access_flags: class.get_raw_access_flags(), - superclass_idx: if let Some(sup) = &class.superclass { - *self.type_ids.get(sup).ok_or(anyhow!( - "Type {} (superclass of class {}) not found in dex builder", - sup.__repr__(), - class.__repr__() - ))? as u32 - } else { - NO_INDEX.0 - }, - interfaces_off: 0, - source_file_idx: if let Some(file) = &class.source_file { - *self.strings.get(file).ok_or(anyhow!( - "String {} (source file of class {}) not found in dex builder", - file.__repr__(), - class.__repr__() - ))? as u32 - } else { - NO_INDEX.0 - }, - annotations_off: 0, // TODO - class_data_off: { - // need relinking once offset of class_data section is known - if class.has_data_section() { - let mut data = ClassDataItem::default(); - - let mut static_fields: Vec = - class.static_fields.keys().cloned().collect(); - static_fields.sort(); - let mut last_field_id = 0; - for id in &static_fields { - let idx = self.field_ids.get(id).ok_or(anyhow!( - "Field {} (field of class {}) not found in dex builder", - id.__repr__(), - class.__repr__() - ))?; - let field_idx_diff = Uleb128((idx - last_field_id) as u32); - last_field_id = *idx; - let access_flags = Uleb128( - class.static_fields.get(id).unwrap().get_raw_access_flags(), - ); - data.static_fields.push(EncodedField { - field_idx_diff, - access_flags, - }); - } - - let mut instance_fields: Vec = - class.instance_fields.keys().cloned().collect(); - instance_fields.sort(); - let mut last_field_id = 0; - for id in &instance_fields { - let idx = self.field_ids.get(id).ok_or(anyhow!( - "Field {} (field of class {}) not found in dex builder", - id.__repr__(), - class.__repr__() - ))?; - let field_idx_diff = Uleb128((idx - last_field_id) as u32); - last_field_id = *idx; - let access_flags = Uleb128( - class - .instance_fields - .get(id) - .unwrap() - .get_raw_access_flags(), - ); - data.instance_fields.push(EncodedField { - field_idx_diff, - access_flags, - }); - } - - let mut direct_methods: Vec = - class.direct_methods.keys().cloned().collect(); - direct_methods.sort(); - let mut last_method_id = 0; - for id in &direct_methods { - let idx = self.method_ids.get(id).ok_or(anyhow!( - "Method {} (method of class {}) not found in dex builder", - id.__repr__(), - class.__repr__() - ))?; - let method_idx_diff = Uleb128((idx - last_method_id) as u32); - last_method_id = *idx; - let access_flags = Uleb128( - class.direct_methods.get(id).unwrap().get_raw_access_flags(), - ); - data.direct_methods.push(EncodedMethod { - method_idx_diff, - access_flags, - code_off: Uleb128(0), // TODO - }); - } - - let mut virtual_methods: Vec = - class.virtual_methods.keys().cloned().collect(); - virtual_methods.sort(); - let mut last_method_id = 0; - for id in &virtual_methods { - let idx = self.method_ids.get(id).ok_or(anyhow!( - "Method {} (method of class {}) not found in dex builder", - id.__repr__(), - class.__repr__() - ))?; - let method_idx_diff = Uleb128((idx - last_method_id) as u32); - last_method_id = *idx; - let access_flags = Uleb128( - class - .virtual_methods - .get(id) - .unwrap() - .get_raw_access_flags(), - ); - data.virtual_methods.push(EncodedMethod { - method_idx_diff, - access_flags, - code_off: Uleb128(0), // TODO - }); - } - - let class_data_off = section_manager.get_size(Section::ClassDataItem); - section_manager.add_elt(Section::ClassDataItem, Some(data.size())); - class_data_list.push(data); - class_data_off - } else { - 0 - } - }, - static_values_off: 0, // TODO }); - section_manager.add_elt(Section::ClassDefItem, None); + } + self.method_ids_list = method_ids_list_aux; + Ok(()) + } + + /// Insert a class_data_item in the class_data section (in data). + fn insert_class_data_item(&mut self, class_id: &IdType) -> Result<()> { + let mut data = ClassDataItem::default(); + let (class, _) = self.class_defs.get(class_id).unwrap(); + + let mut static_fields: Vec = class.static_fields.keys().cloned().collect(); + static_fields.sort(); + let mut last_field_id = 0; + for id in &static_fields { + let idx = self.field_ids.get(id).ok_or(anyhow!( + "Field {} (field of class {}) not found in dex builder", + id.__repr__(), + class.__repr__() + ))?; + let field_idx_diff = Uleb128((idx - last_field_id) as u32); + last_field_id = *idx; + let access_flags = Uleb128(class.static_fields.get(id).unwrap().get_raw_access_flags()); + data.static_fields.push(EncodedField { + field_idx_diff, + access_flags, + }); } - // Method handles are not ordered, nor deduplicated, so they are generated on the fly - let mut _method_handles: Vec = vec![]; + let mut instance_fields: Vec = class.instance_fields.keys().cloned().collect(); + instance_fields.sort(); + let mut last_field_id = 0; + for id in &instance_fields { + let idx = self.field_ids.get(id).ok_or(anyhow!( + "Field {} (field of class {}) not found in dex builder", + id.__repr__(), + class.__repr__() + ))?; + let field_idx_diff = Uleb128((idx - last_field_id) as u32); + last_field_id = *idx; + let access_flags = Uleb128( + class + .instance_fields + .get(id) + .unwrap() + .get_raw_access_flags(), + ); + data.instance_fields.push(EncodedField { + field_idx_diff, + access_flags, + }); + } + let mut direct_methods: Vec = class.direct_methods.keys().cloned().collect(); + direct_methods.sort(); + let mut last_method_id = 0; + for id in &direct_methods { + let idx = self.method_ids.get(id).ok_or(anyhow!( + "Method {} (method of class {}) not found in dex builder", + id.__repr__(), + class.__repr__() + ))?; + let method_idx_diff = Uleb128((idx - last_method_id) as u32); + last_method_id = *idx; + let access_flags = + Uleb128(class.direct_methods.get(id).unwrap().get_raw_access_flags()); + data.direct_methods.push(EncodedMethod { + method_idx_diff, + access_flags, + code_off: Uleb128(0), // TODO + }); + } + + let mut virtual_methods: Vec = class.virtual_methods.keys().cloned().collect(); + virtual_methods.sort(); + let mut last_method_id = 0; + for id in &virtual_methods { + let idx = self.method_ids.get(id).ok_or(anyhow!( + "Method {} (method of class {}) not found in dex builder", + id.__repr__(), + class.__repr__() + ))?; + let method_idx_diff = Uleb128((idx - last_method_id) as u32); + last_method_id = *idx; + let access_flags = Uleb128( + class + .virtual_methods + .get(id) + .unwrap() + .get_raw_access_flags(), + ); + data.virtual_methods.push(EncodedMethod { + method_idx_diff, + access_flags, + code_off: Uleb128(0), // TODO + }); + } + self.section_manager + .add_elt(Section::ClassDataItem, Some(data.size())); + self.class_data_list.push(data); + Ok(()) + } + + /// Insert a class_def_item in the class_defs section **and** the other struct that needs to be + /// generated on the fly. + /// + /// # Warning + /// + /// The class_defs section **MUST** be sorted by inheritance dependencies (parents classes and + /// interfaces must appear **before** child classes). Accordingly, this method must be invoked + /// in the right order. + fn insert_class_def_item(&mut self, class_id: &IdType) -> Result<()> { + let idx = self.class_defs_list.len(); + self.class_defs + .entry(class_id.clone()) + .and_modify(|(_, i)| *i = idx); + let (class, _) = self.class_defs.get(class_id).unwrap(); + let class_data_off = if class.has_data_section() { + let class_data_off = self.section_manager.get_size(Section::ClassDataItem); + self.insert_class_data_item(class_id)?; + class_data_off + } else { + 0 + }; + // & vs &mut cluster-f, this make rust drop the ref so self hold by `class` before + // mutating self with `insert_class_data_item`, and get a new ref afterward + let (class, _) = self.class_defs.get(class_id).unwrap(); + self.class_defs_list.push(ClassDefItem { + class_idx: *self.type_ids.get(class_id).ok_or(anyhow!( + "Type {} (type of class {}) not found in dex builder", + class_id.__repr__(), + class.__repr__() + ))? as u32, + access_flags: class.get_raw_access_flags(), + superclass_idx: if let Some(sup) = &class.superclass { + *self.type_ids.get(sup).ok_or(anyhow!( + "Type {} (superclass of class {}) not found in dex builder", + sup.__repr__(), + class.__repr__() + ))? as u32 + } else { + NO_INDEX.0 + }, + interfaces_off: 0, + source_file_idx: if let Some(file) = &class.source_file { + *self.strings.get(file).ok_or(anyhow!( + "String {} (source file of class {}) not found in dex builder", + file.__repr__(), + class.__repr__() + ))? as u32 + } else { + NO_INDEX.0 + }, + + annotations_off: 0, // TODO + class_data_off, // need relinking once offset of class_data section is known + static_values_off: 0, // TODO + }); + self.section_manager.add_elt(Section::ClassDefItem, None); + Ok(()) + } + + fn gen_type_list_section(&mut self) -> Result<()> { debug!("Generate the type_list section"); - let mut type_lists_index = HashMap::new(); + // Collect all type lists for proto in self.proto_ids.keys() { if !proto.parameters.is_empty() { let type_list = self.gen_type_list(&proto.parameters).with_context(|| { format!("Failed to generate param list for {}", proto.__repr__()) })?; - type_lists_index.insert(type_list, 0); + self.type_lists_index.insert(type_list, 0); } } for (class, _) in self.class_defs.values() { @@ -450,42 +494,54 @@ impl DexWriter { let type_list = self.gen_type_list(&class.interfaces).with_context(|| { format!("Failed to generate interface list for {}", class.__repr__()) })?; - type_lists_index.insert(type_list, 0); + self.type_lists_index.insert(type_list, 0); } } + + // safe type lists with their offset in the section let mut offset = 0; - let mut type_lists_and_local_offsets = vec![]; - for (i, (list, idx)) in type_lists_index.iter_mut().enumerate() { + for (i, (list, idx)) in self.type_lists_index.iter_mut().enumerate() { while offset % 4 != 0 { // Alignment - section_manager.incr_section_size(Section::TypeList, 1); + self.section_manager.incr_section_size(Section::TypeList, 1); offset += 1; } *idx = i; - type_lists_and_local_offsets.push((list.clone(), offset)); - section_manager.add_elt(Section::TypeList, Some(list.size())); - offset += list.size(); + self.type_lists_with_offset.push((list.clone(), offset)); + self.section_manager + .add_elt(Section::TypeList, Some(list.size())); + offset += list.size() as u32; } + Ok(()) + } + /// Generate the map list. + /// + /// # Warning + /// + /// All sections must be generated (but not necessarely linked) before generating the map list. + fn get_map_list(&mut self) -> Result<()> { debug!("Generate the map_list"); // Get the size of a map item - let map_item_size = MapItem { - type_: MapItemType::HeaderItem, - unused: 0, - size: 0, - offset: 0, - } - .size(); + let map_item_size = 12; /* = MapItem { + type_: MapItemType::HeaderItem, + unused: 0, + size: 0, + offset: 0, + } + .size(); */ // Empty map has a size 4, then we add the size of a MapItem for each element - section_manager.add_elt(Section::MapList, Some(4)); + // The size of the map_list must be computed before generating the map list, + // as it affect the offset of some sections. + self.section_manager.add_elt(Section::MapList, Some(4)); 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); + if !section.is_data() && self.section_manager.get_nb_elt(*section) != 0 { + self.section_manager + .incr_section_size(Section::MapList, map_item_size); } } - let mut map_list = MapList::default(); for section in Section::VARIANT_LIST { - if !section.is_data() && section_manager.get_nb_elt(*section) != 0 { + if !section.is_data() && self.section_manager.get_nb_elt(*section) != 0 { /* match section { // Alignment @@ -493,70 +549,124 @@ impl DexWriter { _ => (), } */ - map_list.list.push(MapItem { + self.map_list.list.push(MapItem { type_: section.get_map_item_type(), unused: 0, - size: section_manager.get_nb_elt(*section) as u32, - offset: section_manager.get_offset(*section), + size: self.section_manager.get_nb_elt(*section) as u32, + offset: self.section_manager.get_offset(*section), }); } } + Ok(()) + } + /// Link the offsets in the header. + /// + /// # Warning + /// + /// Linking can only occur once all sections are entirelly generated. + fn link_header(&mut self) { debug!("Link the header section"); - self.header.map_off = section_manager.get_offset(Section::MapList); - self.header.string_ids_size = section_manager.get_nb_elt(Section::StringIdItem) as u32; - self.header.string_ids_off = section_manager.get_offset(Section::StringIdItem); - self.header.type_ids_size = section_manager.get_nb_elt(Section::TypeIdItem) as u32; - self.header.type_ids_off = section_manager.get_offset(Section::TypeIdItem); - self.header.proto_ids_size = section_manager.get_nb_elt(Section::ProtoIdItem) as u32; - self.header.proto_ids_off = section_manager.get_offset(Section::ProtoIdItem); - self.header.field_ids_size = section_manager.get_nb_elt(Section::FieldIdItem) as u32; - self.header.field_ids_off = section_manager.get_offset(Section::FieldIdItem); - self.header.method_ids_size = section_manager.get_nb_elt(Section::MethodIdItem) as u32; - self.header.method_ids_off = section_manager.get_offset(Section::MethodIdItem); - self.header.class_defs_size = section_manager.get_nb_elt(Section::ClassDefItem) as u32; - self.header.class_defs_off = section_manager.get_offset(Section::ClassDefItem); - self.header.data_size = section_manager.get_size(Section::Data); - self.header.data_off = section_manager.get_offset(Section::Data); + self.header.map_off = self.section_manager.get_offset(Section::MapList); + self.header.string_ids_size = self.section_manager.get_nb_elt(Section::StringIdItem) as u32; + self.header.string_ids_off = self.section_manager.get_offset(Section::StringIdItem); + self.header.type_ids_size = self.section_manager.get_nb_elt(Section::TypeIdItem) as u32; + self.header.type_ids_off = self.section_manager.get_offset(Section::TypeIdItem); + self.header.proto_ids_size = self.section_manager.get_nb_elt(Section::ProtoIdItem) as u32; + self.header.proto_ids_off = self.section_manager.get_offset(Section::ProtoIdItem); + self.header.field_ids_size = self.section_manager.get_nb_elt(Section::FieldIdItem) as u32; + self.header.field_ids_off = self.section_manager.get_offset(Section::FieldIdItem); + self.header.method_ids_size = self.section_manager.get_nb_elt(Section::MethodIdItem) as u32; + 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_off = self.section_manager.get_offset(Section::Data); + } + /// Link the offsets of type_lists items in proto_id_items and class_def_items. + /// + /// # Warning + /// + /// Linking can only occur once all sections are entirelly generated. + fn link_type_list_occurences(&mut self) -> Result<()> { debug!("Link the type_list entries in the proto_id_items and class_def_items"); + // Occurences in proto_ids for (proto, idx) in &self.proto_ids { if !proto.parameters.is_empty() { let type_list = self.gen_type_list(&proto.parameters).with_context(|| { format!("Failed to generate param list for {}", proto.__repr__()) })?; - let offset = section_manager.get_offset(Section::TypeList) - + type_lists_and_local_offsets[*type_lists_index.get(&type_list).unwrap()].1 - as u32; - proto_ids_list[*idx].parameters_off = offset; + let offset = self.section_manager.get_offset(Section::TypeList) + + self.type_lists_with_offset[*self.type_lists_index.get(&type_list).unwrap()] + .1; + self.proto_ids_list[*idx].parameters_off = offset; } } + // Occurences in class_defs for (cls, idx) in self.class_defs.values() { if !cls.interfaces.is_empty() { let type_list = self.gen_type_list(&cls.interfaces).with_context(|| { format!("Failed to generate interface list for {}", cls.__repr__()) })?; - let offset = section_manager.get_offset(Section::TypeList) - + type_lists_and_local_offsets[*type_lists_index.get(&type_list).unwrap()].1 - as u32; - class_defs_list[*idx].interfaces_off = offset; + let offset = self.section_manager.get_offset(Section::TypeList) + + self.type_lists_with_offset[*self.type_lists_index.get(&type_list).unwrap()] + .1; + self.class_defs_list[*idx].interfaces_off = offset; } } + Ok(()) + } + /// Link the offsets of class_data_items in class_def_items. + /// + /// # Warning + /// + /// Linking can only occur once all sections are entirelly generated. + fn link_class_data_occurences(&mut self) -> Result<()> { debug!("Link the class_data_item entries in class_def_items"); for (class, idx) in self.class_defs.values() { if class.has_data_section() { - class_defs_list[*idx].class_data_off += - section_manager.get_offset(Section::ClassDataItem); + self.class_defs_list[*idx].class_data_off += + self.section_manager.get_offset(Section::ClassDataItem); } } + Ok(()) + } + + fn write_dex_file(&mut self, writer: &mut dyn Write) -> Result<()> { + // TODO: SPLIT THIS IN METHODS !!! + self.section_manager.reset(); + self.section_manager.add_elt(Section::HeaderItem, None); + + self.gen_string_data_section()?; + self.gen_type_ids_section()?; + self.gen_proto_ids_section()?; + self.gen_field_ids_section()?; + self.gen_method_ids_section()?; + + debug!("Sort classes and generate the class_defs and class_data section"); + for class_id in self.get_sorted_class_def()? { + self.insert_class_def_item(&class_id)?; + } + self.gen_type_list_section()?; + + // TODO: move to attributes + // Method handles are not ordered, nor deduplicated, so they are generated on the fly + let mut _method_handles: Vec = vec![]; + + self.get_map_list()?; + + self.link_header(); + self.link_type_list_occurences()?; + self.link_class_data_occurences()?; debug!("Serialize the dex file"); // TODO: compute checksum, hash, ect self.header.serialize(writer)?; // StringIdItem section - let mut string_off = section_manager.get_offset(Section::StringDataItem); - for string in string_ids_list.iter() { + let mut string_off = self.section_manager.get_offset(Section::StringDataItem); + for string in self.string_data_list.iter() { let str_id = StringIdItem { string_data_off: string_off, }; @@ -564,23 +674,23 @@ impl DexWriter { string_off += string.size() as u32; } // TypeId section - for ty in type_ids_list { + for ty in &self.type_ids_list { ty.serialize(writer)?; } // ProtoId section - for proto in proto_ids_list { + for proto in &self.proto_ids_list { proto.serialize(writer)?; } // FieldIdItem section - for field_id in field_ids_list { + for field_id in &self.field_ids_list { field_id.serialize(writer)?; } // MethodIdItem section - for method_id in method_ids_list { + for method_id in &self.method_ids_list { method_id.serialize(writer)?; } // ClassDefItem section - for class_def in class_defs_list { + for class_def in &self.class_defs_list { class_def.serialize(writer)?; } // TODO: CallSiteIdItem, @@ -589,10 +699,10 @@ impl DexWriter { handle.serialize(writer)?; } // MapList - map_list.serialize(writer)?; + self.map_list.serialize(writer)?; // TypeList, let mut offset = 0; - for (list, _) in type_lists_and_local_offsets { + for (list, _) in &self.type_lists_with_offset { while offset % 4 != 0 { offset += 1; 0u8.serialize(writer)?; @@ -603,11 +713,11 @@ impl DexWriter { // TODO: AnnotationSetRefList, // TODO: AnnotationSetItem, // ClassDataItem section - for data in class_data_list { + for data in &self.class_data_list { data.serialize(writer)?; } // TODO: CodeItem, - for string in string_ids_list { + for string in &self.string_data_list { string.serialize(writer)?; } @@ -694,6 +804,12 @@ impl DexWriter { } Ok(type_list) } + + pub fn gen_dex_file_to_vec(&mut self) -> Result> { + let mut output = Cursor::new(Vec::::new()); + self.write_dex_file(&mut output)?; + Ok(output.into_inner()) + } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -851,7 +967,7 @@ impl Section { } } -#[derive(Debug, Default)] +#[derive(Debug, Default, Clone)] struct SectionManager { sizes: [u32; Self::NB_SECTION], nb_elt: [usize; Self::NB_SECTION], @@ -860,6 +976,11 @@ struct SectionManager { impl SectionManager { const NB_SECTION: usize = 22; + fn reset(&mut self) { + self.sizes = [0; Self::NB_SECTION]; + self.nb_elt = [0; Self::NB_SECTION] + } + fn add_elt(&mut self, section: Section, size: Option) { if section.is_data() { panic!("Cannot add element directly in section data");