[Core] support dex041(Container format)

Author: teble

Core/dexkit/dex_item.cpp

@@ -28,18 +28,10 @@ namespace dexkit {
 
 inline void PushEncodeNumber(dex::InstructionFormat op_format, uint8_t op, const uint16_t *ptr, std::vector<EncodeNumber> *using_numbers);
 
-DexItem::DexItem(uint32_t id, uint8_t *data, size_t size, DexKit *dexkit) :
-        _image(std::make_unique<MemMap>(data, size)),
-        dexkit(dexkit),
-        reader(_image->data(), _image->len()),
-        dex_id(id) {
-    InitBaseCache();
-}
-
-DexItem::DexItem(uint32_t id, std::unique_ptr<MemMap> mmap, DexKit *dexkit) :
+DexItem::DexItem(uint32_t id, std::shared_ptr<MemMap> mmap, uint32_t header_off, DexKit *dexkit) :
         _image(std::move(mmap)),
         dexkit(dexkit),
-        reader(_image->data(), _image->len()),
+        reader(_image->data(), _image->len(), header_off),
         dex_id(id) {
     InitBaseCache();
 }

Core/dexkit/dexkit.cpp

@@ -51,36 +51,68 @@ Error DexKit::InitFullCache() {
     return Error::SUCCESS;
 }
 
+static inline std::vector<uint32_t> ParseLogicalDexOffsets(const std::shared_ptr<MemMap> &image) {
+    std::vector<uint32_t> offs;
+    const uint8_t *base = image->data();
+    const size_t n = image->len();
+
+    size_t off = 0;
+    while (true) {
+        if (off + sizeof(dex::Header) > n) break;
+        const auto header = reinterpret_cast<const dex::Header *>(base + off);
+        uint32_t sz = header->file_size;
+        if (sz < sizeof(dex::Header) || off + sz > n) break;
+        offs.push_back(static_cast<uint32_t>(off));
+        off += sz;
+        if (off >= n) break;
+    }
+    return offs;
+}
+
 Error DexKit::AddDex(uint8_t *data, size_t size) {
     std::lock_guard lock(_mutex);
-    dex_items.emplace_back(std::make_unique<DexItem>(dex_cnt++, data, size, this));
-    std::sort(dex_items.begin(), dex_items.end(), comp);
+    images.emplace_back(std::make_shared<MemMap>(data, size));
+    auto image = images.back();
+    for (auto off : ParseLogicalDexOffsets(image)) {
+        dex_items.emplace_back(std::make_unique<DexItem>(dex_cnt++, image, off, this));
+    }
     return Error::SUCCESS;
 }
 
 Error DexKit::AddImage(std::unique_ptr<MemMap> dex_image) {
     std::lock_guard lock(_mutex);
-    dex_items.emplace_back(std::make_unique<DexItem>(dex_cnt++, std::move(dex_image), this));
-    std::sort(dex_items.begin(), dex_items.end(), comp);
+    images.emplace_back(std::move(dex_image));
+    auto image = images.back();
+    for (auto off : ParseLogicalDexOffsets(image)) {
+        dex_items.emplace_back(std::make_unique<DexItem>(dex_cnt++, image, off, this));
+    }
     return Error::SUCCESS;
 }
 
 Error DexKit::AddImage(std::vector<std::unique_ptr<MemMap>> dex_images) {
     std::lock_guard lock(_mutex);
-    auto old_size = dex_items.size();
-    auto new_size = old_size + dex_images.size();
-    dex_items.resize(new_size);
+    const auto old_size = images.size();
+    const auto new_size = old_size + dex_images.size();
+    images.resize(new_size);
+    std::vector<std::pair<std::shared_ptr<MemMap>, uint32_t>> add_items;
+    for (auto i = old_size; i < new_size; i++) {
+        images[i] = std::move(dex_images[i]);
+        for (auto off : ParseLogicalDexOffsets(images[i])) {
+            add_items.emplace_back(images[i], off);
+        }
+    }
+    dex_items.resize(dex_items.size() + add_items.size());
     {
         ThreadPool pool(_thread_num);
         auto index = old_size;
-        for (auto &dex_image: dex_images) {
-            pool.enqueue([this, &dex_image, index]() {
-                dex_items[index] = std::make_unique<DexItem>(index, std::move(dex_image), this);
+        for (auto &[image, offset]: add_items) {
+            pool.enqueue([this, &image, index, offset]() {
+                dex_items[index] = std::make_unique<DexItem>(index, std::move(image), offset, this);
             });
             index++;
         }
     }
-    std::sort(dex_items.begin(), dex_items.end(), comp);
+    dex_cnt += add_items.size();
     return Error::SUCCESS;
 }
 
@@ -91,39 +123,55 @@ Error DexKit::AddZipPath(std::string_view apk_path, int unzip_thread_num) {
     }
     auto zip_file = ZipArchive::Open(map);
     if (!zip_file) return Error::OPEN_ZIP_FILE_FAILED;
-    std::vector<std::pair<int, const Entry *>> dex_pairs;
+    std::vector<std::pair<int, const Entry *>> image_pairs;
     for (int idx = 1;; ++idx) {
         auto entry_name = "classes" + (idx == 1 ? std::string() : std::to_string(idx)) + ".dex";
         auto entry = zip_file->Find(entry_name);
         if (!entry) {
             break;
         }
-        dex_pairs.emplace_back(idx, entry);
+        image_pairs.emplace_back(idx, entry);
     }
-    int old_size = (int) dex_items.size();
-    int new_size = (int) (old_size + dex_pairs.size());
-    dex_items.resize(new_size);
+    const auto old_size = images.size();
+    const auto new_size = old_size + image_pairs.size();
+    images.resize(new_size);
     {
         ThreadPool pool(unzip_thread_num == 0 ? _thread_num : unzip_thread_num);
-        for (auto &dex_pair: dex_pairs) {
+        for (auto &dex_pair: image_pairs) {
             pool.enqueue([this, &dex_pair, old_size, &zip_file]() {
                 auto dex_image = zip_file->GetUncompressData(*dex_pair.second);
                 auto ptr = std::make_unique<MemMap>(std::move(dex_image));
                 if (!ptr->ok()) {
                     return;
                 }
-                int idx = old_size + dex_pair.first - 1;
-                dex_items[idx] = std::make_unique<DexItem>(idx, std::move(ptr), this);
+                auto idx = old_size + dex_pair.first - 1;
+                images[idx] = std::move(ptr);
             });
         }
     }
-    dex_cnt += (uint32_t) dex_pairs.size();
+    std::vector<std::pair<std::shared_ptr<MemMap>, uint32_t>> add_items;
+    for (auto i = old_size; i < new_size; i++) {
+        for (auto off : ParseLogicalDexOffsets(images[i])) {
+            add_items.emplace_back(images[i], off);
+        }
+    }
+    dex_items.resize(dex_items.size() + add_items.size());
+    {
+        ThreadPool pool(_thread_num);
+        auto index = old_size;
+        for (auto &[image, offset]: add_items) {
+            pool.enqueue([this, &image, index, offset]() {
+                dex_items[index] = std::make_unique<DexItem>(index, std::move(image), offset, this);
+            });
+            index++;
+        }
+    }
+    dex_cnt += add_items.size();
     return Error::SUCCESS;
 }
 
-Error DexKit::ExportDexFile(std::string_view path) {
-    for (auto &dex_item: dex_items) {
-        auto image = dex_item->GetImage();
+Error DexKit::ExportDexFile(std::string_view path) const {
+    for (const auto &image: images) {
         std::string file_name(path);
         if (file_name.back() != '/') {
             file_name += '/';

Core/dexkit/include/dex_item.h

@@ -48,8 +48,7 @@ class DexKit;
 class DexItem {
 public:
 
-    explicit DexItem(uint32_t id, uint8_t *data, size_t size, DexKit *dexkit);
-    explicit DexItem(uint32_t id, std::unique_ptr<MemMap> mmap, DexKit *dexkit);
+    explicit DexItem(uint32_t id, std::shared_ptr<MemMap> mmap, uint32_t header_off, DexKit *dexkit);
     ~DexItem() = default;
     DexItem(DexItem &&) = default;
     DexItem &operator=(DexItem &&) = default;
@@ -241,7 +240,7 @@ class DexItem {
 
 private:
     DexKit *dexkit;
-    std::unique_ptr<MemMap> _image;
+    std::shared_ptr<MemMap> _image;
     dex::Reader reader;
 
     uint32_t dex_cross_flag = 0;

Core/dexkit/include/dexkit.h

@@ -51,7 +51,7 @@ class DexKit {
     Error AddImage(std::unique_ptr<MemMap> dex_image);
     Error AddImage(std::vector<std::unique_ptr<MemMap>> dex_images);
     Error AddZipPath(std::string_view apk_path, int unzip_thread_num = 0);
-    Error ExportDexFile(std::string_view path);
+    [[nodiscard]] Error ExportDexFile(std::string_view path) const;
     [[nodiscard]] int GetDexNum() const;
 
     std::unique_ptr<flatbuffers::FlatBufferBuilder> FindClass(const schema::FindClass *query);
@@ -89,6 +89,7 @@ class DexKit {
     std::shared_mutex _put_class_mutex;
     std::atomic<uint32_t> dex_cnt = 0;
     uint32_t _thread_num = std::thread::hardware_concurrency();
+    std::vector<std::shared_ptr<MemMap>> images;
     std::vector<std::unique_ptr<DexItem>> dex_items;
     phmap::flat_hash_map<std::string_view, std::pair<uint16_t /*dex_id*/, uint32_t /*type_idx*/>> class_declare_dex_map;
 

Core/third_party/slicer/export/slicer/reader.h

@@ -36,7 +36,7 @@ namespace dex {
 //
 class Reader {
  public:
-  Reader(const dex::u1* image, size_t size);
+  Reader(const dex::u1* image, size_t size, dex::u4 header_off);
   ~Reader() = default;
 
   // No copy/move semantics
@@ -109,15 +109,15 @@ class Reader {
   template <class T>
   const T* ptr(int offset) const {
     SLICER_CHECK_GE(offset, 0 && offset + sizeof(T) <= size_);
-    return reinterpret_cast<const T*>(image_ + offset);
+    return reinterpret_cast<const T*>(image_ + header_off_ + offset);
   }
 
   // Convert a data section file pointer (absolute offset) to an in-memory pointer
   // (offset should be inside the data section)
   template <class T>
   const T* dataPtr(int offset) const {
     SLICER_CHECK_GE(offset, header_->data_off && offset + sizeof(T) <= size_);
-    return reinterpret_cast<const T*>(image_ + offset);
+    return reinterpret_cast<const T*>(image_ + header_off_ + offset);
   }
 
   // Map an indexed section to an ArrayView<T>
@@ -141,6 +141,7 @@ class Reader {
 
   // .dex image header
   const dex::Header* header_;
+  dex::u4 header_off_;
 
   // .dex IR associated with the reader
   std::shared_ptr<ir::DexFile> dex_ir_;

Core/third_party/slicer/reader.cc

@@ -27,7 +27,7 @@
 
 namespace dex {
 
-Reader::Reader(const dex::u1* image, size_t size) : image_(image), size_(size) {
+Reader::Reader(const dex::u1* image, size_t size, dex::u4 header_off = 0) : image_(image), size_(size), header_off_(header_off) {
   // init the header reference
   header_ = ptr<dex::Header>(0);
   ValidateHeader();