diff --git a/androscalpel/src/apk.rs b/androscalpel/src/apk.rs index 64e0286..d4cdf02 100644 --- a/androscalpel/src/apk.rs +++ b/androscalpel/src/apk.rs @@ -34,6 +34,68 @@ pub struct DexFile { pub(crate) bin_cache: Option>, // TODO: invalidate the cache !!! } +impl DexFile { + pub fn get_all_types(&self) -> HashSet { + let mut v = TypeCollector::default(); + v.visit_dex_file(self).unwrap(); + v.types + } + + pub fn get_all_method_ids(&self) -> HashSet { + let mut v = MethodIdCollector::default(); + v.visit_dex_file(self).unwrap(); + v.method_ids + } + + pub fn get_all_protos(&self) -> HashSet { + let mut v = MethodTypeCollector::default(); + v.visit_dex_file(self).unwrap(); + v.method_types + } + + pub fn get_all_field_ids(&self) -> HashSet { + let mut v = FieldIdCollector::default(); + v.visit_dex_file(self).unwrap(); + v.field_ids + } + + /// Check if the `DexFile` is exiding the format limits. + pub fn is_overflowing(&self) -> bool { + (self.get_all_types().len() > u16::MAX as usize) + || (self.get_all_protos().len() > u16::MAX as usize) + || (self.get_all_field_ids().len() > u16::MAX as usize) + || (self.get_all_method_ids().len() > u16::MAX as usize) + } + + /// Check if the class can be added to the dex file without exiding the format limits. + pub fn has_space_for_class(&self, class: &Class) -> bool { + (self + .get_all_types() + .union(&class.get_all_types()) + .collect::>() + .len() + <= u16::MAX as usize) + && (self + .get_all_protos() + .union(&class.get_all_protos()) + .collect::>() + .len() + <= u16::MAX as usize) + && (self + .get_all_field_ids() + .union(&class.get_all_field_ids()) + .collect::>() + .len() + <= u16::MAX as usize) + && (self + .get_all_method_ids() + .union(&class.get_all_method_ids()) + .collect::>() + .len() + <= u16::MAX as usize) + } +} + impl Visitable for DexFile { fn default_visit(&self, v: &mut V) -> Result<()> { for (id, class) in &self.classes { @@ -2846,36 +2908,33 @@ impl Apk { .map(|mut dex_writer| dex_writer.gen_dex_file_to_vec()) .collect() */ - let mut bin_dex_files = HashMap::new(); - // TODO: Multithread - for ( + self.dex_files.par_iter().map(|( name, DexFile { classes, not_referenced_strings, bin_cache, }, - ) in self.dex_files.iter() + )| { if let Some(bin_cache) = bin_cache { - bin_dex_files.insert(name.clone(), bin_cache.clone()); + Ok((name.clone(), bin_cache.clone())) } else { let mut writer = DexWriter::new(); for class_ in classes.values() { match writer.add_class(class_) { Ok(()) => (), - Err(DexWritterError::OutOfSpace(_)) => bail!( - "{name} dex file has to many class/method/field to be serialize in dex format" + Err(DexWritterError::OutOfSpace(err)) => bail!( + "{name} dex file has to many class/method/field to be serialize in dex format ({err}). You might want to use `Apk::redistribute_classes() before serializing`" ), } } for string in not_referenced_strings { writer.add_string(string.clone()); } - bin_dex_files.insert(name.clone(), writer.gen_dex_file_to_vec()?); - }; - } - Ok(bin_dex_files) + writer.gen_dex_file_to_vec().map(|bin_dex| (name.clone(), bin_dex)) + } + }).collect::>>() } // TODO: check for android platform classes? @@ -3368,6 +3427,54 @@ impl Apk { } } } + + /// Redistribute classes among dex files. This is needed when a file reference more than 2**16 + /// types, methods, fields or protoypes (restriction imposed by the dalvik format). + /// + ///
+ /// + /// In some edge cases (when several classes in the same application share the same name), this + /// method can change the behavior of the application. + /// + ///
+ pub fn redistribute_classes(&mut self) { + // TODO: handle cases where several classes have the same name. + while let Some(fname) = self.dex_files.iter().find_map(|(name, dex)| { + if dex.is_overflowing() { + Some(name.clone()) + } else { + None + } + }) { + let dex = self.dex_files.get_mut(&fname).unwrap(); + let class_name = dex.classes.keys().find(|_| true).unwrap().clone(); + let class = dex.classes.remove(&class_name).unwrap(); + if let Some(dex) = self + .dex_files + .values_mut() + .find(|dex| dex.has_space_for_class(&class)) + { + dex.classes.insert(class_name, class); + } else { + // If not file has space for the class, create a new file + let mut i = 0; + let new_fname = loop { + let name = if i == 0 { + "classes.dex".into() + } else { + format!("classes{}.dex", i + 1) + }; + if self.dex_files.contains_key(&name) { + break name; + }; + i += 1; + }; + let mut new_dex = DexFile::default(); + new_dex.classes.insert(class_name, class); + self.dex_files.insert(new_fname, new_dex); + } + } + } } /// Parse a .dex file name, and if it is a valid android file, return the index of the file. diff --git a/androscalpel/src/dex_writer.rs b/androscalpel/src/dex_writer.rs index ed57a1e..26e03f4 100644 --- a/androscalpel/src/dex_writer.rs +++ b/androscalpel/src/dex_writer.rs @@ -5,7 +5,7 @@ use std::io; use std::io::{Cursor, Seek, SeekFrom, Write}; use adler::Adler32; -use anyhow::{Context, anyhow, bail}; +use anyhow::{anyhow, bail, Context}; use log::{debug, warn}; use sha1::{Digest, Sha1}; @@ -194,7 +194,7 @@ impl DexWriter { .iter() .filter(|ty| !self.type_ids.contains_key(ty)) .count(); - if new_nb_types >= u16::MAX as usize { + if new_nb_types > u16::MAX as usize { return Err(DexWritterError::OutOfSpace( "To many types for one dex file".into(), )); @@ -206,7 +206,7 @@ impl DexWriter { .iter() .filter(|proto| !self.proto_ids.contains_key(proto)) .count(); - if new_nb_protos >= u16::MAX as usize { + if new_nb_protos > u16::MAX as usize { return Err(DexWritterError::OutOfSpace( "To many prototypes for one dex file".into(), )); @@ -218,7 +218,7 @@ impl DexWriter { .iter() .filter(|field| !self.field_ids.contains_key(field)) .count(); - if new_nb_field_ids >= u16::MAX as usize { + if new_nb_field_ids > u16::MAX as usize { return Err(DexWritterError::OutOfSpace( "To many field ids for one dex file".into(), )); @@ -230,7 +230,7 @@ impl DexWriter { .iter() .filter(|meth| !self.method_ids.contains_key(meth)) .count(); - if new_nb_method_ids >= u16::MAX as usize { + if new_nb_method_ids > u16::MAX as usize { return Err(DexWritterError::OutOfSpace( "To many method ids for one dex file".into(), )); @@ -1769,12 +1769,12 @@ impl DexWriter { debug!("Generate the map_list"); // Get the size of a map item let map_item_size = 12; /* = MapItem { - type_: MapItemType::HeaderItem, - unused: 0, - size: 0, - offset: 0, - } - .size(); */ + 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 // The size of the map_list must be computed before generating the map list, // as it affect the offset of some sections. @@ -2679,12 +2679,12 @@ impl SectionManager { } let mut map_list_size = 4; let map_item_size = 12; /* = MapItem { - type_: MapItemType::HeaderItem, - unused: 0, - size: 0, - offset: 0, - } - .size(); */ + 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)