[res] Make TargetResourcesContainer valid after an error

ApkResourcesContainer class has a variant that is either a zip
file object, or a full blown AssetManager with ApkAssets with
the same zip file inside. But if transition from one to the
other fails, it could end up in a state with neither, causing a
crash if the users try accessing it.

This change ensures that:
1. The zip file object is only moved into ApkAssets if the
   loading succeeds

2. If any part of the further initialization of ResState fails,
   we take the zip file out and put it back into the variant

+ Add unit tests for both ApkAssets and libidmap2 to ensure
  they don't crash in any of those scenarios (they used to)

+ Enable root tests in libidmap2 as they were never running

Flag: EXEMPT bugfix
Bug: 381108280
Test: atest libandroidfw_tests idmap2_tests + boot
Change-Id: I8f4803fdf03a41ba7a6892e6aed07f04b77788ce

8 files changed, 117 insertions, 28 deletions

diff --git a/cmds/idmap2/Android.bp b/cmds/idmap2/Android.bp
index 58763a7f9aca..d9ff19051de9 100644
--- a/cmds/idmap2/Android.bp
+++ b/cmds/idmap2/Android.bp
@@ -165,6 +165,7 @@ cc_test {
     ],
     host_supported: true,
     test_suites: ["general-tests"],
+    require_root: true,
     srcs: [
         "tests/BinaryStreamVisitorTests.cpp",
         "tests/CommandLineOptionsTests.cpp",
diff --git a/cmds/idmap2/libidmap2/ResourceContainer.cpp b/cmds/idmap2/libidmap2/ResourceContainer.cpp
index 57ae3548123b..f22a481c1f28 100644
--- a/cmds/idmap2/libidmap2/ResourceContainer.cpp
+++ b/cmds/idmap2/libidmap2/ResourceContainer.cpp
@@ -22,6 +22,7 @@
 #include <utility>
 #include <vector>
 
+#include "android-base/scopeguard.h"
 #include "androidfw/ApkAssets.h"
 #include "androidfw/AssetManager.h"
 #include "androidfw/Util.h"
@@ -269,27 +270,40 @@ struct ResState {
   std::unique_ptr<AssetManager2> am;
   ZipAssetsProvider* zip_assets;
 
-  static Result<ResState> Initialize(std::unique_ptr<ZipAssetsProvider> zip,
+  static Result<ResState> Initialize(std::unique_ptr<ZipAssetsProvider>&& zip,
                                      package_property_t flags) {
     ResState state;
     state.zip_assets = zip.get();
     if ((state.apk_assets = ApkAssets::Load(std::move(zip), flags)) == nullptr) {
-      return Error("failed to load apk asset");
+      return Error("failed to load apk asset for '%s'",
+                   state.zip_assets->GetDebugName().c_str());
     }
 
+    // Make sure we put ZipAssetsProvider where we took it if initialization fails, so the
+    // original object stays valid for any next call it may get.
+    auto scoped_restore_zip_assets = android::base::ScopeGuard([&zip, &state]() {
+      zip = std::unique_ptr<ZipAssetsProvider>(
+          static_cast<ZipAssetsProvider*>(
+              std::move(const_cast<ApkAssets&>(*state.apk_assets)).TakeAssetsProvider().release()));
+    });
+
     if ((state.arsc = state.apk_assets->GetLoadedArsc()) == nullptr) {
-      return Error("failed to retrieve loaded arsc");
+      return Error("failed to retrieve loaded arsc for '%s'",
+                   state.zip_assets->GetDebugName().c_str());
     }
 
     if ((state.package = GetPackageAtIndex0(state.arsc)) == nullptr) {
-      return Error("failed to retrieve loaded package at index 0");
+      return Error("failed to retrieve loaded package at index 0 for '%s'",
+                   state.zip_assets->GetDebugName().c_str());
     }
 
     state.am = std::make_unique<AssetManager2>();
     if (!state.am->SetApkAssets({state.apk_assets}, false)) {
-      return Error("failed to create asset manager");
+      return Error("failed to create asset manager for '%s'",
+                   state.zip_assets->GetDebugName().c_str());
     }
 
+    scoped_restore_zip_assets.Disable();
     return state;
   }
 };
diff --git a/cmds/idmap2/tests/IdmapTests.cpp b/cmds/idmap2/tests/IdmapTests.cpp
index 1b656e8c2088..7093614f4047 100644
--- a/cmds/idmap2/tests/IdmapTests.cpp
+++ b/cmds/idmap2/tests/IdmapTests.cpp
@@ -214,6 +214,20 @@ TEST(IdmapTests, CreateIdmapHeaderFromApkAssets) {
   ASSERT_EQ(idmap->GetHeader()->GetOverlayName(), TestConstants::OVERLAY_NAME_ALL_POLICIES);
 }
 
+TEST(IdmapTests, TargetContainerWorksAfterError) {
+  auto target = TargetResourceContainer::FromPath(GetTestDataPath() + "/target/target-bad.apk");
+  ASSERT_TRUE(target);
+
+  auto crc = target->get()->GetCrc();
+  ASSERT_TRUE(crc);
+
+  // This call tries to construct the full ApkAssets state, and fails.
+  ASSERT_FALSE(target->get()->DefinesOverlayable());
+  auto crc2 = target->get()->GetCrc();
+  ASSERT_TRUE(crc2);
+  EXPECT_EQ(*crc, *crc2);
+}
+
 TEST(IdmapTests, CreateIdmapDataFromApkAssets) {
   std::string target_apk_path = GetTestDataPath() + "/target/target.apk";
   std::string overlay_apk_path = GetTestDataPath() + "/overlay/overlay.apk";
diff --git a/cmds/idmap2/tests/data/target/target-bad.apk b/cmds/idmap2/tests/data/target/target-bad.apk
new file mode 100644
index 000000000000..fd8678238c4d
--- /dev/null
+++ b/cmds/idmap2/tests/data/target/target-bad.apk
diff --git a/libs/androidfw/ApkAssets.cpp b/libs/androidfw/ApkAssets.cpp
index 49254d1c6f6e..dbb891455ddd 100644
--- a/libs/androidfw/ApkAssets.cpp
+++ b/libs/androidfw/ApkAssets.cpp
@@ -40,20 +40,21 @@ ApkAssets::ApkAssets(PrivateConstructorUtil, std::unique_ptr<Asset> resources_as
 }
 
 ApkAssetsPtr ApkAssets::Load(const std::string& path, package_property_t flags) {
-  return Load(ZipAssetsProvider::Create(path, flags), flags);
+  return LoadImpl(ZipAssetsProvider::Create(path, flags), flags);
 }
 
 ApkAssetsPtr ApkAssets::LoadFromFd(base::unique_fd fd, const std::string& debug_name,
                                    package_property_t flags, off64_t offset, off64_t len) {
-  return Load(ZipAssetsProvider::Create(std::move(fd), debug_name, offset, len), flags);
+  return LoadImpl(ZipAssetsProvider::Create(std::move(fd), debug_name, offset, len), flags);
 }
 
-ApkAssetsPtr ApkAssets::Load(std::unique_ptr<AssetsProvider> assets, package_property_t flags) {
+ApkAssetsPtr ApkAssets::LoadImpl(std::unique_ptr<AssetsProvider>&& assets,
+                                 package_property_t flags) {
   return LoadImpl(std::move(assets), flags, nullptr /* idmap_asset */, nullptr /* loaded_idmap */);
 }
 
-ApkAssetsPtr ApkAssets::LoadTable(std::unique_ptr<Asset> resources_asset,
-                                  std::unique_ptr<AssetsProvider> assets,
+ApkAssetsPtr ApkAssets::LoadTable(std::unique_ptr<Asset>&& resources_asset,
+                                  std::unique_ptr<AssetsProvider>&& assets,
                                   package_property_t flags) {
   if (resources_asset == nullptr) {
     return {};
@@ -97,10 +98,10 @@ ApkAssetsPtr ApkAssets::LoadOverlay(const std::string& idmap_path, package_prope
                   std::move(loaded_idmap));
 }
 
-ApkAssetsPtr ApkAssets::LoadImpl(std::unique_ptr<AssetsProvider> assets,
+ApkAssetsPtr ApkAssets::LoadImpl(std::unique_ptr<AssetsProvider>&& assets,
                                  package_property_t property_flags,
-                                 std::unique_ptr<Asset> idmap_asset,
-                                 std::unique_ptr<LoadedIdmap> loaded_idmap) {
+                                 std::unique_ptr<Asset>&& idmap_asset,
+                                 std::unique_ptr<LoadedIdmap>&& loaded_idmap) {
   if (assets == nullptr) {
     return {};
   }
@@ -119,11 +120,11 @@ ApkAssetsPtr ApkAssets::LoadImpl(std::unique_ptr<AssetsProvider> assets,
                   std::move(idmap_asset), std::move(loaded_idmap));
 }
 
-ApkAssetsPtr ApkAssets::LoadImpl(std::unique_ptr<Asset> resources_asset,
-                                 std::unique_ptr<AssetsProvider> assets,
+ApkAssetsPtr ApkAssets::LoadImpl(std::unique_ptr<Asset>&& resources_asset,
+                                 std::unique_ptr<AssetsProvider>&& assets,
                                  package_property_t property_flags,
-                                 std::unique_ptr<Asset> idmap_asset,
-                                 std::unique_ptr<LoadedIdmap> loaded_idmap) {
+                                 std::unique_ptr<Asset>&& idmap_asset,
+                                 std::unique_ptr<LoadedIdmap>&& loaded_idmap) {
   if (assets == nullptr ) {
     return {};
   }
diff --git a/libs/androidfw/include/androidfw/ApkAssets.h b/libs/androidfw/include/androidfw/ApkAssets.h
index 1fa67528c78b..231808beb718 100644
--- a/libs/androidfw/include/androidfw/ApkAssets.h
+++ b/libs/androidfw/include/androidfw/ApkAssets.h
@@ -47,13 +47,37 @@ class ApkAssets : public RefBase {
                                  package_property_t flags = 0U, off64_t offset = 0,
                                  off64_t len = AssetsProvider::kUnknownLength);
 
+  //
   // Creates an ApkAssets from an AssetProvider.
-  // The ApkAssets will take care of destroying the AssetsProvider when it is destroyed.
-  static ApkAssetsPtr Load(std::unique_ptr<AssetsProvider> assets, package_property_t flags = 0U);
+  // The ApkAssets will take care of destroying the AssetsProvider when it is destroyed;
+  // the original argument is not moved from if loading fails.
+  //
+  // Note: this function takes care of the case when you pass a move(unique_ptr<Derived>)
+  //    that would create a temporary unique_ptr<AssetsProvider> by moving your pointer into
+  //    it before the function call, making it impossible to not move from the parameter
+  //    on loading failure. The two overloads take care of moving the pointer back if needed.
+  //
+
+  template <class T>
+  static ApkAssetsPtr Load(std::unique_ptr<T>&& assets, package_property_t flags = 0U)
+      requires(std::is_same_v<T, AssetsProvider>) {
+    return LoadImpl(std::move(assets), flags);
+  }
+
+  template <class T>
+  static ApkAssetsPtr Load(std::unique_ptr<T>&& assets, package_property_t flags = 0U)
+      requires(!std::is_same_v<T, AssetsProvider> && std::is_base_of_v<AssetsProvider, T>) {
+    std::unique_ptr<AssetsProvider> base_assets(std::move(assets));
+    auto res = LoadImpl(std::move(base_assets), flags);
+    if (!res) {
+      assets.reset(static_cast<T*>(base_assets.release()));
+    }
+    return res;
+  }
 
   // Creates an ApkAssets from the given asset file representing a resources.arsc.
-  static ApkAssetsPtr LoadTable(std::unique_ptr<Asset> resources_asset,
-                                std::unique_ptr<AssetsProvider> assets,
+  static ApkAssetsPtr LoadTable(std::unique_ptr<Asset>&& resources_asset,
+                                std::unique_ptr<AssetsProvider>&& assets,
                                 package_property_t flags = 0U);
 
   // Creates an ApkAssets from an IDMAP, which contains the original APK path, and the overlay
@@ -94,17 +118,29 @@ class ApkAssets : public RefBase {
 
   bool IsUpToDate() const;
 
+  // DANGER!
+  // This is a destructive method that rips the assets provider out of ApkAssets object.
+  // It is only useful when one knows this assets object can't be used anymore, and they
+  // need the underlying assets provider back (e.g. when initialization fails for some
+  // reason).
+  std::unique_ptr<AssetsProvider> TakeAssetsProvider() && {
+    return std::move(assets_provider_);
+  }
+
  private:
-  static ApkAssetsPtr LoadImpl(std::unique_ptr<AssetsProvider> assets,
+  static ApkAssetsPtr LoadImpl(std::unique_ptr<AssetsProvider>&& assets,
                                package_property_t property_flags,
-                               std::unique_ptr<Asset> idmap_asset,
-                               std::unique_ptr<LoadedIdmap> loaded_idmap);
+                               std::unique_ptr<Asset>&& idmap_asset,
+                               std::unique_ptr<LoadedIdmap>&& loaded_idmap);
 
-  static ApkAssetsPtr LoadImpl(std::unique_ptr<Asset> resources_asset,
-                               std::unique_ptr<AssetsProvider> assets,
+  static ApkAssetsPtr LoadImpl(std::unique_ptr<Asset>&& resources_asset,
+                               std::unique_ptr<AssetsProvider>&& assets,
                                package_property_t property_flags,
-                               std::unique_ptr<Asset> idmap_asset,
-                               std::unique_ptr<LoadedIdmap> loaded_idmap);
+                               std::unique_ptr<Asset>&& idmap_asset,
+                               std::unique_ptr<LoadedIdmap>&& loaded_idmap);
+
+  static ApkAssetsPtr LoadImpl(std::unique_ptr<AssetsProvider>&& assets,
+                               package_property_t flags = 0U);
 
   // Allows us to make it possible to call make_shared from inside the class but still keeps the
   // ctor 'private' for all means and purposes.
diff --git a/libs/androidfw/tests/ApkAssets_test.cpp b/libs/androidfw/tests/ApkAssets_test.cpp
index 70326b71da28..c36d9908032f 100644
--- a/libs/androidfw/tests/ApkAssets_test.cpp
+++ b/libs/androidfw/tests/ApkAssets_test.cpp
@@ -28,6 +28,7 @@ using ::android::base::unique_fd;
 using ::com::android::basic::R;
 using ::testing::Eq;
 using ::testing::Ge;
+using ::testing::IsNull;
 using ::testing::NotNull;
 using ::testing::SizeIs;
 using ::testing::StrEq;
@@ -108,4 +109,26 @@ TEST(ApkAssetsTest, OpenUncompressedAssetFd) {
   EXPECT_THAT(buffer, StrEq("This should be uncompressed.\n\n"));
 }
 
+TEST(ApkAssetsTest, TakeAssetsProviderNotCrashing) {
+  // Make sure the apk assets object can survive taking its assets provider and doesn't crash
+  // the process.
+  {
+    auto loaded_apk = ApkAssets::Load(GetTestDataPath() + "/basic/basic.apk");
+    ASSERT_THAT(loaded_apk, NotNull());
+
+    auto provider = std::move(*loaded_apk).TakeAssetsProvider();
+    ASSERT_THAT(provider, NotNull());
+  }
+  // If this test doesn't crash by this point we're all good.
+}
+
+TEST(ApkAssetsTest, AssetsProviderNotMovedOnError) {
+  auto assets_provider
+      = ZipAssetsProvider::Create(GetTestDataPath() + "/bad/bad.apk", 0);
+  ASSERT_THAT(assets_provider, NotNull());
+  auto loaded_apk = ApkAssets::Load(std::move(assets_provider));
+  ASSERT_THAT(loaded_apk, IsNull());
+  ASSERT_THAT(assets_provider, NotNull());
+}
+
 }  // namespace android
diff --git a/libs/androidfw/tests/data/bad/bad.apk b/libs/androidfw/tests/data/bad/bad.apk
new file mode 100644
index 000000000000..3226bcd52e99
--- /dev/null
+++ b/libs/androidfw/tests/data/bad/bad.apk