8 files changed, 821 insertions, 264 deletions

diff --git a/src/arch/arm64/Makefile.inc b/src/arch/arm64/Makefile.inc
index f57ef720b9..3c3cd20370 100644
--- a/src/arch/arm64/Makefile.inc
+++ b/src/arch/arm64/Makefile.inc
@@ -136,6 +136,7 @@ ramstage-y += memcpy.S
 ramstage-y += memmove.S
 ramstage-$(CONFIG_ARM64_USE_ARM_TRUSTED_FIRMWARE) += arm_tf.c
 ramstage-y += transition.c transition_asm.S
+ramstage-$(CONFIG_PAYLOAD_FIT_SUPPORT) += fit_payload.c
 
 rmodules_arm64-y += memset.S
 rmodules_arm64-y += memcpy.S
diff --git a/src/arch/arm64/fit_payload.c b/src/arch/arm64/fit_payload.c
new file mode 100644
index 0000000000..c4bbceee4c
--- /dev/null
+++ b/src/arch/arm64/fit_payload.c
@@ -0,0 +1,262 @@
+/*
+ * Copyright 2013 Google Inc.
+ * Copyright 2018 Facebook, Inc.
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License as
+ * published by the Free Software Foundation; either version 2 of
+ * the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ */
+
+#include <console/console.h>
+#include <bootmem.h>
+#include <stdlib.h>
+#include <program_loading.h>
+#include <string.h>
+#include <commonlib/compression.h>
+#include <commonlib/cbfs_serialized.h>
+#include <lib.h>
+#include <fit.h>
+#include <endian.h>
+
+#define MAX_KERNEL_SIZE (64*MiB)
+
+struct arm64_kernel_header {
+	u32 code0;
+	u32 code1;
+	u64 text_offset;
+	u64 image_size;
+	u64 flags;
+	u64 res2;
+	u64 res3;
+	u64 res4;
+	u32 magic;
+#define KERNEL_HEADER_MAGIC  0x644d5241
+	u32 res5;
+};
+
+static struct {
+	union {
+		struct arm64_kernel_header header;
+		u8 raw[sizeof(struct arm64_kernel_header) + 0x100];
+	};
+#define SCRATCH_CANARY_VALUE 0xdeadbeef
+	u32 canary;
+} scratch;
+
+/* Returns true if decompressing was successful and it looks like a kernel. */
+static bool decompress_kernel_header(const struct fit_image_node *node)
+{
+	/* Partially decompress to get text_offset. Can't check for errors. */
+	scratch.canary = SCRATCH_CANARY_VALUE;
+	switch (node->compression) {
+	case CBFS_COMPRESS_NONE:
+		memcpy(scratch.raw, node->data, sizeof(scratch.raw));
+		break;
+	case CBFS_COMPRESS_LZMA:
+		ulzman(node->data, node->size,
+		       scratch.raw, sizeof(scratch.raw));
+		break;
+	case CBFS_COMPRESS_LZ4:
+		ulz4fn(node->data, node->size,
+		       scratch.raw, sizeof(scratch.raw));
+		break;
+	default:
+		printk(BIOS_ERR, "ERROR: Unsupported compression algorithm!\n");
+		return false;
+	}
+
+	/* Should never happen, but if it does we'll want to know. */
+	if (scratch.canary != SCRATCH_CANARY_VALUE)
+		die("ERROR: Partial decompression ran over scratchbuf!\n");
+
+	if (scratch.header.magic != KERNEL_HEADER_MAGIC) {
+		printk(BIOS_ERR,
+		       "ERROR: Invalid kernel magic: %#.8x\n != %#.8x\n",
+		       scratch.header.magic, KERNEL_HEADER_MAGIC);
+		return false;
+	}
+
+	/**
+	 * Prior to v3.17, the endianness of text_offset was not specified.  In
+	 * these cases image_size is zero and text_offset is 0x80000 in the
+	 * endianness of the kernel.  Where image_size is non-zero image_size is
+	 * little-endian and must be respected.  Where image_size is zero,
+	 * text_offset can be assumed to be 0x80000.
+	 */
+	if (!scratch.header.image_size)
+		scratch.header.text_offset = cpu_to_le64(0x80000);
+
+	return true;
+}
+
+static size_t get_kernel_size(const struct fit_image_node *node)
+{
+	if (scratch.header.image_size)
+		return le64_to_cpu(scratch.header.image_size);
+
+	/**
+	 * When image_size is zero, a bootloader should attempt to keep as much
+	 * memory as possible free for use by the kernel immediately after the
+	 * end of the kernel image. The amount of space required will vary
+	 * depending on selected features, and is effectively unbound.
+	 */
+
+	printk(BIOS_WARNING, "FIT: image_size not set in kernel header.\n"
+	       "Leaving additional %u MiB of free space after kernel.\n",
+	       MAX_KERNEL_SIZE >> 20);
+
+	return node->size + MAX_KERNEL_SIZE;
+}
+
+static bool fit_place_kernel(const struct range_entry *r, void *arg)
+{
+	struct region *region = arg;
+	resource_t start;
+
+	if (range_entry_tag(r) != BM_MEM_RAM)
+		return true;
+
+	/**
+	 * The Image must be placed text_offset bytes from a 2MB aligned base
+	 * address anywhere in usable system RAM and called there. The region
+	 * between the 2 MB aligned base address and the start of the image has
+	 * no special significance to the kernel, and may be used for other
+	 * purposes.
+	 *
+	 * If the reserved memory (BL31 for example) is smaller than text_offset
+	 * we can use the 2 MiB base address, otherwise use the next 2 MiB page.
+	 * It's not mandatory, but wastes less memory below the kernel.
+	 */
+	start = ALIGN_DOWN(range_entry_base(r), 2 * MiB) +
+		le64_to_cpu(scratch.header.text_offset);
+
+	if (start < range_entry_base(r))
+		start += 2 * MiB;
+	/**
+	 * At least image_size bytes from the start of the image must be free
+	 * for use by the kernel.
+	 */
+	if (start + region->size < range_entry_end(r)) {
+		region->offset = (size_t)start;
+		return false;
+	}
+
+	return true;
+}
+
+/**
+ * Place the region in free memory range.
+ *
+ * The caller has to set region->offset to the minimum allowed address.
+ * The region->offset is usually 0 on kernel >v4.6 and kernel_base + kernel_size
+ * on kernel <v4.6.
+ */
+static bool fit_place_mem(const struct range_entry *r, void *arg)
+{
+	struct region *region = arg;
+	resource_t start;
+
+	if (range_entry_tag(r) != BM_MEM_RAM)
+		return true;
+
+	/* Linux 4.15 doesn't like 4KiB alignment. Align to 1 MiB for now. */
+	start = ALIGN_UP(MAX(region->offset, range_entry_base(r)), 1 * MiB);
+
+	if (start + region->size < range_entry_end(r)) {
+		region->offset = (size_t)start;
+		return false;
+	}
+
+	return true;
+}
+
+bool fit_payload_arch(struct prog *payload, struct fit_config_node *config,
+		      struct region *kernel,
+		      struct region *fdt,
+		      struct region *initrd)
+{
+	bool place_anywhere;
+	void *arg = NULL;
+
+	if (!config->fdt || !fdt) {
+		printk(BIOS_CRIT, "CRIT: Providing a valid FDT is mandatory to "
+		       "boot an ARM64 kernel!\n");
+		return false;
+	}
+
+	if (!decompress_kernel_header(config->kernel_node)) {
+		printk(BIOS_CRIT, "CRIT: Payload doesn't look like an ARM64"
+		       " kernel Image.\n");
+		return false;
+	}
+
+	/* Update kernel size from image header, if possible */
+	kernel->size = get_kernel_size(config->kernel_node);
+	printk(BIOS_DEBUG, "FIT: Using kernel size of 0x%zx bytes\n",
+	       kernel->size);
+
+	/**
+	 * The code assumes that bootmem_walk provides a sorted list of memory
+	 * regions, starting from the lowest address.
+	 * The order of the calls here doesn't matter, as the placement is
+	 * enforced in the called functions.
+	 * For details check code on top.
+	 */
+
+	if (!bootmem_walk(fit_place_kernel, kernel))
+		return false;
+
+	/* Mark as reserved for future allocations. */
+	bootmem_add_range(kernel->offset, kernel->size, BM_MEM_PAYLOAD);
+
+	/**
+	 * NOTE: versions prior to v4.6 cannot make use of memory below the
+	 * physical offset of the Image so it is recommended that the Image be
+	 * placed as close as possible to the start of system RAM.
+	 *
+	 * For kernel <v4.6 the INITRD and FDT can't be placed below the kernel.
+	 * In that case set region offset to an address on top of kernel.
+	 */
+	place_anywhere = !!(le64_to_cpu(scratch.header.flags) & (1 << 3));
+	printk(BIOS_DEBUG, "FIT: Placing FDT and INITRD %s\n",
+	       place_anywhere ? "anywhere" : "on top of kernel");
+
+	/* Place INITRD */
+	if (config->ramdisk) {
+		if (place_anywhere)
+			initrd->offset = 0;
+		else
+			initrd->offset = kernel->offset + kernel->size;
+
+		if (!bootmem_walk(fit_place_mem, initrd))
+			return false;
+		/* Mark as reserved for future allocations. */
+		bootmem_add_range(initrd->offset, initrd->size, BM_MEM_PAYLOAD);
+	}
+
+	/* Place FDT */
+	if (place_anywhere)
+		fdt->offset = 0;
+	else
+		fdt->offset = kernel->offset + kernel->size;
+
+	if (!bootmem_walk(fit_place_mem, fdt))
+		return false;
+	/* Mark as reserved for future allocations. */
+	bootmem_add_range(fdt->offset, fdt->size, BM_MEM_PAYLOAD);
+
+	/* Kernel expects FDT as argument */
+	arg = (void *)fdt->offset;
+
+	prog_set_entry(payload, (void *)kernel->offset, arg);
+
+	bootmem_dump_ranges();
+
+	return true;
+}
diff --git a/src/include/fit.h b/src/include/fit.h
index 1b2f975042..eb51b50edb 100644
--- a/src/include/fit.h
+++ b/src/include/fit.h
@@ -1,8 +1,8 @@
 /*
  * Copyright 2013 Google Inc.
+ * Copyright 2018-present Facebook, Inc.
  *
- * See file CREDITS for list of people who contributed to this
- * project.
+ * Taken from depthcharge: src/boot/fit.h
  *
  * This program is free software; you can redistribute it and/or
  * modify it under the terms of the GNU General Public License as
@@ -15,64 +15,85 @@
  * GNU General Public License for more details.
  */
 
-#ifndef __BOOT_FIT_H__
-#define __BOOT_FIT_H__
+#ifndef __LIB_FIT_H__
+#define __LIB_FIT_H__
 
 #include <stddef.h>
 #include <stdint.h>
+#include <device_tree.h>
+#include <list.h>
+#include <program_loading.h>
 
-#include "base/device_tree.h"
-#include "base/list.h"
-
-typedef enum CompressionType
-{
-	CompressionInvalid,
-	CompressionNone,
-	CompressionLzma,
-	CompressionLz4,
-} CompressionType;
-
-typedef struct FitImageNode
+struct fit_image_node
 {
 	const char *name;
 	void *data;
 	uint32_t size;
-	CompressionType compression;
+	int compression;
 
-	ListNode list_node;
-} FitImageNode;
+	struct list_node list_node;
+};
 
-typedef struct FitConfigNode
+struct fit_config_node
 {
 	const char *name;
 	const char *kernel;
-	FitImageNode *kernel_node;
+	struct fit_image_node *kernel_node;
 	const char *fdt;
-	FitImageNode *fdt_node;
+	struct fit_image_node *fdt_node;
 	const char *ramdisk;
-	FitImageNode *ramdisk_node;
-	FdtProperty compat;
+	struct fit_image_node *ramdisk_node;
+	struct fdt_property compat;
 	int compat_rank;
 	int compat_pos;
+	const char *compat_string;
 
-	ListNode list_node;
-} FitConfigNode;
+	struct list_node list_node;
+};
 
 /*
- * Unpack a FIT image into memory, choosing the right configuration through the
- * compatible string set by fit_add_compat() and unflattening the corresponding
- * kernel device tree.
+ * Updates the cmdline in the devicetree.
+ */
+void fit_update_chosen(struct device_tree *tree, char *cmd_line);
+
+/*
+ * Add a compat string to the list of supported board ids.
+ * Has to be called before fit_load().
+ * The most common use-case would be to implement it on board level.
+ * Strings that were added first have a higher priority on finding a match.
  */
-FitImageNode *fit_load(void *fit, char *cmd_line, DeviceTree **dt);
+void fit_add_compat_string(const char *str);
 
 /*
- * Add a compatible string for the preferred kernel DT to the list for this
- * platform. This should be called before the first fit_load() so it will be
- * ranked as a better match than the default compatible strings. |compat| must
- * stay accessible throughout depthcharge's runtime (i.e. not stack-allocated)!
+ * Updates the memory section in the devicetree.
+ */
+void fit_update_memory(struct device_tree *tree);
+
+/*
+ * Do architecture specific payload placements and fixups.
+ * Set entrypoint and first argument (if any).
+ * @param payload The payload, to set the entry point
+ * @param config The extracted FIT config
+ * @param kernel out-argument where to place the kernel
+ * @param fdt out-argument where to place the devicetree
+ * @param initrd out-argument where to place the initrd (optional)
+ * @return True if all config nodes could be placed, the corresponding
+ *         regions have been updated and the entry point has been set.
+ *         False on error.
+ */
+bool fit_payload_arch(struct prog *payload, struct fit_config_node *config,
+		      struct region *kernel,
+		      struct region *fdt,
+		      struct region *initrd);
+
+/*
+ * Unpack a FIT image into memory, choosing the right configuration through the
+ * compatible string set by fit_add_compat() and return the selected config
+ * node.
  */
-void fit_add_compat(const char *compat);
+struct fit_config_node *fit_load(void *fit);
 
-void fit_add_ramdisk(DeviceTree *tree, void *ramdisk_addr, size_t ramdisk_size);
+void fit_add_ramdisk(struct device_tree *tree, void *ramdisk_addr,
+		     size_t ramdisk_size);
 
-#endif /* __BOOT_FIT_H__ */
+#endif /* __LIB_FIT_H__ */
diff --git a/src/include/fit_payload.h b/src/include/fit_payload.h
new file mode 100644
index 0000000000..dd66289853
--- /dev/null
+++ b/src/include/fit_payload.h
@@ -0,0 +1,26 @@
+/*
+ * Copyright 2013 Google Inc.
+ *
+ * See file CREDITS for list of people who contributed to this
+ * project.
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License as
+ * published by the Free Software Foundation; either version 2 of
+ * the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but without any warranty; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ */
+
+#ifndef __FIT_PAYLOAD_H_
+#define __FIT_PAYLOAD_H_
+
+#include <program_loading.h>
+#include <stdint.h>
+
+void fit_payload(struct prog *payload);
+
+#endif /* __FIT_PAYLOAD_H_ */
diff --git a/src/lib/Makefile.inc b/src/lib/Makefile.inc
index 08ad9b2e4e..4ae80d3961 100644
--- a/src/lib/Makefile.inc
+++ b/src/lib/Makefile.inc
@@ -150,6 +150,8 @@ ramstage-y += b64_decode.c
 ramstage-$(CONFIG_ACPI_NHLT) += nhlt.c
 ramstage-y += list.c
 ramstage-$(CONFIG_FLATTENED_DEVICE_TREE) += device_tree.c
+ramstage-$(CONFIG_PAYLOAD_FIT_SUPPORT) += fit.c
+ramstage-$(CONFIG_PAYLOAD_FIT_SUPPORT) += fit_payload.c
 
 romstage-y += cbmem_common.c
 romstage-y += imd_cbmem.c
diff --git a/src/lib/fit.c b/src/lib/fit.c
index 79af36c0b4..fe8a82ef97 100644
--- a/src/lib/fit.c
+++ b/src/lib/fit.c
@@ -1,8 +1,8 @@
 /*
  * Copyright 2013 Google Inc.
+ * Copyright 2018-present Facebook, Inc.
  *
- * See file CREDITS for list of people who contributed to this
- * project.
+ * Taken from depthcharge: src/boot/fit.c
  *
  * This program is free software; you can redistribute it and/or
  * modify it under the terms of the GNU General Public License as
@@ -17,93 +17,99 @@
 
 #include <assert.h>
 #include <endian.h>
-#include <libpayload.h>
 #include <stdint.h>
-
-#include "base/ranges.h"
-#include "boot/fit.h"
-
-
-
-static ListNode image_nodes;
-static ListNode config_nodes;
-
-static const char *fit_kernel_compat[10] = { NULL };
-static int num_fit_kernel_compat = 0;
-
-void fit_add_compat(const char *compat)
+#include <bootmem.h>
+#include <stdlib.h>
+#include <string.h>
+#include <cbfs.h>
+#include <program_loading.h>
+#include <timestamp.h>
+#include <memrange.h>
+#include <fit.h>
+#include <boardid.h>
+#include <commonlib/include/commonlib/stdlib.h>
+
+static struct list_node image_nodes;
+static struct list_node config_nodes;
+static struct list_node compat_strings;
+
+struct compat_string_entry {
+	const char *compat_string;
+	struct list_node list_node;
+};
+
+/* Convert string to lowercase and replace '_' with '-'. */
+static char *clean_compat_string(char *str)
 {
-	assert(num_fit_kernel_compat < ARRAY_SIZE(fit_kernel_compat));
-	fit_kernel_compat[num_fit_kernel_compat++] = compat;
+	for (size_t i = 0; i < strlen(str); i++) {
+		str[i] = tolower(str[i]);
+		if (str[i] == '_')
+			str[i] = '-';
+	}
+
+	return str;
 }
 
-static void fit_add_default_compats(void)
+static void fit_add_default_compat_strings(void)
 {
-	const char pattern[] = "google,%s-rev%u-sku%u";
-	u32 rev = lib_sysinfo.board_id;
-	u32 sku = lib_sysinfo.sku_id;
+	char compat_string[80] = {};
 
-	static int done = 0;
-	if (done)
-		return;
-	done = 1;
+	if ((board_id() != UNDEFINED_STRAPPING_ID) &&
+	    (sku_id() != UNDEFINED_STRAPPING_ID)) {
+		snprintf(compat_string, sizeof(compat_string),
+			 "%s,%s-rev%u-sku%u", CONFIG_MAINBOARD_VENDOR,
+			 CONFIG_MAINBOARD_PART_NUMBER, board_id(), sku_id());
 
-	char *compat = xmalloc(sizeof(pattern) + sizeof(CONFIG_BOARD) + 20);
-	sprintf(compat, pattern, CONFIG_BOARD,
-		lib_sysinfo.board_id, lib_sysinfo.sku_id);
+		fit_add_compat_string(compat_string);
+	}
 
-	char *c;
-	for (c = compat; *c != '\0'; c++)
-		if (*c == '_')
-			*c = '-';
+	if (board_id() != UNDEFINED_STRAPPING_ID) {
+		snprintf(compat_string, sizeof(compat_string), "%s,%s-rev%u",
+			 CONFIG_MAINBOARD_VENDOR, CONFIG_MAINBOARD_PART_NUMBER,
+			 board_id());
 
-	if (sku != UNDEFINED_STRAPPING_ID && rev != UNDEFINED_STRAPPING_ID)
-		fit_add_compat(strdup(compat));
+		fit_add_compat_string(compat_string);
+	}
 
-	*strrchr(compat, '-') = '\0';
-	if (rev != UNDEFINED_STRAPPING_ID)
-		fit_add_compat(strdup(compat));
+	snprintf(compat_string, sizeof(compat_string), "%s,%s",
+		 CONFIG_MAINBOARD_VENDOR, CONFIG_MAINBOARD_PART_NUMBER);
 
-	*strrchr(compat, '-') = '\0';
-	fit_add_compat(compat);
+	fit_add_compat_string(compat_string);
 }
 
-
-
-
-static void image_node(DeviceTreeNode *node)
+static void image_node(struct device_tree_node *node)
 {
-	FitImageNode *image = xzalloc(sizeof(*image));
-	image->compression = CompressionNone;
+	struct fit_image_node *image = xzalloc(sizeof(*image));
 
+	image->compression = CBFS_COMPRESS_NONE;
 	image->name = node->name;
 
-	DeviceTreeProperty *prop;
+	struct device_tree_property *prop;
 	list_for_each(prop, node->properties, list_node) {
 		if (!strcmp("data", prop->prop.name)) {
 			image->data = prop->prop.data;
 			image->size = prop->prop.size;
 		} else if (!strcmp("compression", prop->prop.name)) {
 			if (!strcmp("none", prop->prop.data))
-				image->compression = CompressionNone;
+				image->compression = CBFS_COMPRESS_NONE;
 			else if (!strcmp("lzma", prop->prop.data))
-				image->compression = CompressionLzma;
+				image->compression = CBFS_COMPRESS_LZMA;
 			else if (!strcmp("lz4", prop->prop.data))
-				image->compression = CompressionLz4;
+				image->compression = CBFS_COMPRESS_LZ4;
 			else
-				image->compression = CompressionInvalid;
+				image->compression = -1;
 		}
 	}
 
 	list_insert_after(&image->list_node, &image_nodes);
 }
 
-static void config_node(DeviceTreeNode *node)
+static void config_node(struct device_tree_node *node)
 {
-	FitConfigNode *config = xzalloc(sizeof(*config));
+	struct fit_config_node *config = xzalloc(sizeof(*config));
 	config->name = node->name;
 
-	DeviceTreeProperty *prop;
+	struct device_tree_property *prop;
 	list_for_each(prop, node->properties, list_node) {
 		if (!strcmp("kernel", prop->prop.name))
 			config->kernel = prop->prop.data;
@@ -116,21 +122,20 @@ static void config_node(DeviceTreeNode *node)
 	list_insert_after(&config->list_node, &config_nodes);
 }
 
-static void fit_unpack(DeviceTree *tree, const char **default_config)
+static void fit_unpack(struct device_tree *tree, const char **default_config)
 {
 	assert(tree && tree->root);
 
-	DeviceTreeNode *top;
+	struct device_tree_node *top;
 	list_for_each(top, tree->root->children, list_node) {
-		DeviceTreeNode *child;
+		struct device_tree_node *child;
 		if (!strcmp("images", top->name)) {
 
 			list_for_each(child, top->children, list_node)
 				image_node(child);
 
 		} else if (!strcmp("configurations", top->name)) {
-
-			DeviceTreeProperty *prop;
+			struct device_tree_property *prop;
 			list_for_each(prop, top->properties, list_node) {
 				if (!strcmp("default", prop->prop.name) &&
 						default_config)
@@ -143,9 +148,9 @@ static void fit_unpack(DeviceTree *tree, const char **default_config)
 	}
 }
 
-static FitImageNode *find_image(const char *name)
+static struct fit_image_node *find_image(const char *name)
 {
-	FitImageNode *image;
+	struct fit_image_node *image;
 	list_for_each(image, image_nodes, list_node) {
 		if (!strcmp(image->name, name))
 			return image;
@@ -153,7 +158,8 @@ static FitImageNode *find_image(const char *name)
 	return NULL;
 }
 
-static int fdt_find_compat(void *blob, uint32_t start_offset, FdtProperty *prop)
+static int fdt_find_compat(void *blob, uint32_t start_offset,
+			   struct fdt_property *prop)
 {
 	int offset = start_offset;
 	int size;
@@ -174,7 +180,8 @@ static int fdt_find_compat(void *blob, uint32_t start_offset, FdtProperty *prop)
 	return -1;
 }
 
-static int fit_check_compat(FdtProperty *compat_prop, const char *compat_name)
+static int fit_check_compat(struct fdt_property *compat_prop,
+			    const char *compat_name)
 {
 	int bytes = compat_prop->size;
 	const char *compat_str = compat_prop->data;
@@ -189,18 +196,21 @@ static int fit_check_compat(FdtProperty *compat_prop, const char *compat_name)
 	return -1;
 }
 
-static void update_chosen(DeviceTree *tree, char *cmd_line)
+void fit_update_chosen(struct device_tree *tree, char *cmd_line)
 {
 	const char *path[] = { "chosen", NULL };
-	DeviceTreeNode *node = dt_find_node(tree->root, path, NULL, NULL, 1);
+	struct device_tree_node *node;
+	node = dt_find_node(tree->root, path, NULL, NULL, 1);
 
 	dt_add_string_prop(node, "bootargs", cmd_line);
 }
 
-void fit_add_ramdisk(DeviceTree *tree, void *ramdisk_addr, size_t ramdisk_size)
+void fit_add_ramdisk(struct device_tree *tree, void *ramdisk_addr,
+		     size_t ramdisk_size)
 {
 	const char *path[] = { "chosen", NULL };
-	DeviceTreeNode *node = dt_find_node(tree->root, path, NULL, NULL, 1);
+	struct device_tree_node *node;
+	node = dt_find_node(tree->root, path, NULL, NULL, 1);
 
 	/* Warning: this assumes the ramdisk is currently located below 4GiB. */
 	u32 start = (uintptr_t)ramdisk_addr;
@@ -210,49 +220,40 @@ void fit_add_ramdisk(DeviceTree *tree, void *ramdisk_addr, size_t ramdisk_size)
 	dt_add_u32_prop(node, "linux,initrd-end", end);
 }
 
-static void update_reserve_map(uint64_t start, uint64_t end, void *data)
+static void update_reserve_map(uint64_t start, uint64_t end,
+			       struct device_tree *tree)
 {
-	DeviceTree *tree = (DeviceTree *)data;
+	struct device_tree_reserve_map_entry *entry = xzalloc(sizeof(*entry));
 
-	DeviceTreeReserveMapEntry *entry = xzalloc(sizeof(*entry));
 	entry->start = start;
 	entry->size = end - start;
 
 	list_insert_after(&entry->list_node, &tree->reserve_map);
 }
 
-typedef struct EntryParams
-{
+struct entry_params {
 	unsigned addr_cells;
 	unsigned size_cells;
 	void *data;
-} EntryParams;
+};
 
 static uint64_t max_range(unsigned size_cells)
 {
-	// Split up ranges who's sizes are too large to fit in #size-cells.
-	// The largest value we can store isn't a power of two, so we'll round
-	// down to make the math easier.
+	/*
+	 * Split up ranges who's sizes are too large to fit in #size-cells.
+	 * The largest value we can store isn't a power of two, so we'll round
+	 * down to make the math easier.
+	 */
 	return 0x1ULL << (size_cells * 32 - 1);
 }
 
-static void count_entries(u64 start, u64 end, void *pdata)
-{
-	EntryParams *params = (EntryParams *)pdata;
-	unsigned *count = (unsigned *)params->data;
-	u64 size = end - start;
-	u64 max_size = max_range(params->size_cells);
-	*count += ALIGN_UP(size, max_size) / max_size;
-}
-
-static void update_mem_property(u64 start, u64 end, void *pdata)
+static void update_mem_property(u64 start, u64 end, struct entry_params *params)
 {
-	EntryParams *params = (EntryParams *)pdata;
 	u8 *data = (u8 *)params->data;
 	u64 full_size = end - start;
 	while (full_size) {
 		const u64 max_size = max_range(params->size_cells);
-		const u32 size = MIN(max_size, full_size);
+		const u64 size = MIN(max_size, full_size);
 
 		dt_write_int(data, start, params->addr_cells * sizeof(u32));
 		data += params->addr_cells * sizeof(uint32_t);
@@ -265,109 +266,189 @@ static void update_mem_property(u64 start, u64 end, void *pdata)
 	params->data = data;
 }
 
-static void update_memory(DeviceTree *tree)
+struct mem_map {
+	struct memranges mem;
+	struct memranges reserved;
+};
+
+static bool walk_memory_table(const struct range_entry *r, void *arg)
 {
-	Ranges mem;
-	Ranges reserved;
-	DeviceTreeNode *node;
+	struct mem_map *arg_map = arg;
+
+	/*
+	 * Kernel likes its available memory areas at least 1MB
+	 * aligned, let's trim the regions such that unaligned padding
+	 * is added to reserved memory.
+	 */
+	if (range_entry_tag(r) == BM_MEM_RAM) {
+		uint64_t new_start = ALIGN_UP(range_entry_base(r), 1 * MiB);
+		uint64_t new_end = ALIGN_DOWN(range_entry_end(r), 1 * MiB);
+
+		if (new_start != range_entry_base(r))
+			memranges_insert(&arg_map->reserved,
+					 range_entry_base(r),
+					 new_start - range_entry_base(r),
+					 BM_MEM_RESERVED);
+
+		if (new_start != new_end)
+			memranges_insert(&arg_map->mem, new_start,
+					 new_end - new_start, BM_MEM_RAM);
+
+		if (new_end != range_entry_end(r))
+			memranges_insert(&arg_map->reserved, new_end,
+					 range_entry_end(r) - new_end,
+					 BM_MEM_RESERVED);
+	} else
+		memranges_insert(&arg_map->reserved, range_entry_base(r),
+				 range_entry_size(r),
+				 BM_MEM_RESERVED);
+
+	return true;
+}
+
+void fit_add_compat_string(const char *str)
+{
+	struct compat_string_entry *compat_node;
+
+	compat_node = xzalloc(sizeof(*compat_node));
+	compat_node->compat_string = strdup(str);
+
+	clean_compat_string((char *)compat_node->compat_string);
+
+	list_insert_after(&compat_node->list_node, &compat_strings);
+}
+
+void fit_update_memory(struct device_tree *tree)
+{
+	const struct range_entry *r;
+	struct device_tree_node *node;
 	u32 addr_cells = 1, size_cells = 1;
+	struct mem_map map;
+
+	printk(BIOS_INFO, "FIT: Updating devicetree memory entries\n");
+
 	dt_read_cell_props(tree->root, &addr_cells, &size_cells);
 
-	// First remove all existing device_type="memory" nodes, then add ours.
+	/*
+	 * First remove all existing device_type="memory" nodes, then add ours.
+	 */
 	list_for_each(node, tree->root->children, list_node) {
 		const char *devtype = dt_find_string_prop(node, "device_type");
 		if (devtype && !strcmp(devtype, "memory"))
 			list_remove(&node->list_node);
 	}
+
 	node = xzalloc(sizeof(*node));
+
 	node->name = "memory";
 	list_insert_after(&node->list_node, &tree->root->children);
-	dt_add_string_prop(node, "device_type", "memory");
-
-	// Read memory info from coreboot (ranges are merged automatically).
-	ranges_init(&mem);
-	ranges_init(&reserved);
-
-#define MEMORY_ALIGNMENT (1 << 20)
-	for (int i = 0; i < lib_sysinfo.n_memranges; i++) {
-		struct memrange *range = &lib_sysinfo.memrange[i];
-		uint64_t start = range->base;
-		uint64_t end = range->base + range->size;
-
-		/*
-		 * Kernel likes its availabe memory areas at least 1MB
-		 * aligned, let's trim the regions such that unaligned padding
-		 * is added to reserved memory.
-		 */
-		if (range->type == CB_MEM_RAM) {
-			uint64_t new_start = ALIGN_UP(start, MEMORY_ALIGNMENT);
-			uint64_t new_end = ALIGN_DOWN(end, MEMORY_ALIGNMENT);
-
-			if (new_start != start)
-				ranges_add(&reserved, start, new_start);
-
-			if (new_start != new_end)
-				ranges_add(&mem, new_start, new_end);
-
-			if (new_end != end)
-				ranges_add(&reserved, new_end, end);
-		} else {
-			ranges_add(&reserved, start, end);
-		}
-	}
+	dt_add_string_prop(node, "device_type", (char *)"memory");
+
+	memranges_init_empty(&map.mem, NULL, 0);
+	memranges_init_empty(&map.reserved, NULL, 0);
 
-	// CBMEM regions are both carved out and explicitly reserved.
-	ranges_for_each(&reserved, &update_reserve_map, tree);
+	bootmem_walk_os_mem(walk_memory_table, &map);
 
-	// Count the amount of 'reg' entries we need (account for size limits).
-	unsigned count = 0;
-	EntryParams count_params = { addr_cells, size_cells, &count };
-	ranges_for_each(&mem, &count_entries, &count_params);
+	/* CBMEM regions are both carved out and explicitly reserved. */
+	memranges_each_entry(r, &map.reserved) {
+		update_reserve_map(range_entry_base(r), range_entry_end(r),
+				   tree);
+	}
+
+	/*
+	 * Count the amount of 'reg' entries we need (account for size limits).
+	 */
+	size_t count = 0;
+	memranges_each_entry(r, &map.mem) {
+		uint64_t size = range_entry_size(r);
+		uint64_t max_size = max_range(size_cells);
+		count += DIV_ROUND_UP(size, max_size);
+	}
 
-	// Allocate the right amount of space and fill up the entries.
+	/* Allocate the right amount of space and fill up the entries. */
 	size_t length = count * (addr_cells + size_cells) * sizeof(u32);
-	void *data = xmalloc(length);
-	EntryParams add_params = { addr_cells, size_cells, data };
-	ranges_for_each(&mem, &update_mem_property, &add_params);
+
+	void *data = xzalloc(length);
+
+	struct entry_params add_params = { addr_cells, size_cells, data };
+	memranges_each_entry(r, &map.mem) {
+		update_mem_property(range_entry_base(r), range_entry_end(r),
+				    &add_params);
+	}
 	assert(add_params.data - data == length);
 
-	// Assemble the final property and add it to the device tree.
+	/* Assemble the final property and add it to the device tree. */
 	dt_add_bin_prop(node, "reg", data, length);
+
+	memranges_teardown(&map.mem);
+	memranges_teardown(&map.reserved);
+}
+
+/*
+ * Finds a compat string and updates the compat position and rank.
+ * @param fdt_blob Pointer to FDT
+ * @param config The current config node to operate on
+ */
+static void fit_update_compat(void *fdt_blob, struct fit_config_node *config)
+{
+	struct compat_string_entry *compat_node;
+	struct fdt_header *fdt_header = (struct fdt_header *)fdt_blob;
+	uint32_t fdt_offset = be32_to_cpu(fdt_header->structure_offset);
+	size_t i = 0;
+
+	if (!fdt_find_compat(fdt_blob, fdt_offset, &config->compat)) {
+		list_for_each(compat_node, compat_strings, list_node) {
+			int pos = fit_check_compat(&config->compat,
+						   compat_node->compat_string);
+			if (pos >= 0) {
+				config->compat_pos = pos;
+				config->compat_rank = i;
+				config->compat_string =
+					compat_node->compat_string;
+				break;
+			}
+			i++;
+		}
+	}
 }
 
-FitImageNode *fit_load(void *fit, char *cmd_line, DeviceTree **dt)
+struct fit_config_node *fit_load(void *fit)
 {
-	FdtHeader *header = (FdtHeader *)fit;
-	FitImageNode *image;
-	FitConfigNode *config;
-	int i;
+	struct fdt_header *header = (struct fdt_header *)fit;
+	struct fit_image_node *image;
+	struct fit_config_node *config;
+	struct compat_string_entry *compat_node;
 
-	printf("Loading FIT.\n");
+	printk(BIOS_DEBUG, "FIT: Loading FIT from %p\n", fit);
 
-	if (betohl(header->magic) != FdtMagic) {
-		printf("Bad FIT header magic value 0x%08x.\n",
-			betohl(header->magic));
+	if (be32toh(header->magic) != FDT_HEADER_MAGIC) {
+		printk(BIOS_ERR, "FIT: Bad header magic value 0x%08x.\n",
+		       be32toh(header->magic));
 		return NULL;
 	}
-
-	DeviceTree *tree = fdt_unflatten(fit);
+	struct device_tree *tree = fdt_unflatten(fit);
 
 	const char *default_config_name = NULL;
-	FitConfigNode *default_config = NULL;
-	FitConfigNode *compat_config = NULL;
+	struct fit_config_node *default_config = NULL;
+	struct fit_config_node *compat_config = NULL;
 
 	fit_unpack(tree, &default_config_name);
 
-	// List the images we found.
+	/* List the images we found. */
 	list_for_each(image, image_nodes, list_node)
-		printf("Image %s has %d bytes.\n", image->name, image->size);
-
-	fit_add_default_compats();
-	printf("Compat preference:");
-	for (i = 0; i < num_fit_kernel_compat; i++)
-		printf(" %s", fit_kernel_compat[i]);
-	printf("\n");
-	// Process and list the configs.
+		printk(BIOS_DEBUG, "FIT: Image %s has %d bytes.\n", image->name,
+		       image->size);
+
+	fit_add_default_compat_strings();
+
+	printk(BIOS_DEBUG, "FIT: Compat preference "
+	       "(lowest to highest priority) :");
+
+	list_for_each(compat_node, compat_strings, list_node) {
+		printk(BIOS_DEBUG, " %s", compat_node->compat_string);
+	}
+	printk(BIOS_DEBUG, "\n");
+	/* Process and list the configs. */
 	list_for_each(config, config_nodes, list_node) {
 		if (config->kernel)
 			config->kernel_node = find_image(config->kernel);
@@ -376,112 +457,82 @@ FitImageNode *fit_load(void *fit, char *cmd_line, DeviceTree **dt)
 		if (config->ramdisk)
 			config->ramdisk_node = find_image(config->ramdisk);
 
+		if (config->ramdisk_node &&
+		    config->ramdisk_node->compression < 0) {
+			printk(BIOS_WARNING, "WARN: Ramdisk is compressed with "
+			       "an unsupported algorithm, discarding config %s."
+			       "\n", config->name);
+			list_remove(&config->list_node);
+			continue;
+		}
+
 		if (!config->kernel_node ||
-				(config->fdt && !config->fdt_node)) {
-			printf("Missing image, discarding config %s.\n",
-				config->name);
+		    (config->fdt && !config->fdt_node)) {
+			printk(BIOS_DEBUG, "FIT: Missing image, discarding "
+			       "config %s.\n", config->name);
 			list_remove(&config->list_node);
 			continue;
 		}
 
 		if (config->fdt_node) {
-			if (config->fdt_node->compression != CompressionNone) {
-				printf("FDT compression not yet supported, "
-				       "skipping config %s.\n", config->name);
+			if (config->fdt_node->compression !=
+			    CBFS_COMPRESS_NONE) {
+				printk(BIOS_DEBUG,
+				       "FIT: FDT compression not yet supported,"
+				       " skipping config %s.\n", config->name);
 				list_remove(&config->list_node);
 				continue;
 			}
 
-			void *fdt_blob = config->fdt_node->data;
-			FdtHeader *fdt_header = (FdtHeader *)fdt_blob;
-			uint32_t fdt_offset =
-				betohl(fdt_header->structure_offset);
 			config->compat_pos = -1;
 			config->compat_rank = -1;
-			if (!fdt_find_compat(fdt_blob, fdt_offset,
-					    &config->compat)) {
-				for (i = 0; i < num_fit_kernel_compat; i++) {
-					int pos = fit_check_compat(
-							&config->compat,
-							fit_kernel_compat[i]);
-					if (pos >= 0) {
-						config->compat_pos = pos;
-						config->compat_rank = i;
-						break;
-					}
-				}
-			}
-		}
 
-		printf("Config %s", config->name);
+			fit_update_compat(config->fdt_node->data, config);
+		}
+		printk(BIOS_DEBUG, "FIT: config %s", config->name);
 		if (default_config_name &&
-				!strcmp(config->name, default_config_name)) {
-			printf(" (default)");
+		    !strcmp(config->name, default_config_name)) {
+			printk(BIOS_DEBUG, " (default)");
 			default_config = config;
 		}
-		printf(", kernel %s", config->kernel);
 		if (config->fdt)
-			printf(", fdt %s", config->fdt);
+			printk(BIOS_DEBUG, ", fdt %s", config->fdt);
 		if (config->ramdisk)
-			printf(", ramdisk %s", config->ramdisk);
+			printk(BIOS_DEBUG, ", ramdisk %s", config->ramdisk);
 		if (config->compat.name) {
-			printf(", compat");
+			printk(BIOS_DEBUG, ", compat");
 			int bytes = config->compat.size;
 			const char *compat_str = config->compat.data;
 			for (int pos = 0; bytes && compat_str[0]; pos++) {
-				printf(" %s", compat_str);
+				printk(BIOS_DEBUG, " %s", compat_str);
 				if (pos == config->compat_pos)
-					printf(" (match)");
+					printk(BIOS_DEBUG, " (match)");
 				int len = strlen(compat_str) + 1;
 				compat_str += len;
 				bytes -= len;
 			}
 
 			if (config->compat_rank >= 0 && (!compat_config ||
-			    config->compat_rank < compat_config->compat_rank))
+			    config->compat_rank > compat_config->compat_rank))
 				compat_config = config;
 		}
-		printf("\n");
+		printk(BIOS_DEBUG, "\n");
 	}
 
-	FitConfigNode *to_boot = NULL;
+	struct fit_config_node *to_boot = NULL;
 	if (compat_config) {
 		to_boot = compat_config;
-		printf("Choosing best match %s for compat %s.\n",
-		       to_boot->name, fit_kernel_compat[to_boot->compat_rank]);
+		printk(BIOS_INFO, "FIT: Choosing best match %s for compat "
+		       "%s.\n", to_boot->name, to_boot->compat_string);
 	} else if (default_config) {
 		to_boot = default_config;
-		printf("No match, choosing default %s.\n", to_boot->name);
+		printk(BIOS_INFO, "FIT: No match, choosing default %s.\n",
+		       to_boot->name);
 	} else {
-		printf("No compatible or default configs. Giving up.\n");
-		// We're leaking memory here, but at this point we're beyond
-		// saving anyway.
+		printk(BIOS_ERR, "FIT: No compatible or default configs. "
+		       "Giving up.\n");
 		return NULL;
 	}
 
-	if (to_boot->fdt_node) {
-		*dt = fdt_unflatten(to_boot->fdt_node->data);
-		if (!*dt) {
-			printf("Failed to unflatten the kernel's fdt.\n");
-			return NULL;
-		}
-
-		/* Update only if non-NULL cmd line */
-		if (cmd_line)
-			update_chosen(*dt, cmd_line);
-
-		update_memory(*dt);
-
-		if (to_boot->ramdisk_node) {
-			if (to_boot->ramdisk_node->compression
-					!= CompressionNone) {
-				printf("Ramdisk compression not supported.\n");
-				return NULL;
-			}
-			fit_add_ramdisk(*dt, to_boot->ramdisk_node->data,
-					to_boot->ramdisk_node->size);
-		}
-	}
-
-	return to_boot->kernel_node;
+	return to_boot;
 }
diff --git a/src/lib/fit_payload.c b/src/lib/fit_payload.c
new file mode 100644
index 0000000000..ada22e8747
--- /dev/null
+++ b/src/lib/fit_payload.c
@@ -0,0 +1,181 @@
+/*
+ * This file is part of the coreboot project.
+ *
+ * Copyright (C) 2003-2004 Eric Biederman
+ * Copyright (C) 2005-2010 coresystems GmbH
+ * Copyright (C) 2014 Google Inc.
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; version 2 of the License.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ */
+
+#include <console/console.h>
+#include <bootmem.h>
+#include <cbmem.h>
+#include <device/resource.h>
+#include <stdlib.h>
+#include <commonlib/region.h>
+#include <fit.h>
+#include <program_loading.h>
+#include <timestamp.h>
+#include <cbfs.h>
+#include <string.h>
+#include <commonlib/compression.h>
+#include <lib.h>
+#include <fit_payload.h>
+
+/* Pack the device_tree and place it at given position. */
+static void pack_fdt(struct region *fdt, struct device_tree *dt)
+{
+	printk(BIOS_INFO, "FIT: Flattening FDT to %p\n",
+	       (void *)fdt->offset);
+
+	dt_flatten(dt, (void *)fdt->offset);
+	prog_segment_loaded(fdt->offset, fdt->size, 0);
+}
+
+/**
+ * Extract a node to given regions.
+ * Returns true on error, false on success.
+ */
+static bool extract(struct region *region, struct fit_image_node *node)
+{
+	void *dst = (void *)region->offset;
+	const char *comp_name;
+	size_t true_size = 0;
+
+	switch (node->compression) {
+	case CBFS_COMPRESS_NONE:
+		comp_name = "Relocating uncompressed";
+		break;
+	case CBFS_COMPRESS_LZMA:
+		comp_name = "Decompressing LZMA";
+		break;
+	case CBFS_COMPRESS_LZ4:
+		comp_name = "Decompressing LZ4";
+		break;
+	default:
+		printk(BIOS_ERR, "ERROR: Unsupported compression\n");
+		return true;
+	}
+
+	printk(BIOS_INFO, "FIT: %s %s to %p\n", comp_name, node->name, dst);
+
+	switch (node->compression) {
+	case CBFS_COMPRESS_NONE:
+		memcpy(dst, node->data, node->size);
+		true_size = node->size;
+		break;
+	case CBFS_COMPRESS_LZMA:
+		timestamp_add_now(TS_START_ULZMA);
+		true_size = ulzman(node->data, node->size, dst, region->size);
+		timestamp_add_now(TS_END_ULZMA);
+		break;
+	case CBFS_COMPRESS_LZ4:
+		timestamp_add_now(TS_START_ULZ4F);
+		true_size = ulz4fn(node->data, node->size, dst, region->size);
+		timestamp_add_now(TS_END_ULZ4F);
+		break;
+	default:
+		return true;
+	}
+
+	if (!true_size) {
+		printk(BIOS_ERR, "ERROR: %s node failed!\n", comp_name);
+		return true;
+	}
+
+	prog_segment_loaded(region->offset, true_size, 0);
+
+	return false;
+}
+
+/*
+ * Parse the uImage FIT, choose a configuration and extract images.
+ */
+void fit_payload(struct prog *payload)
+{
+	struct device_tree *dt = NULL;
+	struct region kernel = {0}, fdt = {0}, initrd = {0};
+	void *data;
+
+	data = rdev_mmap_full(prog_rdev(payload));
+
+	if (data == NULL)
+		return;
+
+	printk(BIOS_INFO, "FIT: Examine payload %s\n", payload->name);
+
+	struct fit_config_node *config = fit_load(data);
+
+	if (!config || !config->kernel_node) {
+		printk(BIOS_ERR, "ERROR: Could not load FIT\n");
+		rdev_munmap(prog_rdev(payload), data);
+		return;
+	}
+
+	if (config->fdt_node) {
+		dt = fdt_unflatten(config->fdt_node->data);
+		if (!dt) {
+			printk(BIOS_ERR,
+			       "ERROR: Failed to unflatten the FDT.\n");
+			rdev_munmap(prog_rdev(payload), data);
+			return;
+		}
+
+		dt_apply_fixups(dt);
+
+		/* Update device_tree */
+#if defined(CONFIG_LINUX_COMMAND_LINE)
+		fit_update_chosen(dt, (char *)CONFIG_LINUX_COMMAND_LINE);
+#endif
+		fit_update_memory(dt);
+	}
+
+	/* Collect infos for fit_payload_arch */
+	kernel.size = config->kernel_node->size;
+	fdt.size = dt ? dt_flat_size(dt) : 0;
+	initrd.size = config->ramdisk_node ? config->ramdisk_node->size : 0;
+
+	/* Invoke arch specific payload placement and fixups */
+	if (!fit_payload_arch(payload, config, &kernel, &fdt, &initrd)) {
+		printk(BIOS_ERR, "ERROR: Failed to find free memory region\n");
+		bootmem_dump_ranges();
+		rdev_munmap(prog_rdev(payload), data);
+		return;
+	}
+
+	/* Load the images to given position */
+	if (config->fdt_node) {
+		/* Update device_tree */
+		if (config->ramdisk_node)
+			fit_add_ramdisk(dt, (void *)initrd.offset, initrd.size);
+
+		pack_fdt(&fdt, dt);
+	}
+
+	if (config->ramdisk_node &&
+	    extract(&initrd, config->ramdisk_node)) {
+		printk(BIOS_ERR, "ERROR: Failed to extract initrd\n");
+		rdev_munmap(prog_rdev(payload), data);
+		return;
+	}
+
+	timestamp_add_now(TS_KERNEL_DECOMPRESSION);
+
+	if (extract(&kernel, config->kernel_node)) {
+		printk(BIOS_ERR, "ERROR: Failed to extract kernel\n");
+		rdev_munmap(prog_rdev(payload), data);
+		return;
+	}
+
+	timestamp_add_now(TS_START_KERNEL);
+
+	rdev_munmap(prog_rdev(payload), data);
+}
diff --git a/src/lib/prog_loaders.c b/src/lib/prog_loaders.c
index 6811eb0001..02b6590aac 100644
--- a/src/lib/prog_loaders.c
+++ b/src/lib/prog_loaders.c
@@ -31,6 +31,7 @@
 #include <symbols.h>
 #include <timestamp.h>
 #include <cbfs.h>
+#include <fit_payload.h>
 
 /* Only can represent up to 1 byte less than size_t. */
 const struct mem_region_device addrspace_32bit =
@@ -183,7 +184,19 @@ void payload_load(void)
 
 	mirror_payload(payload);
 
-	selfload(payload, true);
+	switch (prog_cbfs_type(payload)) {
+	case CBFS_TYPE_SELF: /* Simple ELF */
+		selfload(payload, true);
+		break;
+	case CBFS_TYPE_FIT: /* Flattened image tree */
+		if (IS_ENABLED(CONFIG_PAYLOAD_FIT_SUPPORT)) {
+			fit_payload(payload);
+			break;
+		} /* else fall-through */
+	default:
+		die("Unsupported payload type.\n");
+		break;
+	}
 
 out:
 	if (prog_entry(payload) == NULL)