From 2e735ca34ecf70f5ebdded9e550ca47f78b5e1e6 Mon Sep 17 00:00:00 2001
From: Wei Wang <wvw@google.com>
Date: Tue, 10 May 2016 22:48:13 -0700
Subject: [PATCH] Include BRILLO_VENDOR_PARTITIONS in target zip generation

This change will enable the target zip to include partition defined in
BRILLO_VENDOR_PARTITIONS. This is necessary for OTA to support
partitions other than boot and system in ab_partitions.txt.

The target zip will include images in BRILLO_VENDOR_PARTITIONS if
defined, under VENDOR_IMAGES/ with path kept. Also any vendor partitions
defined in AB_OTA_PARTITIONS will be copied to IMAGES/ in the target zip.

BUG: 28623063
Change-Id: Ic479048dfb8ac506acf827865e784fcb0432a1d5
---
 core/Makefile                                 | 11 +++++++++++
 tools/releasetools/add_img_to_target_files.py | 13 +++++++++++--
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/core/Makefile b/core/Makefile
index 6fb2458916..b3b456ea80 100644
--- a/core/Makefile
+++ b/core/Makefile
@@ -1873,6 +1873,17 @@ ifeq ($(AB_OTA_UPDATER),true)
 	@# Include the build type in META/misc_info.txt so the server can easily differentiate production builds.
 	$(hide) echo "build_type=$(TARGET_BUILD_VARIANT)" >> $(zip_root)/META/misc_info.txt
 	$(hide) echo "ab_update=true" >> $(zip_root)/META/misc_info.txt
+ifdef BRILLO_VENDOR_PARTITIONS
+	$(hide) mkdir -p $(zip_root)/VENDOR_IMAGES
+	$(hide) for f in $(BRILLO_VENDOR_PARTITIONS); do \
+	  pair1="$$(echo $$f | awk -F':' '{print $$1}')"; \
+	  pair2="$$(echo $$f | awk -F':' '{print $$2}')"; \
+	  src=$${pair1}/$${pair2}; \
+	  dest=$(zip_root)/VENDOR_IMAGES/$${pair2}; \
+	  mkdir -p $$(dirname "$${dest}"); \
+	  $(ACP) $${src} $${dest}; \
+	done;
+endif
 ifdef OSRELEASED_DIRECTORY
 	$(hide) $(ACP) $(TARGET_OUT_ETC)/$(OSRELEASED_DIRECTORY)/product_id $(zip_root)/META/product_id.txt
 	$(hide) $(ACP) $(TARGET_OUT_ETC)/$(OSRELEASED_DIRECTORY)/product_version $(zip_root)/META/product_version.txt
diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index f98a281fe3..9e4426367e 100755
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -395,8 +395,9 @@ def AddImagesToTargetFiles(filename):
     banner("partition-table")
     AddPartitionTable(output_zip)
 
-  # For devices using A/B update, copy over images from RADIO/ to IMAGES/ and
-  # make sure we have all the needed images ready under IMAGES/.
+  # For devices using A/B update, copy over images from RADIO/ and/or
+  # VENDOR_IMAGES/ to IMAGES/ and make sure we have all the needed
+  # images ready under IMAGES/. All images should have '.img' as extension.
   ab_partitions = os.path.join(OPTIONS.input_tmp, "META", "ab_partitions.txt")
   if os.path.exists(ab_partitions):
     with open(ab_partitions, 'r') as f:
@@ -404,9 +405,17 @@ def AddImagesToTargetFiles(filename):
     for line in lines:
       img_name = line.strip() + ".img"
       img_radio_path = os.path.join(OPTIONS.input_tmp, "RADIO", img_name)
+      img_vendor_dir = os.path.join(
+        OPTIONS.input_tmp, "VENDOR_IMAGES")
       if os.path.exists(img_radio_path):
         common.ZipWrite(output_zip, img_radio_path,
                         os.path.join("IMAGES", img_name))
+      else:
+        for root, _, files in os.walk(img_vendor_dir):
+          if img_name in files:
+            common.ZipWrite(output_zip, os.path.join(root, img_name),
+              os.path.join("IMAGES", img_name))
+            break
 
       # Zip spec says: All slashes MUST be forward slashes.
       img_path = 'IMAGES/' + img_name