4 files changed, 158 insertions, 0 deletions
diff --git a/src/crypto/test/CMakeLists.txt b/src/crypto/test/CMakeLists.txt
index 0d5ca81..84a6174 100644
--- a/src/crypto/test/CMakeLists.txt
+++ b/src/crypto/test/CMakeLists.txt
@@ -4,4 +4,5 @@ add_library(
   OBJECT
 
   file_test.cc
+  malloc.cc
 )
diff --git a/src/crypto/test/file_test.cc b/src/crypto/test/file_test.cc
index 12405f2..8df6f9a 100644
--- a/src/crypto/test/file_test.cc
+++ b/src/crypto/test/file_test.cc
@@ -17,6 +17,7 @@
 #include <ctype.h>
 #include <errno.h>
 #include <stdarg.h>
+#include <stdlib.h>
 #include <string.h>
 
 #include <openssl/err.h>
diff --git a/src/crypto/test/malloc.cc b/src/crypto/test/malloc.cc
new file mode 100644
index 0000000..9ffdf01
--- /dev/null
+++ b/src/crypto/test/malloc.cc
@@ -0,0 +1,145 @@
+/* Copyright (c) 2014, Google Inc.
+ *
+ * Permission to use, copy, modify, and/or distribute this software for any
+ * purpose with or without fee is hereby granted, provided that the above
+ * copyright notice and this permission notice appear in all copies.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+ * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+ * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
+ * SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+ * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
+ * OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
+ * CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. */
+
+#include <openssl/base.h>
+
+#if defined(__has_feature)
+#if __has_feature(address_sanitizer) || __has_feature(memory_sanitizer)
+#define OPENSSL_ASAN
+#endif
+#endif
+
+#if defined(__GLIBC__) && !defined(__UCLIBC__)
+#define OPENSSL_GLIBC
+#endif
+
+// This file isn't built on ARM or Aarch64 because we link statically in those
+// builds and trying to override malloc in a static link doesn't work. It also
+// requires glibc. It's also disabled on ASan builds as this interferes with
+// ASan's malloc interceptor.
+//
+// TODO(davidben): See if this and ASan's and MSan's interceptors can be made to
+// coexist.
+#if defined(__linux__) && defined(OPENSSL_GLIBC) && !defined(OPENSSL_ARM) && \
+    !defined(OPENSSL_AARCH64) && !defined(OPENSSL_ASAN)
+
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <unistd.h>
+
+#include <new>
+
+
+/* This file defines overrides for the standard allocation functions that allow
+ * a given allocation to be made to fail for testing. If the program is run
+ * with MALLOC_NUMBER_TO_FAIL set to a base-10 number then that allocation will
+ * return NULL. If MALLOC_ABORT_ON_FAIL is also defined then the allocation
+ * will abort() rather than return NULL.
+ *
+ * This code is not thread safe. */
+
+static uint64_t current_malloc_count = 0;
+static uint64_t malloc_number_to_fail = 0;
+static char failure_enabled = 0, abort_on_fail = 0;
+static int in_call = 0;
+
+extern "C" {
+/* These are other names for the standard allocation functions. */
+extern void *__libc_malloc(size_t size);
+extern void *__libc_calloc(size_t num_elems, size_t size);
+extern void *__libc_realloc(void *ptr, size_t size);
+}
+
+static void exit_handler(void) {
+  if (failure_enabled && current_malloc_count > malloc_number_to_fail) {
+    _exit(88);
+  }
+}
+
+static void cpp_new_handler() {
+  // Return to try again. It won't fail a second time.
+  return;
+}
+
+/* should_fail_allocation returns true if the current allocation should fail. */
+static int should_fail_allocation() {
+  static int init = 0;
+  char should_fail;
+
+  if (in_call) {
+    return 0;
+  }
+
+  in_call = 1;
+
+  if (!init) {
+    const char *env = getenv("MALLOC_NUMBER_TO_FAIL");
+    if (env != NULL && env[0] != 0) {
+      char *endptr;
+      malloc_number_to_fail = strtoull(env, &endptr, 10);
+      if (*endptr == 0) {
+        failure_enabled = 1;
+        atexit(exit_handler);
+        std::set_new_handler(cpp_new_handler);
+      }
+    }
+    abort_on_fail = (NULL != getenv("MALLOC_ABORT_ON_FAIL"));
+    init = 1;
+  }
+
+  in_call = 0;
+
+  if (!failure_enabled) {
+    return 0;
+  }
+
+  should_fail = (current_malloc_count == malloc_number_to_fail);
+  current_malloc_count++;
+
+  if (should_fail && abort_on_fail) {
+    abort();
+  }
+  return should_fail;
+}
+
+extern "C" {
+
+void *malloc(size_t size) {
+  if (should_fail_allocation()) {
+    return NULL;
+  }
+
+  return __libc_malloc(size);
+}
+
+void *calloc(size_t num_elems, size_t size) {
+  if (should_fail_allocation()) {
+    return NULL;
+  }
+
+  return __libc_calloc(num_elems, size);
+}
+
+void *realloc(void *ptr, size_t size) {
+  if (should_fail_allocation()) {
+    return NULL;
+  }
+
+  return __libc_realloc(ptr, size);
+}
+
+}  // extern "C"
+
+#endif  /* defined(linux) && GLIBC && !ARM && !AARCH64 && !ASAN */
diff --git a/src/crypto/test/scoped_types.h b/src/crypto/test/scoped_types.h
index eb04c18..c5c8cfe 100644
--- a/src/crypto/test/scoped_types.h
+++ b/src/crypto/test/scoped_types.h
@@ -16,6 +16,7 @@
 #define OPENSSL_HEADER_CRYPTO_TEST_SCOPED_TYPES_H
 
 #include <stdint.h>
+#include <stdio.h>
 
 #include <openssl/bio.h>
 #include <openssl/bn.h>
@@ -57,6 +58,12 @@ struct OpenSSLFree {
   }
 };
 
+struct FileCloser {
+  void operator()(FILE *file) {
+    fclose(file);
+  }
+};
+
 template<typename T, void (*func)(T*)>
 using ScopedOpenSSLType = bssl::unique_ptr<T, OpenSSLDeleter<T, func>>;
 
@@ -108,6 +115,9 @@ using ScopedX509_ALGOR = ScopedOpenSSLType<X509_ALGOR, X509_ALGOR_free>;
 
 using ScopedX509Stack = ScopedOpenSSLStack<STACK_OF(X509), X509, X509_free>;
 
+using ScopedEVP_CIPHER_CTX = ScopedOpenSSLContext<EVP_CIPHER_CTX, int,
+                                                  EVP_CIPHER_CTX_init,
+                                                  EVP_CIPHER_CTX_cleanup>;
 using ScopedEVP_MD_CTX = ScopedOpenSSLContext<EVP_MD_CTX, int, EVP_MD_CTX_init,
                                               EVP_MD_CTX_cleanup>;
 using ScopedHMAC_CTX = ScopedOpenSSLContext<HMAC_CTX, void, HMAC_CTX_init,
@@ -116,5 +126,6 @@ using ScopedHMAC_CTX = ScopedOpenSSLContext<HMAC_CTX, void, HMAC_CTX_init,
 using ScopedOpenSSLBytes = bssl::unique_ptr<uint8_t, OpenSSLFree<uint8_t>>;
 using ScopedOpenSSLString = bssl::unique_ptr<char, OpenSSLFree<char>>;
 
+using ScopedFILE = bssl::unique_ptr<FILE, FileCloser>;
 
 #endif  // OPENSSL_HEADER_CRYPTO_TEST_SCOPED_TYPES_H