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

17 files changed, 135 insertions, 6 deletions

diff --git a/firmware/stub/vboot_api_stub_sf.c b/firmware/stub/vboot_api_stub_sf.c
index e2f1c1ae..3c5ec85a 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 @@ VbError_t VbExHashFirmwareBody(VbCommonParams *cparams,
 {
 	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 0ebbca68..1ea6bc40 100644
--- a/host/lib/host_signature.c
+++ b/host/lib/host_signature.c
@@ -65,7 +65,7 @@ VbSignature* CalculateChecksum(const uint8_t* data, uint64_t size) {
 
   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 @@ VbSignature* CalculateChecksum(const uint8_t* data, uint64_t size) {
 
   /* 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 @@ VbSignature* CalculateSignature(const uint8_t* data, uint64_t size,
   /* 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 f5d73d1d..4acf5887 100644
--- a/tests/test_common.h
+++ b/tests/test_common.h
@@ -50,4 +50,7 @@ int TEST_FALSE(int result, const char* testname);
 #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 61f0de14..925a146b 100644
--- a/tests/vboot_api_devmode_tests.c
+++ b/tests/vboot_api_devmode_tests.c
@@ -328,6 +328,8 @@ int main(int argc, char* argv[]) {
 
   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 4a97b930..9826e550 100644
--- a/tests/vboot_api_firmware_tests.c
+++ b/tests/vboot_api_firmware_tests.c
@@ -249,6 +249,8 @@ int main(int argc, char* argv[]) {
 
   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 323f5f66..1cb2d835 100644
--- a/tests/vboot_api_kernel2_tests.c
+++ b/tests/vboot_api_kernel2_tests.c
@@ -566,5 +566,8 @@ int main(void)
 	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 d8851252..fa13d4e9 100644
--- a/tests/vboot_api_kernel3_tests.c
+++ b/tests/vboot_api_kernel3_tests.c
@@ -361,5 +361,8 @@ int main(void)
 {
 	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 82b15016..bb827024 100644
--- a/tests/vboot_api_kernel4_tests.c
+++ b/tests/vboot_api_kernel4_tests.c
@@ -245,5 +245,8 @@ int main(void)
 {
 	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 8cfc1cb4..bccad0de 100644
--- a/tests/vboot_api_kernel_tests.c
+++ b/tests/vboot_api_kernel_tests.c
@@ -313,5 +313,8 @@ int main(void)
 {
 	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 d674931d..bc6a3efc 100644
--- a/tests/vboot_audio_tests.c
+++ b/tests/vboot_audio_tests.c
@@ -224,5 +224,8 @@ int main(int argc, char* argv[]) {
   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 564518ac..220391cb 100644
--- a/tests/vboot_common2_tests.c
+++ b/tests/vboot_common2_tests.c
@@ -102,7 +102,7 @@ static void VerifyDigestTest(const VbPublicKey *public_key,
 
 	RSAPublicKeyFree(rsa);
 	free(sig);
-	free(digest);
+	VbExFree(digest);
 }
 
 static void ReSignKernelPreamble(VbKernelPreambleHeader *h,
@@ -277,5 +277,8 @@ int main(int argc, char *argv[]) {
 		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 b5499f66..b5c0cc93 100644
--- a/tests/vboot_common3_tests.c
+++ b/tests/vboot_common3_tests.c
@@ -24,7 +24,7 @@ static void ReChecksumKeyBlock(VbKeyBlockHeader *h)
 				    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 @@ int main(int argc, char *argv[])
 		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 852be795..7cfea90e 100644
--- a/tests/vboot_common_tests.c
+++ b/tests/vboot_common_tests.c
@@ -237,5 +237,8 @@ int main(int argc, char* argv[])
 	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 8086d319..63f68903 100644
--- a/tests/vboot_display_tests.c
+++ b/tests/vboot_display_tests.c
@@ -264,5 +264,8 @@ int main(void)
 	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 724e7f8b..651197cf 100644
--- a/tests/vboot_firmware_tests.c
+++ b/tests/vboot_firmware_tests.c
@@ -451,6 +451,8 @@ int main(int argc, char* argv[]) {
 
   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 fdf81b7a..ee164cea 100644
--- a/tests/vboot_kernel_tests.c
+++ b/tests/vboot_kernel_tests.c
@@ -567,5 +567,8 @@ int main(void)
 	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 7ffec614..63e93d9d 100644
--- a/tests/vboot_nvstorage_test.c
+++ b/tests/vboot_nvstorage_test.c
@@ -176,6 +176,8 @@ int main(int argc, char* argv[]) {
 
   VbNvStorageTest();
 
+  if (vboot_api_stub_check_memory())
+    error_code = 255;
   if (!gTestSuccess)
     error_code = 255;