Merge "Add support for Virtual A/B"
diff --git a/AndroidBoot.mk b/AndroidBoot.mk
index b66ea51..4644514 100644
--- a/AndroidBoot.mk
+++ b/AndroidBoot.mk
@@ -65,6 +65,12 @@
   DYNAMIC_PARTITION_SUPPORT := DYNAMIC_PARTITION_SUPPORT=0
 endif
 
+ifeq ($(PRODUCT_VIRTUAL_AB_OTA),true)
+  VIRTUAL_AB_OTA := VIRTUAL_AB_OTA=1
+else
+  VIRTUAL_AB_OTA := VIRTUAL_AB_OTA=0
+endif
+
 ifeq ($(BOARD_DTBO_NOT_SUPPORTED),true)
   TARGET_DTBO_NOT_SUPPORTED := TARGET_DTBO_NOT_SUPPORTED=1
 else
@@ -141,7 +147,7 @@
 # ELF binary for ABOOT
 TARGET_ABOOT_ELF := $(PRODUCT_OUT)/aboot.elf
 $(TARGET_ABOOT_ELF): $(ABOOT_CLEAN) | $(ABOOT_OUT)
-	$(MAKE) -C $(LK_PATH) TOOLCHAIN_PREFIX=$(CROSS_COMPILE) BOOTLOADER_OUT=$(CROOT_DIR)/$(ABOOT_OUT) $(BOOTLOADER_PLATFORM) $(EMMC_BOOT) $(SIGNED_KERNEL) $(VERIFIED_BOOT) $(VERIFIED_BOOT_2) $(VB1_KEY_USED) $(TARGET_DTBO_NOT_SUPPORTED) $(ENABLE_DISPLAY) $(ENABLE_KASLRSEED) $(ENABLE_BOOTDEVICE_MOUNT) $(DEVICE_STATUS) $(BUILD_VARIANT) $(BOARD_NAME) $(ENABLE_VB_ATTEST) $(OSVERSION_IN_BOOTIMAGE) $(QSEECOM_SECAPP_REGION_2MB) $(TARGET_USE_SYSTEM_AS_ROOT_IMAGE) $(DYNAMIC_PARTITION_SUPPORT)
+	$(MAKE) -C $(LK_PATH) TOOLCHAIN_PREFIX=$(CROSS_COMPILE) BOOTLOADER_OUT=$(CROOT_DIR)/$(ABOOT_OUT) $(BOOTLOADER_PLATFORM) $(EMMC_BOOT) $(SIGNED_KERNEL) $(VERIFIED_BOOT) $(VERIFIED_BOOT_2) $(VB1_KEY_USED) $(TARGET_DTBO_NOT_SUPPORTED) $(ENABLE_DISPLAY) $(ENABLE_KASLRSEED) $(ENABLE_BOOTDEVICE_MOUNT) $(DEVICE_STATUS) $(BUILD_VARIANT) $(BOARD_NAME) $(ENABLE_VB_ATTEST) $(OSVERSION_IN_BOOTIMAGE) $(QSEECOM_SECAPP_REGION_2MB) $(TARGET_USE_SYSTEM_AS_ROOT_IMAGE) $(DYNAMIC_PARTITION_SUPPORT) $(VIRTUAL_AB_OTA)
 
 # NAND variant output
 TARGET_NAND_BOOTLOADER := $(PRODUCT_OUT)/appsboot.mbn
@@ -175,7 +181,7 @@
 
 # Top level for eMMC variant targets
 $(TARGET_EMMC_BOOTLOADER): $(emmc_appsbootldr_clean) | $(EMMC_BOOTLOADER_OUT) $(INSTALLED_KEYSTOREIMAGE_TARGET)
-	$(MAKE) -C $(LK_PATH) TOOLCHAIN_PREFIX=$(CROSS_COMPILE) BOOTLOADER_OUT=$(CROOT_DIR)/$(EMMC_BOOTLOADER_OUT) $(BOOTLOADER_PLATFORM) EMMC_BOOT=1 $(SIGNED_KERNEL) $(VERIFIED_BOOT) $(VERIFIED_BOOT_2) $(VB1_KEY_USED) $(TARGET_DTBO_NOT_SUPPORTED) $(ENABLE_DISPLAY) $(ENABLE_KASLRSEED) $(ENABLE_BOOTDEVICE_MOUNT) $(DEVICE_STATUS) $(BUILD_VARIANT) $(BOARD_NAME) $(ENABLE_VB_ATTEST) $(OSVERSION_IN_BOOTIMAGE) $(ENABLE_BG_SUPPORT) $(QSEECOM_SECAPP_REGION_2MB) $(TARGET_USE_SYSTEM_AS_ROOT_IMAGE) $(DYNAMIC_PARTITION_SUPPORT)
+	$(MAKE) -C $(LK_PATH) TOOLCHAIN_PREFIX=$(CROSS_COMPILE) BOOTLOADER_OUT=$(CROOT_DIR)/$(EMMC_BOOTLOADER_OUT) $(BOOTLOADER_PLATFORM) EMMC_BOOT=1 $(SIGNED_KERNEL) $(VERIFIED_BOOT) $(VERIFIED_BOOT_2) $(VB1_KEY_USED) $(TARGET_DTBO_NOT_SUPPORTED) $(ENABLE_DISPLAY) $(ENABLE_KASLRSEED) $(ENABLE_BOOTDEVICE_MOUNT) $(DEVICE_STATUS) $(BUILD_VARIANT) $(BOARD_NAME) $(ENABLE_VB_ATTEST) $(OSVERSION_IN_BOOTIMAGE) $(ENABLE_BG_SUPPORT) $(QSEECOM_SECAPP_REGION_2MB) $(TARGET_USE_SYSTEM_AS_ROOT_IMAGE) $(DYNAMIC_PARTITION_SUPPORT) $(VIRTUAL_AB_OTA)
 	$(hide) rm -f $(emmc_appsbootldr_clean)
 
 # Keep build NAND & eMMC as default for targets still using TARGET_BOOTLOADER
diff --git a/app/aboot/aboot.c b/app/aboot/aboot.c
index 690e606..405a49b 100644
--- a/app/aboot/aboot.c
+++ b/app/aboot/aboot.c
@@ -2,7 +2,7 @@
  * Copyright (c) 2009, Google Inc.
  * All rights reserved.
  *
- * Copyright (c) 2009-2020, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2009-2021, The Linux Foundation. All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
@@ -285,6 +285,20 @@
 	"devinfo",
 	"partition"};
 
+static const char *VirtualAbCriticalPartitions[] = {
+	"misc",
+	"metadata",
+	"userdata"};
+
+static bool CheckVirtualAbCriticalPartition (const char *PartitionName);
+
+static const char *VabSnapshotMergeStatus[] = {
+	"none",
+	"unknown",
+	"snapshotted",
+	"merging",
+	"cancelled"};
+
 struct atag_ptbl_entry
 {
 	char name[16];
@@ -340,6 +354,7 @@
 char panel_display_mode[MAX_RSP_SIZE];
 char soc_version_str[MAX_RSP_SIZE];
 char block_size_string[MAX_RSP_SIZE];
+static char SnapshotMergeState[MAX_RSP_SIZE];
 #if PRODUCT_IOT
 
 /* For IOT we are using custom version */
@@ -3390,6 +3405,8 @@
 	int index = INVALID_PTN;
 	uint8_t lun = 0;
 	char *footer = NULL;
+	char EraseResultStr[MAX_RSP_SIZE] = "";
+	VirtualAbMergeStatus SnapshotMergeStatus;
 
 #if VERIFIED_BOOT
 	if(!strcmp(arg, KEYSTORE_PTN_NAME))
@@ -3406,6 +3423,29 @@
 	ptn = partition_get_offset(index);
 	size = partition_get_size(index);
 
+	if (target_virtual_ab_supported()) {
+		if (CheckVirtualAbCriticalPartition (arg)) {
+			snprintf(EraseResultStr, MAX_RSP_SIZE,"Erase of %s is not allowed in %s state",
+						arg, SnapshotMergeState);
+			fastboot_fail(EraseResultStr);
+			return;
+		}
+		SnapshotMergeStatus = GetSnapshotMergeStatus ();
+		if (((SnapshotMergeStatus == MERGING) || (SnapshotMergeStatus == SNAPSHOTTED)) &&
+			!strncmp (arg, "super", strlen ("super"))) {
+
+			if(SetSnapshotMergeStatus (CANCELLED))
+			{
+				fastboot_fail("Failed to update snapshot state to cancel");
+				return;
+			}
+
+			//updating fbvar snapshot-merge-state
+			snprintf(SnapshotMergeState,strlen(VabSnapshotMergeStatus[NONE_MERGE_STATUS]) + 1,
+					"%s", VabSnapshotMergeStatus[NONE_MERGE_STATUS]);
+		}
+	}
+
 	if (!strncmp(arg, "avb_custom_key", strlen("avb_custom_key"))) {
                 dprintf(INFO, "erasing avb_custom_key\n");
                 if (erase_userkey()) {
@@ -4131,10 +4171,28 @@
 	return;
 }
 
+static bool CheckVirtualAbCriticalPartition (const char *PartitionName)
+{
+	VirtualAbMergeStatus SnapshotMergeStatus;
+	uint32 Iter = 0;
+
+	SnapshotMergeStatus = GetSnapshotMergeStatus ();
+	if ((SnapshotMergeStatus == MERGING || SnapshotMergeStatus == SNAPSHOTTED)) {
+		for (Iter = 0; Iter < ARRAY_SIZE (VirtualAbCriticalPartitions); Iter++) {
+			if (!strncmp (PartitionName, VirtualAbCriticalPartitions[Iter],
+				strlen (VirtualAbCriticalPartitions[Iter])))
+				return true;
+		}
+	}
+	return false;
+}
+
 void cmd_flash_mmc(const char *arg, void *data, unsigned sz)
 {
 	sparse_header_t *sparse_header;
 	meta_header_t *meta_header;
+	VirtualAbMergeStatus SnapshotMergeStatus;
+	char FlashResultStr[MAX_RSP_SIZE] = "";
 
 #ifdef SSD_ENABLE
 	/* 8 Byte Magic + 2048 Byte xml + Encrypted Data */
@@ -4228,6 +4286,31 @@
 		}
 	}
 #endif
+
+	if (target_virtual_ab_supported())
+	{
+		if (CheckVirtualAbCriticalPartition (arg)) {
+			snprintf(FlashResultStr, MAX_RSP_SIZE,"Flashing of %s is not allowed in %s state",
+					arg, SnapshotMergeState);
+			fastboot_fail(FlashResultStr);
+			return;
+		}
+
+		SnapshotMergeStatus = GetSnapshotMergeStatus ();
+		if (((SnapshotMergeStatus == MERGING) || (SnapshotMergeStatus == SNAPSHOTTED)) &&
+			!strncmp (arg, "super", strlen ("super"))) {
+			if(SetSnapshotMergeStatus (CANCELLED))
+			{
+				fastboot_fail("Failed to update snapshot state to cancel");
+				return;
+			}
+
+			//updating fbvar snapshot-merge-state
+			snprintf(SnapshotMergeState,strlen(VabSnapshotMergeStatus[NONE_MERGE_STATUS]) + 1,
+					"%s", VabSnapshotMergeStatus[NONE_MERGE_STATUS]);
+		}
+	}
+
 	if (!strncmp(arg, "avb_custom_key", strlen("avb_custom_key"))) {
 		dprintf(INFO, "flashing avb_custom_key\n");
 		if (store_userkey(data, sz)) {
@@ -4439,6 +4522,13 @@
 		return;
 	}
 
+	if (target_virtual_ab_supported()) {
+		if (GetSnapshotMergeStatus () == MERGING) {
+			fastboot_fail ("Slot Change is not allowed in merging state");
+			return;
+		}
+	}
+
 	if (arg)
 	{
 		p = strtok_r((char *)arg, ":", &sp);
@@ -4512,6 +4602,50 @@
 	dprintf(CRITICAL, "ERROR: Failed to reboot device\n");
 	return;
 }
+
+#ifdef VIRTUAL_AB_OTA
+void CmdUpdateSnapshot(const char *arg, void *data, unsigned sz)
+{
+	char *command = NULL;
+	const char *delim = ":";
+	char *sp;
+
+	if (arg) {
+		command = strtok_r((char *)arg, delim, &sp);
+		if (command) {
+			command++;
+
+			if(!strncmp (command, "merge", AsciiStrLen ("merge"))) {
+				if (GetSnapshotMergeStatus () == MERGING) {
+					cmd_reboot_fastboot(Arg, Data, Size);
+				}
+				FastbootOkay ("");
+				return;
+			}
+			else if (!strncmp (Command, "cancel", AsciiStrLen ("cancel"))) {
+				if(!device.is_unlocked) {
+					fastboot_fail ("Snapshot Cancel is not allowed in Lock State");
+					return;
+				}
+
+				if (SetSnapshotMergeStatus (CANCELLED))
+				{
+					fastboot_fail("Failed to update snapshot state to cancel");
+					return;
+				}
+
+				//updating fbvar snapshot-merge-state
+				snprintf(SnapshotMergeState, strlen(VabSnapshotMergeStatus[NONE_MERGE_STATUS]) + 1,
+						"%s", VabSnapshotMergeStatus[NONE_MERGE_STATUS]);
+				fastboot_okay("");
+				return;
+			}
+		}
+	}
+	fastboot_fail("Invalid snapshot-update command");
+	return;
+}
+#endif
 #endif
 
 void cmd_reboot_bootloader(const char *arg, void *data, unsigned sz)
@@ -5013,6 +5147,7 @@
 {
 	int i;
 	char hw_platform_buf[MAX_RSP_SIZE];
+	VirtualAbMergeStatus SnapshotMergeStatus;
 
 	struct fastboot_cmd_desc cmd_list[] = {
 						/* By default the enabled list is empty. */
@@ -5047,6 +5182,9 @@
 						{"reboot-fastboot",cmd_reboot_fastboot},
 						{"reboot-recovery",cmd_reboot_recovery},
 #endif
+#ifdef VIRTUAL_AB_OTA
+						{"snapshot-update", CmdUpdateSnapshot},
+#endif
 #if UNITTEST_FW_SUPPORT
 						{"oem run-tests", cmd_oem_runtests},
 #endif
@@ -5133,6 +5271,27 @@
 #endif
         if (target_dynamic_partition_supported())
 		fastboot_publish("is-userspace", "no");
+
+	if (target_virtual_ab_supported()) {
+		SnapshotMergeStatus = GetSnapshotMergeStatus ();
+
+		switch (SnapshotMergeStatus) {
+			case SNAPSHOTTED:
+				SnapshotMergeStatus = SNAPSHOTTED;
+				break;
+			case MERGING:
+				SnapshotMergeStatus = MERGING;
+				break;
+			default:
+				SnapshotMergeStatus = NONE_MERGE_STATUS;
+				break;
+		}
+
+		snprintf(SnapshotMergeState,
+				strlen(VabSnapshotMergeStatus[SnapshotMergeStatus]) + 1,
+				"%s", VabSnapshotMergeStatus[SnapshotMergeStatus]);
+		fastboot_publish("snapshot-update-state", SnapshotMergeState);
+	}
 }
 
 void aboot_init(const struct app_descriptor *app)
diff --git a/app/aboot/recovery.c b/app/aboot/recovery.c
index 389b91d..fd6c281 100644
--- a/app/aboot/recovery.c
+++ b/app/aboot/recovery.c
@@ -1,4 +1,4 @@
-/* Copyright (c) 2010-2017,2019 The Linux Foundation. All rights reserved.
+/* Copyright (c) 2010-2017,2019,2021 The Linux Foundation. All rights reserved.
 
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are
@@ -544,8 +544,11 @@
 
 		ptn = partition_get_offset(index);
 		ptn_size = partition_get_size(index);
-
+#if VIRTUAL_AB_OTA
+		offset = page_offset;
+#else
 		offset = page_offset * BLOCK_SIZE;
+#endif
 		aligned_size = ROUND_TO_PAGE(size, (unsigned)BLOCK_SIZE - 1);
 		if (ptn_size < offset + aligned_size)
 		{
@@ -617,6 +620,63 @@
 	return 0;
 }
 
+static MiscVirtualABMessage *VirtualAbMsg = NULL;
+
+int SetSnapshotMergeStatus (VirtualAbMergeStatus MergeStatus)
+{
+	int Status = 1;
+	VirtualAbMergeStatus OldMergeStatus;
+
+	if (target_is_emmc_boot())
+	{
+		OldMergeStatus = VirtualAbMsg->MergeStatus;
+		VirtualAbMsg->MergeStatus = MergeStatus;
+
+		Status = write_misc(MISC_VIRTUALAB_OFFSET, &VirtualAbMsg, sizeof(VirtualAbMsg));
+		if (Status != 0) {
+			dprintf(CRITICAL, "Write the VirtualAbMsg failed\n");
+			VirtualAbMsg->MergeStatus = OldMergeStatus;
+		}
+	}
+	return Status;
+}
+
+VirtualAbMergeStatus GetSnapshotMergeStatus (void)
+{
+	VirtualAbMergeStatus MergeStatus = NONE_MERGE_STATUS;
+	uint32_t pagesize = get_page_size();
+
+	if (target_is_emmc_boot())
+	{
+	    if (VirtualAbMsg == NULL) {
+	        if(read_misc(MISC_VIRTUALAB_OFFSET, (void *)&VirtualAbMsg,
+					pagesize))
+		{
+			dprintf(CRITICAL,"Error reading virtualab msg from misc partition\n");
+			return MergeStatus;
+		}
+
+		if (VirtualAbMsg->Magic != MISC_VIRTUAL_AB_MAGIC_HEADER ||
+			VirtualAbMsg->Version != MISC_VIRTUAL_AB_MESSAGE_VERSION) {
+
+			dprintf(CRITICAL,"Error read virtualab msg version:%u magic:%u not valid\n",
+					VirtualAbMsg->Version,VirtualAbMsg->Magic);
+
+			free(VirtualAbMsg);
+			VirtualAbMsg = NULL;
+		}
+		else
+		{
+			dprintf(CRITICAL,"read virtualab MergeStatus:%x\n", VirtualAbMsg->MergeStatus);
+		}
+	    }
+
+	    if (VirtualAbMsg)
+	        MergeStatus = VirtualAbMsg->MergeStatus;
+	}
+	return MergeStatus;
+}
+
 int get_ffbm(char *ffbm, unsigned size)
 {
 	const char *ffbm_cmd = "ffbm-";
diff --git a/app/aboot/recovery.h b/app/aboot/recovery.h
index da7d400..77428e8 100644
--- a/app/aboot/recovery.h
+++ b/app/aboot/recovery.h
@@ -1,4 +1,4 @@
-/* Copyright (c) 2010-2017,2019 The Linux Foundation. All rights reserved.
+/* Copyright (c) 2010-2017,2019,2021 The Linux Foundation. All rights reserved.
 
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are
@@ -56,6 +56,34 @@
 	char recovery[1024];
 };
 
+#define MISC_VIRTUAL_AB_MESSAGE_VERSION 2
+#define MISC_VIRTUAL_AB_MAGIC_HEADER 0x56740AB0
+
+/** MISC Partition usage as per AOSP implementation.
+  * 0   - 2K     For bootloader_message
+  * 2K  - 16K    Used by Vendor's bootloader (the 2K - 4K range may be
+  *              optionally used as bootloader_message_ab struct)
+  * 16K - 32K    Used by uncrypt and recovery to store wipe_package
+  *              for A/B devices
+  * 32K - 64K    System space, used for miscellanious AOSP features.
+  **/
+#define MISC_VIRTUALAB_OFFSET (32 * 1024)
+
+typedef enum VirtualAbMergeStatus{
+  NONE_MERGE_STATUS,
+  UNKNOWN_MERGE_STATUS,
+  SNAPSHOTTED,
+  MERGING,
+  CANCELLED
+} VirtualAbMergeStatus;
+
+typedef struct {
+  uint8_t Version;
+  uint32_t Magic;
+  uint8_t MergeStatus;  // IBootControl 1.1, MergeStatus enum.
+  uint8_t SourceStatus;   // Slot number when merge_status was written.
+  uint8_t Reserved[57];
+} __attribute__ ((packed)) MiscVirtualABMessage;
 
 struct update_header {
 	unsigned char MAGIC[UPDATE_MAGIC_SIZE];
@@ -91,7 +119,8 @@
  * invalid.
  */
 int get_ffbm(char *ffbm, unsigned size);
-
+VirtualAbMergeStatus GetSnapshotMergeStatus(void);
+int SetSnapshotMergeStatus (VirtualAbMergeStatus MergeStatus);
 extern unsigned boot_into_recovery;
 
 #endif
diff --git a/include/target.h b/include/target.h
index d4b2416..6886bdf 100644
--- a/include/target.h
+++ b/include/target.h
@@ -1,7 +1,7 @@
 /*
  * Copyright (c) 2008 Travis Geiselbrecht
  *
- * Copyright (c) 2013-2017,2019 The Linux Foundation. All rights reserved.
+ * Copyright (c) 2013-2017,2019,2021 The Linux Foundation. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining
  * a copy of this software and associated documentation files
@@ -105,6 +105,7 @@
 bool is_display_disabled(void);
 bool target_uses_system_as_root(void);
 bool target_dynamic_partition_supported(void);
+bool target_virtual_ab_supported(void);
 struct qmp_reg *target_get_qmp_settings();
 int target_get_qmp_regsize();
 uint32_t target_ddr_cfg_reg();
diff --git a/target/init.c b/target/init.c
index d7b923f..43d5327 100755
--- a/target/init.c
+++ b/target/init.c
@@ -1,7 +1,7 @@
 /*
  * Copyright (c) 2008 Travis Geiselbrecht
  *
- * Copyright (c) 2015-2020, The Linux Foundation. All rights reserved.
+ * Copyright (c) 2015-2021, The Linux Foundation. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining
  * a copy of this software and associated documentation files
@@ -298,6 +298,18 @@
 #endif
 }
 
+#if VIRTUAL_AB_OTA
+bool target_virtual_ab_supported(void)
+{
+	return true;
+}
+#else
+bool target_virtual_ab_supported(void)
+{
+	return false;
+}
+#endif
+
 /* Default CFG register value */
 uint32_t target_ddr_cfg_reg()
 {