these-android-re/androscalpel
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

split write_dex_file into several methods

Browse source
This commit is contained in:
Jean-Marie Mineau 2023-12-05 14:26:20 +01:00
parent bb9f5a94aa
commit 46ad73cbc7
Signed by: histausse
GPG key ID: B66AEEDA9B645AD2
1 changed files with 423 additions and 302 deletions
Download patch file Download diff file Expand all files Collapse all files

725
androscalpel/src/dex_writer.rs
View file

@ -45,6 +45,33 @@ pub struct DexWriter {
// values // values
// annotations // 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<StringDataItem>,
/// The type_ids sections. Once populated, the elements must be sorted according to the spec.
type_ids_list: Vec<TypeIdItem>,
/// 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<ProtoIdItem>,
/// The field_ids section. Once populated, the elements must be sorted according to the spec.
field_ids_list: Vec<FieldIdItem>,
/// The method_ids section. Once populated, the elements must be sorted according to the spec.
method_ids_list: Vec<MethodIdItem>,
/// 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<ClassDefItem>,
/// The class_data section.
class_data_list: Vec<ClassDataItem>,
/// The type lists found in the classes associated to their index in the type_lists section.
type_lists_index: HashMap<TypeList, usize>,
/// 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 { impl Default for DexWriter {
@ -83,6 +110,17 @@ impl Default for DexWriter {
field_ids: HashMap::new(), field_ids: HashMap::new(),
method_ids: HashMap::new(), method_ids: HashMap::new(),
class_defs: 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(()) Ok(())
} }
pub fn gen_dex_file_to_vec(&mut self) -> Result<Vec<u8>> { fn gen_string_data_section(&mut self) -> Result<()> {
let mut output = Cursor::new(Vec::<u8>::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);
debug!("Sort string and generate string_data_item and string_ids sections"); debug!("Sort string and generate string_data_item and string_ids sections");
let mut string_ids_list: Vec<DexString> = self.strings.keys().cloned().collect(); let mut string_ids_list: Vec<DexString> = self.strings.keys().cloned().collect();
string_ids_list.sort(); string_ids_list.sort();
@ -161,35 +190,40 @@ impl DexWriter {
self.strings self.strings
.entry(string.clone()) .entry(string.clone())
.and_modify(|val| *val = idx); .and_modify(|val| *val = idx);
section_manager.add_elt(Section::StringIdItem, None); self.section_manager.add_elt(Section::StringIdItem, None);
section_manager.add_elt(Section::StringDataItem, Some(string.0.size())); self.section_manager
.add_elt(Section::StringDataItem, Some(string.0.size()));
} }
let string_ids_list: Vec<StringDataItem> = string_ids_list self.string_data_list = string_ids_list
.into_iter() .into_iter()
.map(|string| string.into()) .map(|string| string.into())
.collect(); .collect();
Ok(())
}
fn gen_type_ids_section(&mut self) -> Result<()> {
debug!("Sort types and generate type_id_item section"); debug!("Sort types and generate type_id_item section");
let mut type_ids_list: Vec<IdType> = self.type_ids.keys().cloned().collect(); let mut type_ids_list: Vec<IdType> = self.type_ids.keys().cloned().collect();
type_ids_list.sort(); type_ids_list.sort();
for (idx, ty) in type_ids_list.iter().enumerate() { for (idx, ty) in type_ids_list.iter().enumerate() {
self.type_ids.entry(ty.clone()).and_modify(|val| *val = idx); 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<TypeIdItem> = { let mut type_ids_list_aux = vec![];
let mut type_ids_list_aux = vec![]; for ty in type_ids_list.into_iter() {
for ty in type_ids_list.into_iter() { type_ids_list_aux.push(TypeIdItem {
type_ids_list_aux.push(TypeIdItem { descriptor_idx: *self.strings.get(&ty.0).ok_or(anyhow!(
descriptor_idx: *self.strings.get(&ty.0).ok_or(anyhow!( "String {} (name of type {}) not found in dex builder",
"String {} (name of type {}) not found in dex builder", ty.0.__repr__(),
ty.0.__repr__(), ty.__repr__()
ty.__repr__() ))? as u32,
))? as u32, });
}); }
} self.type_ids_list = type_ids_list_aux;
type_ids_list_aux Ok(())
}; }
fn gen_proto_ids_section(&mut self) -> Result<()> {
debug!("Sort prototypes and generate proto_id_item section"); debug!("Sort prototypes and generate proto_id_item section");
let mut proto_ids_list: Vec<IdMethodType> = self.proto_ids.keys().cloned().collect(); let mut proto_ids_list: Vec<IdMethodType> = self.proto_ids.keys().cloned().collect();
proto_ids_list.sort(); proto_ids_list.sort();
@ -197,28 +231,29 @@ impl DexWriter {
self.proto_ids self.proto_ids
.entry(proto.clone()) .entry(proto.clone())
.and_modify(|val| *val = idx); .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![];
let mut proto_ids_list_aux = vec![]; for proto in proto_ids_list {
for proto in proto_ids_list { proto_ids_list_aux.push(ProtoIdItem {
proto_ids_list_aux.push(ProtoIdItem { shorty_idx: *self.strings.get(&proto.shorty).ok_or(anyhow!(
shorty_idx: *self.strings.get(&proto.shorty).ok_or(anyhow!( "String {}, (shorty of prototype {}) not found in dex builder",
"String {}, (shorty of prototype {}) not found in dex builder", proto.shorty.__repr__(),
proto.shorty.__repr__(), proto.__repr__()
proto.__repr__() ))? as u32,
))? as u32, return_type_idx: *self.type_ids.get(&proto.return_type).ok_or(anyhow!(
return_type_idx: *self.type_ids.get(&proto.return_type).ok_or(anyhow!( "Type {}, (return type of prototype {}) not found in dex builder",
"Type {}, (return type of prototype {}) not found in dex builder", proto.shorty.__repr__(),
proto.shorty.__repr__(), proto.__repr__()
proto.__repr__() ))? as u32,
))? as u32, parameters_off: 0, // Will be linked later
parameters_off: 0, // TO BE LINKED LATTER });
}); }
} self.proto_ids_list = proto_ids_list_aux;
proto_ids_list_aux Ok(())
}; }
fn gen_field_ids_section(&mut self) -> Result<()> {
debug!("Sort field ids and generate field_ids_item"); debug!("Sort field ids and generate field_ids_item");
let mut field_ids_list: Vec<IdField> = self.field_ids.keys().cloned().collect(); let mut field_ids_list: Vec<IdField> = self.field_ids.keys().cloned().collect();
field_ids_list.sort(); field_ids_list.sort();
@ -226,32 +261,34 @@ impl DexWriter {
self.field_ids self.field_ids
.entry(field_id.clone()) .entry(field_id.clone())
.and_modify(|val| *val = idx); .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<FieldIdItem> = { let mut field_ids_list_aux = vec![];
let mut field_ids_list_aux = vec![]; for field in field_ids_list.into_iter() {
for field in field_ids_list.into_iter() { field_ids_list_aux.push(FieldIdItem {
field_ids_list_aux.push(FieldIdItem { class_idx: *self.type_ids.get(&field.class_).ok_or(anyhow!(
class_idx: *self.type_ids.get(&field.class_).ok_or(anyhow!( "Type {} (class of field {}) not found in dex builder",
"Type {} (class of field {}) not found in dex builder", field.class_.__repr__(),
field.class_.__repr__(), field.__repr__()
field.__repr__() ))? as u16,
))? as u16, type_idx: *self.type_ids.get(&field.type_).ok_or(anyhow!(
type_idx: *self.type_ids.get(&field.type_).ok_or(anyhow!( "Type {} (type of field {}) not found in dex builder",
"Type {} (type of field {}) not found in dex builder", field.type_.__repr__(),
field.type_.__repr__(), field.__repr__()
field.__repr__() ))? as u16,
))? as u16, name_idx: *self.strings.get(&field.name).ok_or(anyhow!(
name_idx: *self.strings.get(&field.name).ok_or(anyhow!( "String {} (name of field {}) not found in dex builder",
"String {} (name of field {}) not found in dex builder", field.name.__repr__(),
field.name.__repr__(), field.__repr__()
field.__repr__() ))? as u32,
))? as u32, });
}); }
} self.field_ids_list = field_ids_list_aux;
field_ids_list_aux
};
Ok(())
}
fn gen_method_ids_section(&mut self) -> Result<()> {
debug!("Sort method ids and generate method_id_item section"); debug!("Sort method ids and generate method_id_item section");
let mut method_ids_list: Vec<IdMethod> = self.method_ids.keys().cloned().collect(); let mut method_ids_list: Vec<IdMethod> = self.method_ids.keys().cloned().collect();
method_ids_list.sort(); method_ids_list.sort();
@ -259,190 +296,197 @@ impl DexWriter {
self.method_ids self.method_ids
.entry(method_id.clone()) .entry(method_id.clone())
.and_modify(|val| *val = idx); .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<MethodIdItem> = { let mut method_ids_list_aux = vec![];
let mut method_ids_list_aux = vec![]; for method in method_ids_list.into_iter() {
for method in method_ids_list.into_iter() { method_ids_list_aux.push(MethodIdItem {
method_ids_list_aux.push(MethodIdItem { class_idx: *self.type_ids.get(&method.class_).ok_or(anyhow!(
class_idx: *self.type_ids.get(&method.class_).ok_or(anyhow!( "Type {} (class of method {}) not found in dex builder",
"Type {} (class of method {}) not found in dex builder", method.class_.__repr__(),
method.class_.__repr__(), method.__repr__()
method.__repr__() ))? as u16,
))? as u16, proto_idx: *self.proto_ids.get(&method.proto).ok_or(anyhow!(
proto_idx: *self.proto_ids.get(&method.proto).ok_or(anyhow!( "Prototype {} (signature of method {}) not found in dex builder",
"Prototype {} (signature of method {}) not found in dex builder", method.proto.__repr__(),
method.proto.__repr__(), method.__repr__()
method.__repr__() ))? as u16,
))? as u16, name_idx: *self.strings.get(&method.name).ok_or(anyhow!(
name_idx: *self.strings.get(&method.name).ok_or(anyhow!( "String {} (name of method {}) not found in dex builder",
"String {} (name of method {}) not found in dex builder", method.name.__repr__(),
method.name.__repr__(), method.__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__()
))? as u32, ))? 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<IdField> =
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<IdField> =
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<IdMethod> =
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<IdMethod> =
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<IdField> = 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 instance_fields: Vec<IdField> = class.instance_fields.keys().cloned().collect();
let mut _method_handles: Vec<MethodHandleItem> = vec![]; 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<IdMethod> = 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<IdMethod> = 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"); debug!("Generate the type_list section");
let mut type_lists_index = HashMap::new(); // Collect all type lists
for proto in self.proto_ids.keys() { for proto in self.proto_ids.keys() {
if !proto.parameters.is_empty() { if !proto.parameters.is_empty() {
let type_list = self.gen_type_list(&proto.parameters).with_context(|| { let type_list = self.gen_type_list(&proto.parameters).with_context(|| {
format!("Failed to generate param list for {}", proto.__repr__()) 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() { for (class, _) in self.class_defs.values() {
@ -450,42 +494,54 @@ impl DexWriter {
let type_list = self.gen_type_list(&class.interfaces).with_context(|| { let type_list = self.gen_type_list(&class.interfaces).with_context(|| {
format!("Failed to generate interface list for {}", class.__repr__()) 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 offset = 0;
let mut type_lists_and_local_offsets = vec![]; for (i, (list, idx)) in self.type_lists_index.iter_mut().enumerate() {
for (i, (list, idx)) in type_lists_index.iter_mut().enumerate() {
while offset % 4 != 0 { while offset % 4 != 0 {
// Alignment // Alignment
section_manager.incr_section_size(Section::TypeList, 1); self.section_manager.incr_section_size(Section::TypeList, 1);
offset += 1; offset += 1;
} }
*idx = i; *idx = i;
type_lists_and_local_offsets.push((list.clone(), offset)); self.type_lists_with_offset.push((list.clone(), offset));
section_manager.add_elt(Section::TypeList, Some(list.size())); self.section_manager
offset += list.size(); .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"); debug!("Generate the map_list");
// Get the size of a map item // Get the size of a map item
let map_item_size = MapItem { let map_item_size = 12; /* = MapItem {
type_: MapItemType::HeaderItem, type_: MapItemType::HeaderItem,
unused: 0, unused: 0,
size: 0, size: 0,
offset: 0, offset: 0,
} }
.size(); .size(); */
// Empty map has a size 4, then we add the size of a MapItem for each element // 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 { 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 {
section_manager.incr_section_size(Section::MapList, map_item_size); self.section_manager
.incr_section_size(Section::MapList, map_item_size);
} }
} }
let mut map_list = MapList::default();
for section in Section::VARIANT_LIST { 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 { match section {
// Alignment // Alignment
@ -493,70 +549,124 @@ impl DexWriter {
_ => (), _ => (),
} }
*/ */
map_list.list.push(MapItem { self.map_list.list.push(MapItem {
type_: section.get_map_item_type(), type_: section.get_map_item_type(),
unused: 0, unused: 0,
size: section_manager.get_nb_elt(*section) as u32, size: self.section_manager.get_nb_elt(*section) as u32,
offset: section_manager.get_offset(*section), 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"); debug!("Link the header section");
self.header.map_off = section_manager.get_offset(Section::MapList); self.header.map_off = self.section_manager.get_offset(Section::MapList);
self.header.string_ids_size = section_manager.get_nb_elt(Section::StringIdItem) as u32; self.header.string_ids_size = self.section_manager.get_nb_elt(Section::StringIdItem) as u32;
self.header.string_ids_off = section_manager.get_offset(Section::StringIdItem); self.header.string_ids_off = self.section_manager.get_offset(Section::StringIdItem);
self.header.type_ids_size = section_manager.get_nb_elt(Section::TypeIdItem) as u32; self.header.type_ids_size = self.section_manager.get_nb_elt(Section::TypeIdItem) as u32;
self.header.type_ids_off = section_manager.get_offset(Section::TypeIdItem); self.header.type_ids_off = self.section_manager.get_offset(Section::TypeIdItem);
self.header.proto_ids_size = section_manager.get_nb_elt(Section::ProtoIdItem) as u32; self.header.proto_ids_size = self.section_manager.get_nb_elt(Section::ProtoIdItem) as u32;
self.header.proto_ids_off = section_manager.get_offset(Section::ProtoIdItem); self.header.proto_ids_off = self.section_manager.get_offset(Section::ProtoIdItem);
self.header.field_ids_size = section_manager.get_nb_elt(Section::FieldIdItem) as u32; self.header.field_ids_size = self.section_manager.get_nb_elt(Section::FieldIdItem) as u32;
self.header.field_ids_off = section_manager.get_offset(Section::FieldIdItem); self.header.field_ids_off = self.section_manager.get_offset(Section::FieldIdItem);
self.header.method_ids_size = section_manager.get_nb_elt(Section::MethodIdItem) as u32; self.header.method_ids_size = self.section_manager.get_nb_elt(Section::MethodIdItem) as u32;
self.header.method_ids_off = section_manager.get_offset(Section::MethodIdItem); self.header.method_ids_off = self.section_manager.get_offset(Section::MethodIdItem);
self.header.class_defs_size = section_manager.get_nb_elt(Section::ClassDefItem) as u32; self.header.class_defs_size = self.section_manager.get_nb_elt(Section::ClassDefItem) as u32;
self.header.class_defs_off = section_manager.get_offset(Section::ClassDefItem); self.header.class_defs_off = self.section_manager.get_offset(Section::ClassDefItem);
self.header.data_size = section_manager.get_size(Section::Data); self.header.data_size = self.section_manager.get_size(Section::Data);
self.header.data_off = section_manager.get_offset(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"); 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 { for (proto, idx) in &self.proto_ids {
if !proto.parameters.is_empty() { if !proto.parameters.is_empty() {
let type_list = self.gen_type_list(&proto.parameters).with_context(|| { let type_list = self.gen_type_list(&proto.parameters).with_context(|| {
format!("Failed to generate param list for {}", proto.__repr__()) format!("Failed to generate param list for {}", proto.__repr__())
})?; })?;
let offset = section_manager.get_offset(Section::TypeList) let offset = self.section_manager.get_offset(Section::TypeList)
+ type_lists_and_local_offsets[*type_lists_index.get(&type_list).unwrap()].1 + self.type_lists_with_offset[*self.type_lists_index.get(&type_list).unwrap()]
as u32; .1;
proto_ids_list[*idx].parameters_off = offset; self.proto_ids_list[*idx].parameters_off = offset;
} }
} }
// Occurences in class_defs
for (cls, idx) in self.class_defs.values() { for (cls, idx) in self.class_defs.values() {
if !cls.interfaces.is_empty() { if !cls.interfaces.is_empty() {
let type_list = self.gen_type_list(&cls.interfaces).with_context(|| { let type_list = self.gen_type_list(&cls.interfaces).with_context(|| {
format!("Failed to generate interface list for {}", cls.__repr__()) format!("Failed to generate interface list for {}", cls.__repr__())
})?; })?;
let offset = section_manager.get_offset(Section::TypeList) let offset = self.section_manager.get_offset(Section::TypeList)
+ type_lists_and_local_offsets[*type_lists_index.get(&type_list).unwrap()].1 + self.type_lists_with_offset[*self.type_lists_index.get(&type_list).unwrap()]
as u32; .1;
class_defs_list[*idx].interfaces_off = offset; 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"); debug!("Link the class_data_item entries in class_def_items");
for (class, idx) in self.class_defs.values() { for (class, idx) in self.class_defs.values() {
if class.has_data_section() { if class.has_data_section() {
class_defs_list[*idx].class_data_off += self.class_defs_list[*idx].class_data_off +=
section_manager.get_offset(Section::ClassDataItem); 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<MethodHandleItem> = vec![];
self.get_map_list()?;
self.link_header();
self.link_type_list_occurences()?;
self.link_class_data_occurences()?;
debug!("Serialize the dex file"); debug!("Serialize the dex file");
// TODO: compute checksum, hash, ect // TODO: compute checksum, hash, ect
self.header.serialize(writer)?; self.header.serialize(writer)?;
// StringIdItem section // StringIdItem section
let mut string_off = section_manager.get_offset(Section::StringDataItem); let mut string_off = self.section_manager.get_offset(Section::StringDataItem);
for string in string_ids_list.iter() { for string in self.string_data_list.iter() {
let str_id = StringIdItem { let str_id = StringIdItem {
string_data_off: string_off, string_data_off: string_off,
}; };
@ -564,23 +674,23 @@ impl DexWriter {
string_off += string.size() as u32; string_off += string.size() as u32;
} }
// TypeId section // TypeId section
for ty in type_ids_list { for ty in &self.type_ids_list {
ty.serialize(writer)?; ty.serialize(writer)?;
} }
// ProtoId section // ProtoId section
for proto in proto_ids_list { for proto in &self.proto_ids_list {
proto.serialize(writer)?; proto.serialize(writer)?;
} }
// FieldIdItem section // FieldIdItem section
for field_id in field_ids_list { for field_id in &self.field_ids_list {
field_id.serialize(writer)?; field_id.serialize(writer)?;
} }
// MethodIdItem section // MethodIdItem section
for method_id in method_ids_list { for method_id in &self.method_ids_list {
method_id.serialize(writer)?; method_id.serialize(writer)?;
} }
// ClassDefItem section // ClassDefItem section
for class_def in class_defs_list { for class_def in &self.class_defs_list {
class_def.serialize(writer)?; class_def.serialize(writer)?;
} }
// TODO: CallSiteIdItem, // TODO: CallSiteIdItem,
@ -589,10 +699,10 @@ impl DexWriter {
handle.serialize(writer)?; handle.serialize(writer)?;
} }
// MapList // MapList
map_list.serialize(writer)?; self.map_list.serialize(writer)?;
// TypeList, // TypeList,
let mut offset = 0; let mut offset = 0;
for (list, _) in type_lists_and_local_offsets { for (list, _) in &self.type_lists_with_offset {
while offset % 4 != 0 { while offset % 4 != 0 {
offset += 1; offset += 1;
0u8.serialize(writer)?; 0u8.serialize(writer)?;
@ -603,11 +713,11 @@ impl DexWriter {
// TODO: AnnotationSetRefList, // TODO: AnnotationSetRefList,
// TODO: AnnotationSetItem, // TODO: AnnotationSetItem,
// ClassDataItem section // ClassDataItem section
for data in class_data_list { for data in &self.class_data_list {
data.serialize(writer)?; data.serialize(writer)?;
} }
// TODO: CodeItem, // TODO: CodeItem,
for string in string_ids_list { for string in &self.string_data_list {
string.serialize(writer)?; string.serialize(writer)?;
} }
@ -694,6 +804,12 @@ impl DexWriter {
} }
Ok(type_list) Ok(type_list)
} }
pub fn gen_dex_file_to_vec(&mut self) -> Result<Vec<u8>> {
let mut output = Cursor::new(Vec::<u8>::new());
self.write_dex_file(&mut output)?;
Ok(output.into_inner())
}
} }
#[derive(Debug, Clone, Copy, PartialEq, Eq)] #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@ -851,7 +967,7 @@ impl Section {
} }
} }
#[derive(Debug, Default)] #[derive(Debug, Default, Clone)]
struct SectionManager { struct SectionManager {
sizes: [u32; Self::NB_SECTION], sizes: [u32; Self::NB_SECTION],
nb_elt: [usize; Self::NB_SECTION], nb_elt: [usize; Self::NB_SECTION],
@ -860,6 +976,11 @@ struct SectionManager {
impl SectionManager { impl SectionManager {
const NB_SECTION: usize = 22; 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<usize>) { fn add_elt(&mut self, section: Section, size: Option<usize>) {
if section.is_data() { if section.is_data() {
panic!("Cannot add element directly in section data"); panic!("Cannot add element directly in section data");