1 files changed, 71 insertions, 25 deletions

diff --git a/futility/updater.c b/futility/updater.c
index 42ae7707..a7179dac 100644
--- a/futility/updater.c
+++ b/futility/updater.c
@@ -15,6 +15,7 @@
 
 #include "2rsa.h"
 #include "crossystem.h"
+#include "fmap.h"
 #include "futility.h"
 #include "host_misc.h"
 #include "updater.h"
@@ -927,13 +928,13 @@ int preserve_firmware_section(const struct firmware_image *image_from,
 	find_firmware_section(&from, image_from, section_name);
 	find_firmware_section(&to, image_to, section_name);
 	if (!from.data || !to.data) {
-		DEBUG("Cannot find section %s: from=%p, to=%p", section_name,
-		      from.data, to.data);
+		DEBUG("Cannot find section %.*s: from=%p, to=%p", FMAP_NAMELEN,
+		      section_name, from.data, to.data);
 		return -1;
 	}
 	if (from.size > to.size) {
-		WARN("%s: Section %s is truncated after updated.",
-		     __FUNCTION__, section_name);
+		WARN("%s: Section %.*s is truncated after updated.",
+		     __FUNCTION__, FMAP_NAMELEN, section_name);
 	}
 	/* Use memmove in case if we need to deal with sections that overlap. */
 	memmove(to.data, from.data, Min(from.size, to.size));
@@ -1015,22 +1016,50 @@ static int preserve_management_engine(struct updater_config *cfg,
 	return try_apply_quirk(QUIRK_UNLOCK_ME_FOR_UPDATE, cfg);
 }
 
+/* Preserve firmware sections by FMAP area flags. */
+static int preserve_fmap_sections(struct firmware_image *from,
+				  struct firmware_image *to,
+				  int *count)
+{
+	int i, errcnt = 0;
+	FmapHeader *fmap = to->fmap_header;
+	FmapAreaHeader *ah = (FmapAreaHeader*)(
+			(uint8_t *)fmap + sizeof(FmapHeader));
+	*count = 0;
+
+	for (i = 0; i < fmap->fmap_nareas; i++, ah++) {
+		if (!(ah->area_flags & FMAP_AREA_PRESERVE))
+			continue;
+		/* Warning: area_name 'may' not end with NUL. */
+		if (!firmware_section_exists(from, ah->area_name)) {
+			DEBUG("FMAP area does not exist in source: %.*s",
+			      FMAP_NAMELEN, ah->area_name);
+			continue;
+		}
+		DEBUG("Preserve FMAP area: %.*s", FMAP_NAMELEN, ah->area_name);
+		errcnt += preserve_firmware_section(from, to, ah->area_name);
+		(*count)++;
+	}
+
+	return errcnt;
+}
+
 /*
- * Preserves the critical sections from the current (active) firmware.
- * Currently preserved sections: GBB (HWID and flags), x86 ME, {RO,RW}_PRESERVE,
- * {RO,RW}_VPD, RW_NVRAM.
- * Returns 0 if success, non-zero if error.
+ * Preserve old images without "preserve" information in FMAP.
+ * We have to use the legacy hard-coded list of names.
  */
-static int preserve_images(struct updater_config *cfg)
+static int preserve_known_sections(struct firmware_image *from,
+				   struct firmware_image *to)
 {
 	int errcnt = 0, i;
-	struct firmware_image *from = &cfg->image_current, *to = &cfg->image;
-	const char * const optional_sections[] = {
-		FMAP_RO_PRESERVE,
-		FMAP_RW_PRESERVE,
-		FMAP_RW_NVRAM,
-		FMAP_RW_ELOG,
-		FMAP_RW_SMMSTORE,
+	const char * const names[] = {
+		"RO_PRESERVE",
+		"RW_PRESERVE",
+		"RO_VPD",
+		"RW_VPD",
+		"SMMSTORE",
+		"RW_NVRAM",
+		"RW_ELOG",
 		/*
 		 * TODO(hungte): b/116326638: Remove RO_FSG after the migration
 		 * is finished.
@@ -1041,23 +1070,40 @@ static int preserve_images(struct updater_config *cfg)
 		  * both migrated to the new FMAP based preserve method.
 		  */
 		"SI_GBE",
-		"SI_PDR"
+		"SI_PDR",
 	};
 
-	errcnt += preserve_gbb(from, to, !cfg->factory_update);
-	errcnt += preserve_management_engine(cfg, from, to);
-	errcnt += preserve_firmware_section(from, to, FMAP_RO_VPD);
-	errcnt += preserve_firmware_section(from, to, FMAP_RW_VPD);
-	for (i = 0; i < ARRAY_SIZE(optional_sections); i++) {
-		if (!firmware_section_exists(from, optional_sections[i]))
+	for (i = 0; i < ARRAY_SIZE(names); i++) {
+		if (!firmware_section_exists(from, names[i]))
 			continue;
-		errcnt += preserve_firmware_section(
-				from, to, optional_sections[i]);
+		DEBUG("Preserve firmware section: %s", names[i]);
+		errcnt += preserve_firmware_section(from, to, names[i]);
 	}
 	return errcnt;
 }
 
 /*
+ * Preserves the critical sections from the current (active) firmware.
+ * Currently preserved sections: GBB (HWID and flags), x86 ME, and any firmware
+ * sections with FMAP_AREA_PRESERVE flag set (or a list of known names).
+ * Returns 0 if success, non-zero if error.
+ */
+static int preserve_images(struct updater_config *cfg)
+{
+	int errcnt = 0, found;
+	struct firmware_image *from = &cfg->image_current, *to = &cfg->image;
+
+	errcnt += preserve_gbb(from, to, !cfg->factory_update);
+	errcnt += preserve_management_engine(cfg, from, to);
+	errcnt += preserve_fmap_sections(from, to, &found);
+
+	if (!found)
+		errcnt += preserve_known_sections(from, to);
+
+	return errcnt;
+}
+
+/*
  * Compares if two sections have same size and data.
  * Returns 0 if given sections are the same, otherwise non-zero.
  */