15 files changed, 43 insertions, 2 deletions
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;