5 files changed, 143 insertions, 229 deletions

diff --git a/util/cbfstool/cbfs-mkstage.c b/util/cbfstool/cbfs-mkstage.c
index bd1cf54b6b..68eb641bfc 100644
--- a/util/cbfstool/cbfs-mkstage.c
+++ b/util/cbfstool/cbfs-mkstage.c
@@ -10,8 +10,6 @@
 #include "cbfs.h"
 #include "rmodule.h"
 
-#include <commonlib/bsd/compression.h>
-
 /* Checks if program segment contains the ignored section */
 static int is_phdr_ignored(Elf64_Phdr *phdr, Elf64_Shdr *shdr)
 {
@@ -73,19 +71,20 @@ static Elf64_Shdr *find_ignored_section_header(struct parsed_elf *pelf,
 	return NULL;
 }
 
-static void fill_cbfs_stage(struct buffer *outheader, enum cbfs_compression algo,
-				uint64_t entry, uint64_t loadaddr,
-				uint32_t filesize, uint32_t memsize)
+static int fill_cbfs_stageheader(struct cbfs_file_attr_stageheader *stageheader,
+				 uint64_t entry, uint64_t loadaddr,
+				 uint32_t memsize)
 {
-	/* N.B. The original plan was that SELF data was B.E.
-	 * but: this is all L.E.
-	 * Maybe we should just change the spec.
-	 */
-	xdr_le.put32(outheader, algo);
-	xdr_le.put64(outheader, entry);
-	xdr_le.put64(outheader, loadaddr);
-	xdr_le.put32(outheader, filesize);
-	xdr_le.put32(outheader, memsize);
+	if (entry - loadaddr >= memsize) {
+		ERROR("stage entry point out of bounds!\n");
+		return -1;
+	}
+
+	stageheader->loadaddr = htonll(loadaddr);
+	stageheader->memlen = htonl(memsize);
+	stageheader->entry_offset = htonl(entry - loadaddr);
+
+	return  0;
 }
 
 /* returns size of result, or -1 if error.
@@ -93,25 +92,20 @@ static void fill_cbfs_stage(struct buffer *outheader, enum cbfs_compression algo
  * works for all elf files, not just the restricted set.
  */
 int parse_elf_to_stage(const struct buffer *input, struct buffer *output,
-		       enum cbfs_compression algo, const char *ignore_section)
+		       const char *ignore_section,
+		       struct cbfs_file_attr_stageheader *stageheader)
 {
 	struct parsed_elf pelf;
 	Elf64_Phdr *phdr;
 	Elf64_Ehdr *ehdr;
 	Elf64_Shdr *shdr_ignored;
 	Elf64_Addr virt_to_phys;
-	char *buffer;
-	struct buffer outheader;
 	int ret = -1;
 
 	int headers;
-	int i, outlen;
+	int i;
 	uint64_t data_start, data_end, mem_end;
 
-	comp_func_ptr compress = compression_function(algo);
-	if (!compress)
-		return -1;
-
 	int flags = ELF_PARSE_PHDR | ELF_PARSE_SHDR | ELF_PARSE_STRTAB;
 
 	if (parse_elf(input, &pelf, flags)) {
@@ -178,15 +172,13 @@ int parse_elf_to_stage(const struct buffer *input, struct buffer *output,
 		exit(1);
 	}
 
-	/* allocate an intermediate buffer for the data */
-	buffer = calloc(data_end - data_start, 1);
-
-	if (buffer == NULL) {
+	if (buffer_create(output, data_end - data_start, input->name) != 0) {
 		ERROR("Unable to allocate memory: %m\n");
 		goto err;
 	}
+	memset(output->data, 0, output->size);
 
-	/* Copy the file data into the buffer */
+	/* Copy the file data into the output buffer */
 
 	for (i = 0; i < headers; i++) {
 		if (phdr[i].p_type != PT_LOAD)
@@ -207,90 +199,17 @@ int parse_elf_to_stage(const struct buffer *input, struct buffer *output,
 			ERROR("Underflow copying out the segment."
 			      "File has %zu bytes left, segment end is %zu\n",
 			      input->size, (size_t)(phdr[i].p_offset + phdr[i].p_filesz));
-			free(buffer);
 			goto err;
 		}
-		memcpy(buffer + (phdr[i].p_paddr - data_start),
+		memcpy(&output->data[phdr[i].p_paddr - data_start],
 		       &input->data[phdr[i].p_offset],
 		       phdr[i].p_filesz);
 	}
 
-	/* Now make the output buffer */
-	if (buffer_create(output, sizeof(struct cbfs_stage) + data_end - data_start,
-			  input->name) != 0) {
-		ERROR("Unable to allocate memory: %m\n");
-		free(buffer);
-		goto err;
-	}
-	memset(output->data, 0, output->size);
-
-	/* Compress the data, at which point we'll know information
-	 * to fill out the header. This seems backward but it works because
-	 * - the output header is a known size (not always true in many xdr's)
-	 * - we do need to know the compressed output size first
-	 * If compression fails or makes the data bigger, we'll warn about it
-	 * and use the original data.
-	 */
-	if (compress(buffer, data_end - data_start,
-		     (output->data + sizeof(struct cbfs_stage)),
-		     &outlen) < 0 || (unsigned)outlen > data_end - data_start) {
-		WARN("Compression failed or would make the data bigger "
-		     "- disabled.\n");
-		memcpy(output->data + sizeof(struct cbfs_stage),
-		       buffer, data_end - data_start);
-		outlen = data_end - data_start;
-		algo = CBFS_COMPRESS_NONE;
-	}
-
-	/* Check for enough BSS scratch space to decompress LZ4 in-place. */
-	if (algo == CBFS_COMPRESS_LZ4) {
-		size_t result;
-		size_t memlen = mem_end - data_start;
-		size_t compressed_size = outlen;
-		char *compare_buffer = malloc(memlen);
-		char *start = compare_buffer + memlen - compressed_size;
-
-		if (compare_buffer == NULL) {
-			ERROR("Can't allocate memory!\n");
-			free(buffer);
-			goto err;
-		}
-
-		memcpy(start, output->data + sizeof(struct cbfs_stage),
-		       compressed_size);
-		result = ulz4fn(start, compressed_size, compare_buffer, memlen);
-
-		if (result == 0) {
-			ERROR("Not enough scratch space to decompress LZ4 in-place -- increase BSS size or disable compression!\n");
-			free(compare_buffer);
-			free(buffer);
-			goto err;
-		}
-		if (result != data_end - data_start ||
-		    memcmp(compare_buffer, buffer, data_end - data_start)) {
-			ERROR("LZ4 compression BUG! Report to mailing list.\n");
-			free(compare_buffer);
-			free(buffer);
-			goto err;
-		}
-		free(compare_buffer);
-	}
-
-	free(buffer);
-
-	/* Set up for output marshaling. */
-	outheader.data = output->data;
-	outheader.size = 0;
-
 	/* coreboot expects entry point to be physical address. Thus, adjust the
-	 * entry point accordingly.
-	 */
-	fill_cbfs_stage(&outheader, algo, ehdr->e_entry + virt_to_phys,
-			data_start, outlen, mem_end - data_start);
-
-	output->size = sizeof(struct cbfs_stage) + outlen;
-	ret = 0;
-
+	   entry point accordingly. */
+	ret = fill_cbfs_stageheader(stageheader, ehdr->e_entry + virt_to_phys,
+				    data_start, mem_end - data_start);
 err:
 	parsed_elf_destroy(&pelf);
 	return ret;
@@ -341,13 +260,13 @@ static int rmod_filter(struct reloc_filter *f, const Elf64_Rela *r)
 }
 
 int parse_elf_to_xip_stage(const struct buffer *input, struct buffer *output,
-				uint32_t *location, const char *ignore_section)
+			   uint32_t *location, const char *ignore_section,
+			   struct cbfs_file_attr_stageheader *stageheader)
 {
 	struct xip_context xipctx;
 	struct rmod_context *rmodctx;
 	struct reloc_filter filter;
 	struct parsed_elf *pelf;
-	size_t output_sz;
 	uint32_t adjustment;
 	struct buffer binput;
 	struct buffer boutput;
@@ -381,13 +300,12 @@ int parse_elf_to_xip_stage(const struct buffer *input, struct buffer *output,
 	if (rmodule_collect_relocations(rmodctx, &filter))
 		goto out;
 
-	output_sz = sizeof(struct cbfs_stage) + pelf->phdr->p_filesz;
-	if (buffer_create(output, output_sz, input->name) != 0) {
+	if (buffer_create(output, pelf->phdr->p_filesz, input->name) != 0) {
 		ERROR("Unable to allocate memory: %m\n");
 		goto out;
 	}
 	buffer_clone(&boutput, output);
-	memset(buffer_get(&boutput), 0, output_sz);
+	memset(buffer_get(&boutput), 0, pelf->phdr->p_filesz);
 	buffer_set_size(&boutput, 0);
 
 	/* Single loadable segment. The entire segment moves to final
@@ -395,17 +313,16 @@ int parse_elf_to_xip_stage(const struct buffer *input, struct buffer *output,
 	adjustment = *location - pelf->phdr->p_vaddr;
 	DEBUG("Relocation adjustment: %08x\n", adjustment);
 
-	fill_cbfs_stage(&boutput, CBFS_COMPRESS_NONE,
-			(uint32_t)pelf->ehdr.e_entry + adjustment,
-			(uint32_t)pelf->phdr->p_vaddr + adjustment,
-			pelf->phdr->p_filesz, pelf->phdr->p_memsz);
+	fill_cbfs_stageheader(stageheader,
+			      (uint32_t)pelf->ehdr.e_entry + adjustment,
+			      (uint32_t)pelf->phdr->p_vaddr + adjustment,
+			      pelf->phdr->p_memsz);
 	/* Need an adjustable buffer. */
 	buffer_clone(&binput, input);
 	buffer_seek(&binput, pelf->phdr->p_offset);
 	bputs(&boutput, buffer_get(&binput), pelf->phdr->p_filesz);
 
 	buffer_clone(&boutput, output);
-	buffer_seek(&boutput, sizeof(struct cbfs_stage));
 
 	/* Make adjustments to all the relocations within the program. */
 	for (i = 0; i < rmodctx->nrelocs; i++) {
@@ -431,8 +348,6 @@ int parse_elf_to_xip_stage(const struct buffer *input, struct buffer *output,
 		xdr_le.put32(&out, val + adjustment);
 	}
 
-	/* Need to back up the location to include cbfs stage metadata. */
-	*location -= sizeof(struct cbfs_stage);
 	ret = 0;
 
 out:
diff --git a/util/cbfstool/cbfs_image.c b/util/cbfstool/cbfs_image.c
index d3c6c94d48..1fb19bacd6 100644
--- a/util/cbfstool/cbfs_image.c
+++ b/util/cbfstool/cbfs_image.c
@@ -824,69 +824,6 @@ struct cbfs_file *cbfs_get_entry(struct cbfs_image *image, const char *name)
 	return NULL;
 }
 
-static int cbfs_stage_decompress(struct cbfs_stage *stage, struct buffer *buff)
-{
-	struct buffer reader;
-	char *orig_buffer;
-	char *new_buffer;
-	size_t new_buff_sz;
-	decomp_func_ptr decompress;
-
-	buffer_clone(&reader, buff);
-
-	/* The stage metadata is in little endian. */
-	stage->compression = xdr_le.get32(&reader);
-	stage->entry = xdr_le.get64(&reader);
-	stage->load = xdr_le.get64(&reader);
-	stage->len = xdr_le.get32(&reader);
-	stage->memlen = xdr_le.get32(&reader);
-
-	/* Create a buffer just with the uncompressed program now that the
-	 * struct cbfs_stage has been peeled off. */
-	if (stage->compression == CBFS_COMPRESS_NONE) {
-		new_buff_sz = buffer_size(buff) - sizeof(struct cbfs_stage);
-
-		orig_buffer = buffer_get(buff);
-		new_buffer = calloc(1, new_buff_sz);
-		memcpy(new_buffer, orig_buffer + sizeof(struct cbfs_stage),
-			new_buff_sz);
-		buffer_init(buff, buff->name, new_buffer, new_buff_sz);
-		free(orig_buffer);
-		return 0;
-	}
-
-	decompress = decompression_function(stage->compression);
-	if (decompress == NULL)
-		return -1;
-
-	orig_buffer = buffer_get(buff);
-
-	/* This can be too big of a buffer needed, but there's no current
-	 * field indicating decompressed size of data. */
-	new_buff_sz = stage->memlen;
-	new_buffer = calloc(1, new_buff_sz);
-
-	if (decompress(orig_buffer + sizeof(struct cbfs_stage),
-			(int)(buffer_size(buff) - sizeof(struct cbfs_stage)),
-			new_buffer, (int)new_buff_sz, &new_buff_sz)) {
-		ERROR("Couldn't decompress stage.\n");
-		free(new_buffer);
-		return -1;
-	}
-
-	/* Include correct size for full stage info. */
-	buffer_init(buff, buff->name, new_buffer, new_buff_sz);
-
-	/* True decompressed size is just the data size -- no metadata. */
-	stage->len = new_buff_sz;
-	/* Stage is not compressed. */
-	stage->compression = CBFS_COMPRESS_NONE;
-
-	free(orig_buffer);
-
-	return 0;
-}
-
 static int cbfs_payload_decompress(struct cbfs_payload_segment *segments,
 		struct buffer *buff, int num_seg)
 {
@@ -1020,11 +957,11 @@ static int init_elf_from_arch(Elf64_Ehdr *ehdr, uint32_t cbfs_arch)
 	return 0;
 }
 
-static int cbfs_stage_make_elf(struct buffer *buff, uint32_t arch)
+static int cbfs_stage_make_elf(struct buffer *buff, uint32_t arch,
+			       struct cbfs_file *entry)
 {
 	Elf64_Ehdr ehdr;
 	Elf64_Shdr shdr;
-	struct cbfs_stage stage;
 	struct elf_writer *ew;
 	struct buffer elf_out;
 	size_t empty_sz;
@@ -1035,16 +972,23 @@ static int cbfs_stage_make_elf(struct buffer *buff, uint32_t arch)
 		return -1;
 	}
 
-	if (cbfs_stage_decompress(&stage, buff)) {
-		ERROR("Failed to decompress stage.\n");
+	struct cbfs_file_attr_stageheader *stage = NULL;
+	for (struct cbfs_file_attribute *attr = cbfs_file_first_attr(entry);
+	     attr != NULL; attr = cbfs_file_next_attr(entry, attr)) {
+		if (ntohl(attr->tag) == CBFS_FILE_ATTR_TAG_STAGEHEADER) {
+			stage = (struct cbfs_file_attr_stageheader *)attr;
+			break;
+		}
+	}
+
+	if (stage == NULL) {
+		ERROR("Stage header not found for %s\n", entry->filename);
 		return -1;
 	}
 
 	if (init_elf_from_arch(&ehdr, arch))
 		return -1;
 
-	ehdr.e_entry = stage.entry;
-
 	/* Attempt rmodule translation first. */
 	rmod_ret = rmodule_stage_to_elf(&ehdr, buff);
 
@@ -1056,6 +1000,8 @@ static int cbfs_stage_make_elf(struct buffer *buff, uint32_t arch)
 
 	/* Rmodule couldn't do anything with the data. Continue on with SELF. */
 
+	ehdr.e_entry = ntohll(stage->loadaddr) + ntohl(stage->entry_offset);
+
 	ew = elf_writer_init(&ehdr);
 	if (ew == NULL) {
 		ERROR("Unable to init ELF writer.\n");
@@ -1065,9 +1011,9 @@ static int cbfs_stage_make_elf(struct buffer *buff, uint32_t arch)
 	memset(&shdr, 0, sizeof(shdr));
 	shdr.sh_type = SHT_PROGBITS;
 	shdr.sh_flags = SHF_WRITE | SHF_ALLOC | SHF_EXECINSTR;
-	shdr.sh_addr = stage.load;
-	shdr.sh_size = stage.len;
-	empty_sz = stage.memlen - stage.len;
+	shdr.sh_addr = ntohll(stage->loadaddr);
+	shdr.sh_size = buffer_size(buff);
+	empty_sz = ntohl(stage->memlen) - buffer_size(buff);
 
 	if (elf_writer_add_section(ew, &shdr, buff, ".program")) {
 		ERROR("Unable to add ELF section: .program\n");
@@ -1082,7 +1028,7 @@ static int cbfs_stage_make_elf(struct buffer *buff, uint32_t arch)
 		memset(&shdr, 0, sizeof(shdr));
 		shdr.sh_type = SHT_NOBITS;
 		shdr.sh_flags = SHF_WRITE | SHF_ALLOC;
-		shdr.sh_addr = stage.load + stage.len;
+		shdr.sh_addr = ntohl(stage->loadaddr) + buffer_size(buff);
 		shdr.sh_size = empty_sz;
 		if (elf_writer_add_section(ew, &shdr, &b, ".empty")) {
 			ERROR("Unable to add ELF section: .empty\n");
@@ -1106,7 +1052,8 @@ static int cbfs_stage_make_elf(struct buffer *buff, uint32_t arch)
 	return 0;
 }
 
-static int cbfs_payload_make_elf(struct buffer *buff, uint32_t arch)
+static int cbfs_payload_make_elf(struct buffer *buff, uint32_t arch,
+				 unused struct cbfs_file *entry)
 {
 	Elf64_Ehdr ehdr;
 	Elf64_Shdr shdr;
@@ -1258,7 +1205,7 @@ static int cbfs_payload_make_elf(struct buffer *buff, uint32_t arch)
 	}
 
 	if (elf_writer_serialize(ew, &elf_out)) {
-		ERROR("Unable to create ELF file from stage.\n");
+		ERROR("Unable to create ELF file from payload.\n");
 		goto out;
 	}
 
@@ -1320,13 +1267,13 @@ int cbfs_export_entry(struct cbfs_image *image, const char *entry_name,
 	}
 
 	/*
-	 * The stage metadata is never compressed proper for cbfs_stage
-	 * files. The contents of the stage data can be though. Therefore
-	 * one has to do a second pass for stages to potentially decompress
-	 * the stage data to make it more meaningful.
+	 * We want to export stages and payloads as ELFs, not with coreboot's
+	 * custom stage/SELF binary formats, so we need to do extra processing
+	 * to turn them back into an ELF.
 	 */
 	if (do_processing) {
-		int (*make_elf)(struct buffer *, uint32_t) = NULL;
+		int (*make_elf)(struct buffer *, uint32_t,
+				struct cbfs_file *) = NULL;
 		switch (ntohl(entry->type)) {
 		case CBFS_TYPE_STAGE:
 			make_elf = cbfs_stage_make_elf;
@@ -1335,7 +1282,7 @@ int cbfs_export_entry(struct cbfs_image *image, const char *entry_name,
 			make_elf = cbfs_payload_make_elf;
 			break;
 		}
-		if (make_elf && make_elf(&buffer, arch)) {
+		if (make_elf && make_elf(&buffer, arch, entry)) {
 			ERROR("Failed to write %s into %s.\n",
 			      entry_name, filename);
 			buffer_delete(&buffer);
@@ -1387,17 +1334,29 @@ int cbfs_print_header_info(struct cbfs_image *image)
 	return 0;
 }
 
-static int cbfs_print_stage_info(struct cbfs_stage *stage, FILE* fp)
+static int cbfs_print_stage_info(struct cbfs_file *entry, FILE* fp)
 {
+
+	struct cbfs_file_attr_stageheader *stage = NULL;
+	for (struct cbfs_file_attribute *attr = cbfs_file_first_attr(entry);
+	     attr != NULL; attr = cbfs_file_next_attr(entry, attr)) {
+		if (ntohl(attr->tag) == CBFS_FILE_ATTR_TAG_STAGEHEADER) {
+			stage = (struct cbfs_file_attr_stageheader *)attr;
+			break;
+		}
+	}
+
+	if (stage == NULL) {
+		fprintf(fp, "    ERROR: stage header not found!\n");
+		return -1;
+	}
+
 	fprintf(fp,
-		"    %s compression, entry: 0x%" PRIx64 ", load: 0x%" PRIx64 ", "
-		"length: %d/%d\n",
-		lookup_name_by_type(types_cbfs_compression,
-				    stage->compression, "(unknown)"),
-		stage->entry,
-		stage->load,
-		stage->len,
-		stage->memlen);
+		"    entry: 0x%" PRIx64 ", load: 0x%" PRIx64 ", "
+		"memlen: %d\n",
+		ntohll(stage->loadaddr) + ntohl(stage->entry_offset),
+		ntohll(stage->loadaddr),
+		ntohl(stage->memlen));
 	return 0;
 }
 
@@ -1519,8 +1478,7 @@ int cbfs_print_entry_info(struct cbfs_image *image, struct cbfs_file *entry,
 	/* note the components of the subheader may be in host order ... */
 	switch (ntohl(entry->type)) {
 		case CBFS_TYPE_STAGE:
-			cbfs_print_stage_info((struct cbfs_stage *)
-					      CBFS_SUBHEADER(entry), fp);
+			cbfs_print_stage_info(entry, fp);
 			break;
 
 		case CBFS_TYPE_SELF:
diff --git a/util/cbfstool/cbfstool.c b/util/cbfstool/cbfstool.c
index 3e80712ca2..6133536b4b 100644
--- a/util/cbfstool/cbfstool.c
+++ b/util/cbfstool/cbfstool.c
@@ -14,7 +14,9 @@
 #include "cbfs_sections.h"
 #include "elfparsing.h"
 #include "partitioned_file.h"
+#include "lz4/lib/xxhash.h"
 #include <commonlib/bsd/cbfs_private.h>
+#include <commonlib/bsd/compression.h>
 #include <commonlib/bsd/metadata_hash.h>
 #include <commonlib/fsp.h>
 #include <commonlib/endian.h>
@@ -911,16 +913,7 @@ static int cbfs_add_component(const char *filename,
 					sizeof(struct cbfs_file_attr_position));
 			if (attrs == NULL)
 				goto error;
-			/* If we add a stage or a payload, we need to take  */
-			/* care about the additional metadata that is added */
-			/* to the cbfs file and therefore set the position  */
-			/* the real beginning of the data. */
-			if (type == CBFS_TYPE_STAGE)
-				attrs->position = htonl(offset - sizeof(struct cbfs_stage));
-			else if (type == CBFS_TYPE_SELF)
-				attrs->position = htonl(offset - sizeof(struct cbfs_payload));
-			else
-				attrs->position = htonl(offset);
+			attrs->position = htonl(offset);
 		}
 		/* Add alignment attribute if used */
 		if (param.alignment) {
@@ -1118,11 +1111,18 @@ static int cbfstool_convert_mkstage(struct buffer *buffer, uint32_t *offset,
 	 * stages that would actually fit once compressed.
 	 */
 	if ((param.alignment || param.stage_xip) &&
-	     do_cbfs_locate(offset, sizeof(struct cbfs_stage), data_size))  {
+	     do_cbfs_locate(offset, sizeof(struct cbfs_file_attr_stageheader),
+			    data_size))  {
 		ERROR("Could not find location for stage.\n");
 		return 1;
 	}
 
+	struct cbfs_file_attr_stageheader *stageheader = (void *)
+		cbfs_add_file_attr(header, CBFS_FILE_ATTR_TAG_STAGEHEADER,
+				   sizeof(struct cbfs_file_attr_stageheader));
+	if (!stageheader)
+		return -1;
+
 	if (param.stage_xip) {
 		/*
 		 * Ensure the address is a memory mapped one. This assumes
@@ -1132,19 +1132,57 @@ static int cbfstool_convert_mkstage(struct buffer *buffer, uint32_t *offset,
 		*offset = convert_addr_space(param.image_region, *offset);
 
 		ret = parse_elf_to_xip_stage(buffer, &output, offset,
-						param.ignore_section);
+					     param.ignore_section,
+					     stageheader);
 	} else {
-		ret = parse_elf_to_stage(buffer, &output, param.compression,
-					 param.ignore_section);
+		ret = parse_elf_to_stage(buffer, &output, param.ignore_section,
+					 stageheader);
 	}
-
 	if (ret != 0)
 		return -1;
+
+	/* Store a hash of original uncompressed stage to compare later. */
+	size_t decmp_size = buffer_size(&output);
+	uint32_t decmp_hash = XXH32(buffer_get(&output), decmp_size, 0);
+
+	/* Chain to base conversion routine to handle compression. */
+	ret = cbfstool_convert_raw(&output, offset, header);
+	if (ret != 0)
+		goto fail;
+
+	/* Special care must be taken for LZ4-compressed stages that the BSS is
+	   large enough to provide scratch space for in-place decompression. */
+	if (!param.precompression && param.compression == CBFS_COMPRESS_LZ4) {
+		size_t memlen = ntohl(stageheader->memlen);
+		size_t compressed_size = buffer_size(&output);
+		uint8_t *compare_buffer = malloc(memlen);
+		uint8_t *start = compare_buffer + memlen - compressed_size;
+		if (!compare_buffer) {
+			ERROR("Out of memory\n");
+			goto fail;
+		}
+		memcpy(start, buffer_get(&output), compressed_size);
+		ret = ulz4fn(start, compressed_size, compare_buffer, memlen);
+		if  (ret == 0) {
+			ERROR("Not enough scratch space to decompress LZ4 in-place -- increase BSS size or disable compression!\n");
+			free(compare_buffer);
+			goto fail;
+		} else if (ret != (int)decmp_size ||
+			   decmp_hash != XXH32(compare_buffer, decmp_size, 0)) {
+			ERROR("LZ4 compression BUG! Report to mailing list.\n");
+			free(compare_buffer);
+			goto fail;
+		}
+		free(compare_buffer);
+	}
+
 	buffer_delete(buffer);
-	// Direct assign, no dupe.
-	memcpy(buffer, &output, sizeof(*buffer));
-	header->len = htonl(output.size);
+	buffer_clone(buffer, &output);
 	return 0;
+
+fail:
+	buffer_delete(&output);
+	return -1;
 }
 
 static int cbfstool_convert_mkpayload(struct buffer *buffer,
diff --git a/util/cbfstool/common.h b/util/cbfstool/common.h
index db9c7e7297..479fd713dd 100644
--- a/util/cbfstool/common.h
+++ b/util/cbfstool/common.h
@@ -174,10 +174,12 @@ int parse_flat_binary_to_payload(const struct buffer *input,
 				 enum cbfs_compression algo);
 /* cbfs-mkstage.c */
 int parse_elf_to_stage(const struct buffer *input, struct buffer *output,
-		       enum cbfs_compression algo, const char *ignore_section);
+		       const char *ignore_section,
+		       struct cbfs_file_attr_stageheader *stageheader);
 /* location is TOP aligned. */
 int parse_elf_to_xip_stage(const struct buffer *input, struct buffer *output,
-				uint32_t *location, const char *ignore_section);
+			   uint32_t *location, const char *ignore_section,
+			   struct cbfs_file_attr_stageheader *stageheader);
 
 void print_supported_architectures(void);
 void print_supported_filetypes(void);
diff --git a/util/cbfstool/rmodule.c b/util/cbfstool/rmodule.c
index 258a4d8803..4ac2951f72 100644
--- a/util/cbfstool/rmodule.c
+++ b/util/cbfstool/rmodule.c
@@ -498,8 +498,9 @@ write_elf(const struct rmod_context *ctx, const struct buffer *in,
 	/* Program contents. */
 	buffer_splice(&program, in, ctx->phdr->p_offset, ctx->phdr->p_filesz);
 
-	/* Create ELF writer with modified entry point. */
+	/* Create ELF writer. Set entry point to 0 to match section offsets. */
 	memcpy(&ehdr, &ctx->pelf.ehdr, sizeof(ehdr));
+	ehdr.e_entry = 0;
 	ew = elf_writer_init(&ehdr);
 
 	if (ew == NULL) {