Add memory leak checking

Add checks that the vboot library does not leak memory. This works by
tracking VbExMalloc() calls and making sure that they have an associated
VbExFree().

Adjust host_signature to use VbExFree() instead of free(), so that this
scheme works correctly for existing code.

BUG=chrome-os-partner:21115
BRANCH=pit
TEST=FEATURES=test emerge-peach_pit vboot_reference

Change-Id: I6ccccfbcc162fc43fb75862cd0eddad78ce8b18a
Signed-off-by: Simon Glass <sjg@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/66175
diff --git a/firmware/stub/vboot_api_stub_sf.c b/firmware/stub/vboot_api_stub_sf.c
index e2f1c1a..3c5ec85 100644
--- a/firmware/stub/vboot_api_stub_sf.c
+++ b/firmware/stub/vboot_api_stub_sf.c
@@ -5,6 +5,7 @@
  * Stub implementations of firmware-provided API functions.
  */
 
+#include <execinfo.h>
 #include <stdint.h>
 
 #define _STUB_IMPLEMENTATION_
@@ -17,18 +18,81 @@
 
 #include "vboot_api.h"
 
+#define MAX_STACK_LEVELS 10
+
+
+/* Keep track of nodes that are currently allocated */
+struct alloc_node {
+	struct alloc_node *next;
+	void *ptr;
+	size_t size;
+	void *bt_buffer[MAX_STACK_LEVELS];
+	int bt_levels;
+};
+
+static struct alloc_node *alloc_head;
+
+static void print_stacktrace(void)
+{
+	void *buffer[MAX_STACK_LEVELS];
+	int levels = backtrace(buffer, MAX_STACK_LEVELS);
+
+	// print to stderr (fd = 2), and remove this function from the trace
+	backtrace_symbols_fd(buffer + 1, levels - 1, 2);
+}
+
 void *VbExMalloc(size_t size)
 {
+	struct alloc_node *node;
 	void *p = malloc(size);
+
 	if (!p) {
 		/* Fatal Error. We must abort. */
 		abort();
 	}
+
+	node = malloc(sizeof(*node));
+	if (!node)
+		abort();
+	node->next = alloc_head;
+	node->ptr = p;
+	node->size = size;
+	node->bt_levels = backtrace(node->bt_buffer, MAX_STACK_LEVELS);
+	alloc_head = node;
+
 	return p;
 }
 
+static struct alloc_node **find_node(void *ptr)
+{
+	struct alloc_node **nodep;
+
+	for (nodep = &alloc_head; *nodep; nodep = &(*nodep)->next)
+		if ((*nodep)->ptr == ptr)
+			return nodep;
+
+	return NULL;
+}
+
 void VbExFree(void *ptr)
 {
+	struct alloc_node **nodep, *next;
+
+	nodep = find_node(ptr);
+	if (nodep) {
+		next = (*nodep)->next;
+		free(*nodep);
+		*nodep = next;
+	} else {
+		fprintf(stderr, "\n>>>>>> Invalid VbExFree() %p\n", ptr);
+		fflush(stderr);
+		print_stacktrace();
+		/*
+		 * Fall through and do the free() so we get normal error
+		 * handling.
+		 */
+	}
+
 	free(ptr);
 }
 
@@ -37,3 +101,27 @@
 {
 	return VBERROR_SUCCESS;
 }
+
+int vboot_api_stub_check_memory(void)
+{
+	struct alloc_node *node, *next;
+
+	if (!alloc_head)
+		return 0;
+
+	/*
+	 * Make sure we free all our memory so that valgrind doesn't complain
+	 * about leaked memory.
+	 */
+	fprintf(stderr, "\nWarning, some allocations not freed:");
+	for (node = alloc_head; node; node = next) {
+		next = node->next;
+		fprintf(stderr, "\nptr=%p, size=%zd\n", node->ptr, node->size);
+		fflush(stderr);
+		backtrace_symbols_fd(node->bt_buffer + 1, node->bt_levels - 1,
+				     2);
+		free(node);
+	}
+
+	return -1;
+}
diff --git a/host/lib/host_signature.c b/host/lib/host_signature.c
index 0ebbca6..1ea6bc4 100644
--- a/host/lib/host_signature.c
+++ b/host/lib/host_signature.c
@@ -65,7 +65,7 @@
 
   sig = SignatureAlloc(SHA512_DIGEST_SIZE, 0);
   if (!sig) {
-    free(header_checksum);
+    VbExFree(header_checksum);
     return NULL;
   }
   sig->sig_offset = sizeof(VbSignature);
@@ -74,7 +74,7 @@
 
   /* Signature data immediately follows the header */
   Memcpy(GetSignatureData(sig), header_checksum, SHA512_DIGEST_SIZE);
-  free(header_checksum);
+  VbExFree(header_checksum);
   return sig;
 }
 
@@ -128,12 +128,12 @@
   /* Prepend the digest info to the digest */
   signature_digest = malloc(signature_digest_len);
   if (!signature_digest) {
-    free(digest);
+    VbExFree(digest);
     return NULL;
   }
   Memcpy(signature_digest, digestinfo, digestinfo_size);
   Memcpy(signature_digest + digestinfo_size, digest, digest_size);
-  free(digest);
+  VbExFree(digest);
 
   /* Allocate output signature */
   sig = SignatureAlloc(siglen_map[key->algorithm], size);
diff --git a/tests/test_common.h b/tests/test_common.h
index f5d73d1..4acf588 100644
--- a/tests/test_common.h
+++ b/tests/test_common.h
@@ -50,4 +50,7 @@
 #define COL_RED "\x1b[0;31m"
 #define COL_STOP "\x1b[m"
 
+/* Check that all memory allocations were freed */
+int vboot_api_stub_check_memory(void);
+
 #endif  /* VBOOT_REFERENCE_TEST_COMMON_H_ */
diff --git a/tests/vboot_api_devmode_tests.c b/tests/vboot_api_devmode_tests.c
index 61f0de1..925a146 100644
--- a/tests/vboot_api_devmode_tests.c
+++ b/tests/vboot_api_devmode_tests.c
@@ -328,6 +328,8 @@
 
   if (!gTestSuccess)
     error_code = 255;
+  if (vboot_api_stub_check_memory())
+    error_code = 255;
 
   return error_code;
 }
diff --git a/tests/vboot_api_firmware_tests.c b/tests/vboot_api_firmware_tests.c
index 4a97b93..9826e55 100644
--- a/tests/vboot_api_firmware_tests.c
+++ b/tests/vboot_api_firmware_tests.c
@@ -249,6 +249,8 @@
 
   VbSelectFirmwareTest();
 
+  if (vboot_api_stub_check_memory())
+    error_code = 255;
   if (!gTestSuccess)
     error_code = 255;
 
diff --git a/tests/vboot_api_kernel2_tests.c b/tests/vboot_api_kernel2_tests.c
index 323f5f6..1cb2d83 100644
--- a/tests/vboot_api_kernel2_tests.c
+++ b/tests/vboot_api_kernel2_tests.c
@@ -566,5 +566,8 @@
 	VbBootDevTest();
 	VbBootRecTest();
 
+	if (vboot_api_stub_check_memory())
+		return 255;
+
 	return gTestSuccess ? 0 : 255;
 }
diff --git a/tests/vboot_api_kernel3_tests.c b/tests/vboot_api_kernel3_tests.c
index d885125..fa13d4e 100644
--- a/tests/vboot_api_kernel3_tests.c
+++ b/tests/vboot_api_kernel3_tests.c
@@ -361,5 +361,8 @@
 {
 	VbSoftwareSyncTest();
 
+	if (vboot_api_stub_check_memory())
+		return 255;
+
 	return gTestSuccess ? 0 : 255;
 }
diff --git a/tests/vboot_api_kernel4_tests.c b/tests/vboot_api_kernel4_tests.c
index 82b1501..bb82702 100644
--- a/tests/vboot_api_kernel4_tests.c
+++ b/tests/vboot_api_kernel4_tests.c
@@ -245,5 +245,8 @@
 {
 	VbSlkTest();
 
+	if (vboot_api_stub_check_memory())
+		return 255;
+
 	return gTestSuccess ? 0 : 255;
 }
diff --git a/tests/vboot_api_kernel_tests.c b/tests/vboot_api_kernel_tests.c
index 8cfc1cb..bccad0d 100644
--- a/tests/vboot_api_kernel_tests.c
+++ b/tests/vboot_api_kernel_tests.c
@@ -313,5 +313,8 @@
 {
 	VbTryLoadKernelTest();
 
+	if (vboot_api_stub_check_memory())
+		return 255;
+
 	return gTestSuccess ? 0 : 255;
 }
diff --git a/tests/vboot_audio_tests.c b/tests/vboot_audio_tests.c
index d674931..bc6a3ef 100644
--- a/tests/vboot_audio_tests.c
+++ b/tests/vboot_audio_tests.c
@@ -224,5 +224,8 @@
   if (!gTestSuccess)
     error_code = 255;
 
+  if (vboot_api_stub_check_memory())
+    error_code = 255;
+
   return error_code;
 }
diff --git a/tests/vboot_common2_tests.c b/tests/vboot_common2_tests.c
index 564518a..220391c 100644
--- a/tests/vboot_common2_tests.c
+++ b/tests/vboot_common2_tests.c
@@ -102,7 +102,7 @@
 
 	RSAPublicKeyFree(rsa);
 	free(sig);
-	free(digest);
+	VbExFree(digest);
 }
 
 static void ReSignKernelPreamble(VbKernelPreambleHeader *h,
@@ -277,5 +277,8 @@
 		return -1;
 	}
 
+	if (vboot_api_stub_check_memory())
+		return 255;
+
 	return gTestSuccess ? 0 : 255;
 }
diff --git a/tests/vboot_common3_tests.c b/tests/vboot_common3_tests.c
index b5499f6..b5c0cc9 100644
--- a/tests/vboot_common3_tests.c
+++ b/tests/vboot_common3_tests.c
@@ -24,7 +24,7 @@
 				    SHA512_DIGEST_ALGORITHM);
 	Memcpy(GetSignatureData(&h->key_block_checksum), newchk,
 	       SHA512_DIGEST_SIZE);
-	free(newchk);
+	VbExFree(newchk);
 }
 
 static void KeyBlockVerifyTest(const VbPublicKey *public_key,
@@ -376,5 +376,8 @@
 		return -1;
 	}
 
+	if (vboot_api_stub_check_memory())
+		return 255;
+
 	return gTestSuccess ? 0 : 255;
 }
diff --git a/tests/vboot_common_tests.c b/tests/vboot_common_tests.c
index 852be79..7cfea90 100644
--- a/tests/vboot_common_tests.c
+++ b/tests/vboot_common_tests.c
@@ -237,5 +237,8 @@
 	PublicKeyTest();
 	VbSharedDataTest();
 
+	if (vboot_api_stub_check_memory())
+		return 255;
+
 	return gTestSuccess ? 0 : 255;
 }
diff --git a/tests/vboot_display_tests.c b/tests/vboot_display_tests.c
index 8086d31..63f6890 100644
--- a/tests/vboot_display_tests.c
+++ b/tests/vboot_display_tests.c
@@ -264,5 +264,8 @@
 	DisplayKeyTest();
 	FontTest();
 
+	if (vboot_api_stub_check_memory())
+		return 255;
+
 	return gTestSuccess ? 0 : 255;
 }
diff --git a/tests/vboot_firmware_tests.c b/tests/vboot_firmware_tests.c
index 724e7f8..651197c 100644
--- a/tests/vboot_firmware_tests.c
+++ b/tests/vboot_firmware_tests.c
@@ -451,6 +451,8 @@
 
   LoadFirmwareTest();
 
+  if (vboot_api_stub_check_memory())
+    error_code = 255;
   if (!gTestSuccess)
     error_code = 255;
 
diff --git a/tests/vboot_kernel_tests.c b/tests/vboot_kernel_tests.c
index fdf81b7..ee164ce 100644
--- a/tests/vboot_kernel_tests.c
+++ b/tests/vboot_kernel_tests.c
@@ -567,5 +567,8 @@
 	InvalidParamsTest();
 	LoadKernelTest();
 
+	if (vboot_api_stub_check_memory())
+		return 255;
+
 	return gTestSuccess ? 0 : 255;
 }
diff --git a/tests/vboot_nvstorage_test.c b/tests/vboot_nvstorage_test.c
index 7ffec61..63e93d9 100644
--- a/tests/vboot_nvstorage_test.c
+++ b/tests/vboot_nvstorage_test.c
@@ -176,6 +176,8 @@
 
   VbNvStorageTest();
 
+  if (vboot_api_stub_check_memory())
+    error_code = 255;
   if (!gTestSuccess)
     error_code = 255;