3 files changed, 39 insertions, 25 deletions

diff --git a/src/lib/cbfs.c b/src/lib/cbfs.c
index 7df9dc671e..a274551c33 100644
--- a/src/lib/cbfs.c
+++ b/src/lib/cbfs.c
@@ -382,19 +382,29 @@ void *_cbfs_cbmem_allocator(void *arg, size_t size, union cbfs_mdata *unused)
 	return cbmem_add((uintptr_t)arg, size);
 }
 
-int cbfs_prog_stage_load(struct prog *pstage)
+cb_err_t cbfs_prog_stage_load(struct prog *pstage)
 {
+	union cbfs_mdata mdata;
+	struct region_device rdev;
 	struct cbfs_stage stage;
 	uint8_t *load;
 	void *entry;
 	size_t fsize;
 	size_t foffset;
-	const struct region_device *fh = prog_rdev(pstage);
+	cb_err_t err;
 
-	if (rdev_readat(fh, &stage, 0, sizeof(stage)) != sizeof(stage))
-		return -1;
+	prog_locate_hook(pstage);
+
+	if ((err = cbfs_boot_lookup(prog_name(pstage), false, &mdata, &rdev)))
+		return err;
+
+	assert(be32toh(mdata.h.type) == CBFS_TYPE_STAGE);
+	pstage->cbfs_type = CBFS_TYPE_STAGE;
 
-	fsize = region_device_sz(fh);
+	if (rdev_readat(&rdev, &stage, 0, sizeof(stage)) != sizeof(stage))
+		return CB_CBFS_IO;
+
+	fsize = region_device_sz(&rdev);
 	fsize -= sizeof(stage);
 	foffset = 0;
 	foffset += sizeof(stage);
@@ -416,16 +426,16 @@ int cbfs_prog_stage_load(struct prog *pstage)
 	 * that would hit this path initialize themselves. */
 	if ((ENV_BOOTBLOCK || ENV_SEPARATE_VERSTAGE) &&
 	    !CONFIG(NO_XIP_EARLY_STAGES) && CONFIG(BOOT_DEVICE_MEMORY_MAPPED)) {
-		void *mapping = rdev_mmap(fh, foffset, fsize);
-		rdev_munmap(fh, mapping);
+		void *mapping = rdev_mmap(&rdev, foffset, fsize);
+		rdev_munmap(&rdev, mapping);
 		if (mapping == load)
 			goto out;
 	}
 
-	fsize = cbfs_stage_load_and_decompress(fh, foffset, fsize, load,
-					stage.memlen, stage.compression);
+	fsize = cbfs_stage_load_and_decompress(&rdev, foffset, fsize, load,
+					       stage.memlen, stage.compression);
 	if (!fsize)
-		return -1;
+		return CB_ERR;
 
 	/* Clear area not covered by file. */
 	memset(&load[fsize], 0, stage.memlen - fsize);
@@ -436,7 +446,7 @@ out:
 	prog_set_area(pstage, load, stage.memlen);
 	prog_set_entry(pstage, entry, NULL);
 
-	return 0;
+	return CB_SUCCESS;
 }
 
 void cbfs_boot_device_find_mcache(struct cbfs_boot_device *cbd, uint32_t id)
diff --git a/src/lib/prog_loaders.c b/src/lib/prog_loaders.c
index 4722e54501..b31d79341f 100644
--- a/src/lib/prog_loaders.c
+++ b/src/lib/prog_loaders.c
@@ -45,19 +45,16 @@ void run_romstage(void)
 
 	vboot_run_logic();
 
+	timestamp_add_now(TS_START_COPYROM);
+
 	if (ENV_X86 && CONFIG(BOOTBLOCK_NORMAL)) {
-		if (legacy_romstage_selector(&romstage))
+		if (legacy_romstage_select_and_load(&romstage))
 			goto fail;
 	} else {
-		if (prog_locate(&romstage))
+		if (cbfs_prog_stage_load(&romstage))
 			goto fail;
 	}
 
-	timestamp_add_now(TS_START_COPYROM);
-
-	if (cbfs_prog_stage_load(&romstage))
-		goto fail;
-
 	timestamp_add_now(TS_END_COPYROM);
 
 	console_time_report();
@@ -78,6 +75,7 @@ static void run_ramstage_from_resume(struct prog *ramstage)
 	/* Load the cached ramstage to runtime location. */
 	stage_cache_load_stage(STAGE_RAMSTAGE, ramstage);
 
+	ramstage->cbfs_type = CBFS_TYPE_STAGE;
 	prog_set_arg(ramstage, cbmem_top());
 
 	if (prog_entry(ramstage) != NULL) {
@@ -120,9 +118,6 @@ void run_ramstage(void)
 
 	vboot_run_logic();
 
-	if (prog_locate(&ramstage))
-		goto fail;
-
 	timestamp_add_now(TS_START_COPYRAM);
 
 	if (ENV_X86) {
diff --git a/src/lib/rmodule.c b/src/lib/rmodule.c
index 794453527b..6ea9db724b 100644
--- a/src/lib/rmodule.c
+++ b/src/lib/rmodule.c
@@ -2,6 +2,7 @@
 #include <assert.h>
 #include <cbmem.h>
 #include <cbfs.h>
+#include <cbfs_private.h>
 #include <stdint.h>
 #include <stdlib.h>
 #include <string.h>
@@ -244,14 +245,22 @@ int rmodule_stage_load(struct rmod_stage_load *rsl)
 	int load_offset;
 	struct cbfs_stage stage;
 	void *rmod_loc;
-	struct region_device *fh;
+	struct region_device rdev;
+	union cbfs_mdata mdata;
 
 	if (rsl->prog == NULL || prog_name(rsl->prog) == NULL)
 		return -1;
 
-	fh = prog_rdev(rsl->prog);
+	if (prog_locate_hook(rsl->prog))
+		return -1;
+
+	if (cbfs_boot_lookup(prog_name(rsl->prog), false, &mdata, &rdev) != CB_SUCCESS)
+		return -1;
+
+	assert(be32toh(mdata.h.type) == CBFS_TYPE_STAGE);
+	rsl->prog->cbfs_type = CBFS_TYPE_STAGE;
 
-	if (rdev_readat(fh, &stage, 0, sizeof(stage)) != sizeof(stage))
+	if (rdev_readat(&rdev, &stage, 0, sizeof(stage)) != sizeof(stage))
 		return -1;
 
 	rmodule_offset =
@@ -268,7 +277,7 @@ int rmodule_stage_load(struct rmod_stage_load *rsl)
 	printk(BIOS_INFO, "Decompressing stage %s @ %p (%d bytes)\n",
 	       prog_name(rsl->prog), rmod_loc, stage.memlen);
 
-	if (!cbfs_load_and_decompress(fh, sizeof(stage), stage.len, rmod_loc,
+	if (!cbfs_load_and_decompress(&rdev, sizeof(stage), stage.len, rmod_loc,
 				      stage.memlen, stage.compression))
 		return -1;