5 files changed, 38 insertions, 22 deletions

diff --git a/util/cbfstool/cbfs-mkpayload.c b/util/cbfstool/cbfs-mkpayload.c
index adbe313b18..dabe32262b 100644
--- a/util/cbfstool/cbfs-mkpayload.c
+++ b/util/cbfstool/cbfs-mkpayload.c
@@ -205,8 +205,11 @@ int parse_elf_to_payload(const struct buffer *input,
 		segs[segments].offset = doffset;
 
 		int len;
-		compress((char *)&header[phdr[i].p_offset],
-			 phdr[i].p_filesz, output->data + doffset, &len);
+		if (compress((char *)&header[phdr[i].p_offset],
+			     phdr[i].p_filesz, output->data + doffset, &len)) {
+			buffer_delete(output);
+			return -1;
+		}
 		segs[segments].len = len;
 
 		/* If the compressed section is larger, then use the
@@ -261,7 +264,10 @@ int parse_flat_binary_to_payload(const struct buffer *input,
 	segs[0].mem_len = input->size;
 	segs[0].offset = doffset;
 
-	compress(input->data, input->size, output->data + doffset, &len);
+	if (compress(input->data, input->size, output->data + doffset, &len)) {
+		buffer_delete(output);
+		return -1;
+	}
 	segs[0].compression = algo;
 	segs[0].len = len;
 
@@ -387,7 +393,10 @@ int parse_fv_to_payload(const struct buffer *input,
 	segs[0].mem_len = input->size;
 	segs[0].offset = doffset;
 
-	compress(input->data, input->size, output->data + doffset, &len);
+	if (compress(input->data, input->size, output->data + doffset, &len)) {
+		buffer_delete(output);
+		return -1;
+	}
 	segs[0].compression = algo;
 	segs[0].len = len;
 
diff --git a/util/cbfstool/cbfs-mkstage.c b/util/cbfstool/cbfs-mkstage.c
index d795c7c0e2..8c77ee5e65 100644
--- a/util/cbfstool/cbfs-mkstage.c
+++ b/util/cbfstool/cbfs-mkstage.c
@@ -156,9 +156,12 @@ int parse_elf_to_stage(const struct buffer *input, struct buffer *output,
 	 * - the output header is a known size (not always true in many xdr's)
 	 * - we do need to know the compressed output size first
 	 */
-	compress(buffer, data_end - data_start,
-		 (output->data + sizeof(struct cbfs_stage)),
-		 &outlen);
+	if (compress(buffer, data_end - data_start,
+		     (output->data + sizeof(struct cbfs_stage)),
+		     &outlen) < 0) {
+		free(buffer);
+		return -1;
+	}
 	free(buffer);
 
 	/* Set up for output marshaling. */
diff --git a/util/cbfstool/common.h b/util/cbfstool/common.h
index 455fb8b283..02a088b98d 100644
--- a/util/cbfstool/common.h
+++ b/util/cbfstool/common.h
@@ -117,7 +117,7 @@ uint32_t string_to_arch(const char *arch_string);
 
 #define ALIGN(val, by) (((val) + (by)-1)&~((by)-1))
 
-typedef void (*comp_func_ptr) (char *, int, char *, int *);
+typedef int (*comp_func_ptr) (char *, int, char *, int *);
 typedef enum { CBFS_COMPRESS_NONE = 0, CBFS_COMPRESS_LZMA = 1 } comp_algo;
 
 comp_func_ptr compression_function(comp_algo algo);
@@ -145,8 +145,8 @@ void print_supported_filetypes(void);
 
 #define ARRAY_SIZE(a) (int)(sizeof(a) / sizeof((a)[0]))
 /* lzma/lzma.c */
-void do_lzma_compress(char *in, int in_len, char *out, int *out_len);
-void do_lzma_uncompress(char *dst, int dst_len, char *src, int src_len);
+int do_lzma_compress(char *in, int in_len, char *out, int *out_len);
+int do_lzma_uncompress(char *dst, int dst_len, char *src, int src_len);
 /* xdr.c */
 struct xdr {
 	uint8_t (*get8)(struct buffer *input);
diff --git a/util/cbfstool/compress.c b/util/cbfstool/compress.c
index 38fa03dbed..5b916cee11 100644
--- a/util/cbfstool/compress.c
+++ b/util/cbfstool/compress.c
@@ -26,15 +26,16 @@
 #include <stdio.h>
 #include "common.h"
 
-static void lzma_compress(char *in, int in_len, char *out, int *out_len)
+static int lzma_compress(char *in, int in_len, char *out, int *out_len)
 {
-	do_lzma_compress(in, in_len, out, out_len);
+	return do_lzma_compress(in, in_len, out, out_len);
 }
 
-static void none_compress(char *in, int in_len, char *out, int *out_len)
+static int none_compress(char *in, int in_len, char *out, int *out_len)
 {
 	memcpy(out, in, in_len);
 	*out_len = in_len;
+	return 0;
 }
 
 comp_func_ptr compression_function(comp_algo algo)
diff --git a/util/cbfstool/lzma/lzma.c b/util/cbfstool/lzma/lzma.c
index 7a6081557a..abd1bfd9b5 100644
--- a/util/cbfstool/lzma/lzma.c
+++ b/util/cbfstool/lzma/lzma.c
@@ -83,11 +83,11 @@ static struct ISeqOutStream os = { Write };
  * @param out_len a pointer to the compressed length of in
  */
 
-void do_lzma_compress(char *in, int in_len, char *out, int *out_len)
+int do_lzma_compress(char *in, int in_len, char *out, int *out_len)
 {
 	if (in_len == 0) {
 		ERROR("LZMA: Input length is zero.\n");
-		return;
+		return -1;
 	}
 
 	struct CLzmaEncProps props;
@@ -119,7 +119,7 @@ void do_lzma_compress(char *in, int in_len, char *out, int *out_len)
 	int res = LzmaEnc_SetProps(p, &props);
 	if (res != SZ_OK) {
 		ERROR("LZMA: LzmaEnc_SetProps failed.\n");
-		return;
+		return -1;
 	}
 
 	unsigned char propsEncoded[LZMA_PROPS_SIZE + 8];
@@ -127,7 +127,7 @@ void do_lzma_compress(char *in, int in_len, char *out, int *out_len)
 	res = LzmaEnc_WriteProperties(p, propsEncoded, &propsSize);
 	if (res != SZ_OK) {
 		ERROR("LZMA: LzmaEnc_WriteProperties failed.\n");
-		return;
+		return -1;
 	}
 
 	instream.p = in;
@@ -144,17 +144,18 @@ void do_lzma_compress(char *in, int in_len, char *out, int *out_len)
 	res = LzmaEnc_Encode(p, &os, &is, 0, &LZMAalloc, &LZMAalloc);
 	if (res != SZ_OK) {
 		ERROR("LZMA: LzmaEnc_Encode failed %d.\n", res);
-		return;
+		return -1;
 	}
 
 	*out_len = outstream.pos;
+	return 0;
 }
 
-void do_lzma_uncompress(char *dst, int dst_len, char *src, int src_len)
+int do_lzma_uncompress(char *dst, int dst_len, char *src, int src_len)
 {
 	if (src_len <= LZMA_PROPS_SIZE + 8) {
 		ERROR("LZMA: Input length is too small.\n");
-		return;
+		return -1;
 	}
 
 	uint64_t out_sizemax = get_64(&src[LZMA_PROPS_SIZE]);
@@ -162,7 +163,7 @@ void do_lzma_uncompress(char *dst, int dst_len, char *src, int src_len)
 	if (out_sizemax > (size_t) dst_len) {
 		ERROR("Not copying %d bytes to %d-byte buffer!\n",
 			(unsigned int)out_sizemax, dst_len);
-		return;
+		return -1;
 	}
 
 	enum ELzmaStatus status;
@@ -179,6 +180,8 @@ void do_lzma_uncompress(char *dst, int dst_len, char *src, int src_len)
 
 	if (res != SZ_OK) {
 		ERROR("Error while decompressing.\n");
-		return;
+		return -1;
 	}
+
+	return 0;
 }