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

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;
+}

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);

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_ */

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;
 }

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;
 

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;
 }

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;
 }

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;
 }

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;
 }

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;
 }

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;
 }

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;
 }

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;
 }

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;
 }

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;
 

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;
 }

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;