From 9c3c14e0a645ded1f6d90d1e3d35efade2d51720 Mon Sep 17 00:00:00 2001
From: "A. Cody Schuffelen" <schuffelen@google.com>
Date: Fri, 8 Aug 2025 18:45:01 -0700
Subject: [PATCH] Download META/ab_partitions.txt from the target files zip.

The full target files zip file is not downloaded unless it is requested
by other flags.

Bug: b/439877215
---
 .../host/commands/cvd/fetch/BUILD.bazel       |  3 +++
 .../host/commands/cvd/fetch/fetch_cvd.cc      | 20 +++++++++++++++++--
 base/cvd/cuttlefish/host/libs/zip/zip_cc.cc   |  4 ++++
 base/cvd/cuttlefish/host/libs/zip/zip_cc.h    |  1 +
 4 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/base/cvd/cuttlefish/host/commands/cvd/fetch/BUILD.bazel b/base/cvd/cuttlefish/host/commands/cvd/fetch/BUILD.bazel
index 5981c7c24e4..6b0ffded618 100644
--- a/base/cvd/cuttlefish/host/commands/cvd/fetch/BUILD.bazel
+++ b/base/cvd/cuttlefish/host/commands/cvd/fetch/BUILD.bazel
@@ -138,10 +138,13 @@ cf_cc_library(
         "//cuttlefish/host/libs/web:android_build",
         "//cuttlefish/host/libs/web:android_build_api",
         "//cuttlefish/host/libs/web:android_build_string",
+        "//cuttlefish/host/libs/web:build_api_zip",
         "//cuttlefish/host/libs/web:build_zip_name",
         "//cuttlefish/host/libs/web:chrome_os_build_string",
         "//cuttlefish/host/libs/web:luci_build_api",
         "//cuttlefish/host/libs/web/http_client:curl_global_init",
+        "//cuttlefish/host/libs/zip:zip_cc",
+        "//cuttlefish/host/libs/zip:zip_file",
         "//libbase",
     ],
 )
diff --git a/base/cvd/cuttlefish/host/commands/cvd/fetch/fetch_cvd.cc b/base/cvd/cuttlefish/host/commands/cvd/fetch/fetch_cvd.cc
index 7616617ee3f..e9328590bc3 100644
--- a/base/cvd/cuttlefish/host/commands/cvd/fetch/fetch_cvd.cc
+++ b/base/cvd/cuttlefish/host/commands/cvd/fetch/fetch_cvd.cc
@@ -48,10 +48,13 @@
 #include "cuttlefish/host/libs/web/android_build.h"
 #include "cuttlefish/host/libs/web/android_build_api.h"
 #include "cuttlefish/host/libs/web/android_build_string.h"
+#include "cuttlefish/host/libs/web/build_api_zip.h"
 #include "cuttlefish/host/libs/web/build_zip_name.h"
 #include "cuttlefish/host/libs/web/chrome_os_build_string.h"
 #include "cuttlefish/host/libs/web/http_client/curl_global_init.h"
 #include "cuttlefish/host/libs/web/luci_build_api.h"
+#include "cuttlefish/host/libs/zip/zip_cc.h"
+#include "cuttlefish/host/libs/zip/zip_file.h"
 
 namespace cuttlefish {
 namespace {
@@ -248,10 +251,11 @@ Result<void> FetchDefaultTarget(BuildApi& build_api, const Builds& builds,
     trace.CompletePhase("Desparse image files");
   }
 
+  std::string target_files_name =
+      GetBuildZipName(*builds.default_build, "target_files");
+  ReadableZip target_files_zip = ReadableZip::Invalid();
   if (builds.system || flags.download_target_files_zip) {
     LOG(INFO) << "Downloading target files zip for " << *builds.default_build;
-    std::string target_files_name =
-        GetBuildZipName(*builds.default_build, "target_files");
     std::string target_files = CF_EXPECT(build_api.DownloadFile(
         *builds.default_build, target_directories.default_target_files,
         target_files_name));
@@ -260,7 +264,19 @@ Result<void> FetchDefaultTarget(BuildApi& build_api, const Builds& builds,
     CF_EXPECT(config.AddFilesToConfig(FileSource::DEFAULT_BUILD,
                                       default_build_id, default_build_target,
                                       {target_files}, target_directories.root));
+    target_files_zip = CF_EXPECT(ZipOpenRead(target_files));
+  } else {
+    target_files_zip =
+        CF_EXPECT(OpenZip(build_api, *builds.default_build, target_files_name));
   }
+  std::string ab_partitions = "ab_partitions.txt";
+  std::string ab_partitions_txt_out =
+      target_directories.default_target_files + "/" + ab_partitions;
+  CF_EXPECT(ExtractFile(target_files_zip, "META/" + ab_partitions,
+                        ab_partitions_txt_out));
+  CF_EXPECT(config.AddFilesToConfig(
+      FileSource::DEFAULT_BUILD, default_build_id, default_build_target,
+      {ab_partitions_txt_out}, target_directories.root));
   return {};
 }
 
diff --git a/base/cvd/cuttlefish/host/libs/zip/zip_cc.cc b/base/cvd/cuttlefish/host/libs/zip/zip_cc.cc
index 56ee7802b3e..712aeeb1e0a 100644
--- a/base/cvd/cuttlefish/host/libs/zip/zip_cc.cc
+++ b/base/cvd/cuttlefish/host/libs/zip/zip_cc.cc
@@ -464,6 +464,10 @@ Result<ReadableZip> ReadableZip::FromSource(SeekableZipSource source) {
                                             std::move(fake_writable_source)));
 }
 
+ReadableZip ReadableZip::Invalid() {
+  return ReadableZip(std::unique_ptr<Impl>());
+}
+
 ReadableZip::ReadableZip(ReadableZip&&) = default;
 ReadableZip::~ReadableZip() = default;
 ReadableZip& ReadableZip::operator=(ReadableZip&&) = default;
diff --git a/base/cvd/cuttlefish/host/libs/zip/zip_cc.h b/base/cvd/cuttlefish/host/libs/zip/zip_cc.h
index 91eb33b9ad4..a373dac75e0 100644
--- a/base/cvd/cuttlefish/host/libs/zip/zip_cc.h
+++ b/base/cvd/cuttlefish/host/libs/zip/zip_cc.h
@@ -202,6 +202,7 @@ class ReadableZip {
   friend class WritableZipSource;
 
   static Result<ReadableZip> FromSource(SeekableZipSource);
+  static ReadableZip Invalid();
 
   ReadableZip(ReadableZip&&);
   virtual ~ReadableZip();