Add malloc failure tests.

This commit fixes a number of crashes caused by malloc failures. They
were found using the -malloc-test=0 option to runner.go which runs tests
many times, causing a different allocation call to fail in each case.

(This test only works on Linux and only looks for crashes caused by
allocation failures, not memory leaks or other errors.)

This is not the complete set of crashes! More can be found by collecting
core dumps from running with -malloc-test=0.

Change-Id: Ia61d19f51e373bccb7bc604642c51e043a74bd83
Reviewed-on: https://boringssl-review.googlesource.com/2320
Reviewed-by: Adam Langley <agl@google.com>
diff --git a/ssl/d1_clnt.c b/ssl/d1_clnt.c
index dcee9f8..3af756b 100644
--- a/ssl/d1_clnt.c
+++ b/ssl/d1_clnt.c
@@ -242,7 +242,12 @@
 			s->shutdown=0;
 
 			/* every DTLS ClientHello resets Finished MAC */
-			ssl3_init_finished_mac(s);
+			if (!ssl3_init_finished_mac(s))
+				{
+				OPENSSL_PUT_ERROR(SSL, dtls1_connect, ERR_R_INTERNAL_ERROR);
+				ret = -1;
+				goto end;
+				}
 
 			dtls1_start_timer(s);
 			ret=ssl3_send_client_hello(s);
diff --git a/ssl/d1_srvr.c b/ssl/d1_srvr.c
index bf13620..0bb99b6 100644
--- a/ssl/d1_srvr.c
+++ b/ssl/d1_srvr.c
@@ -239,9 +239,15 @@
 				 * the output is sent in a way that TCP likes :-)
 				 * ...but not with SCTP :-)
 				 */
-					if (!ssl_init_wbio_buffer(s,1)) { ret= -1; goto end; }
+				if (!ssl_init_wbio_buffer(s,1)) { ret= -1; goto end; }
 
-				ssl3_init_finished_mac(s);
+				if (!ssl3_init_finished_mac(s))
+					{
+					OPENSSL_PUT_ERROR(SSL, dtls1_accept, ERR_R_INTERNAL_ERROR);
+					ret = -1;
+					goto end;
+					}
+
 				s->state=SSL3_ST_SR_CLNT_HELLO_A;
 				s->ctx->stats.sess_accept++;
 				}
@@ -267,7 +273,12 @@
 			s->state=SSL3_ST_SW_FLUSH;
 			s->init_num=0;
 
-			ssl3_init_finished_mac(s);
+			if (!ssl3_init_finished_mac(s))
+				{
+				OPENSSL_PUT_ERROR(SSL, dtls1_accept, ERR_R_INTERNAL_ERROR);
+				ret = -1;
+				goto end;
+				}
 			break;
 
 		case SSL3_ST_SW_HELLO_REQ_C:
@@ -301,7 +312,12 @@
 			s->s3->tmp.next_state=SSL3_ST_SR_CLNT_HELLO_A;
 
 			/* HelloVerifyRequest resets Finished MAC */
-			ssl3_init_finished_mac(s);
+			if (!ssl3_init_finished_mac(s))
+				{
+				OPENSSL_PUT_ERROR(SSL, dtls1_accept, ERR_R_INTERNAL_ERROR);
+				ret = -1;
+				goto end;
+				}
 			break;
 			
 
diff --git a/ssl/s23_clnt.c b/ssl/s23_clnt.c
index efa3cd2..ebbd810 100644
--- a/ssl/s23_clnt.c
+++ b/ssl/s23_clnt.c
@@ -188,7 +188,12 @@
 
 			if (!ssl3_setup_buffers(s)) { ret= -1; goto end; }
 
-			ssl3_init_finished_mac(s);
+			if (!ssl3_init_finished_mac(s))
+				{
+				OPENSSL_PUT_ERROR(SSL, ssl23_connect, ERR_R_INTERNAL_ERROR);
+				ret = -1;
+				goto end;
+				}
 
 			s->state=SSL23_ST_CW_CLNT_HELLO_A;
 			s->ctx->stats.sess_connect++;
diff --git a/ssl/s23_srvr.c b/ssl/s23_srvr.c
index ca8e5d6..b4290bb 100644
--- a/ssl/s23_srvr.c
+++ b/ssl/s23_srvr.c
@@ -191,7 +191,12 @@
 				s->init_buf=buf;
 				}
 
-			ssl3_init_finished_mac(s);
+			if (!ssl3_init_finished_mac(s))
+				{
+				OPENSSL_PUT_ERROR(SSL, ssl23_accept, ERR_R_INTERNAL_ERROR);
+				ret = -1;
+				goto end;
+				}
 
 			s->state=SSL23_ST_SR_CLNT_HELLO;
 			s->ctx->stats.sess_accept++;
diff --git a/ssl/s3_cbc.c b/ssl/s3_cbc.c
index 44f408d..510d480 100644
--- a/ssl/s3_cbc.c
+++ b/ssl/s3_cbc.c
@@ -333,7 +333,7 @@
  *   ctx: the EVP_MD_CTX from which we take the hash function.
  *     ssl3_cbc_record_digest_supported must return true for this EVP_MD_CTX.
  *   md_out: the digest output. At most EVP_MAX_MD_SIZE bytes will be written.
- *   md_out_size: if non-NULL, the number of output bytes is written here.
+ *   md_out_size: the number of output bytes is written here.
  *   header: the 13-byte, TLS record header.
  *   data: the record data itself, less any preceeding explicit IV.
  *   data_plus_mac_size: the secret, reported length of the data and MAC
@@ -346,7 +346,7 @@
  * functions, above, we know that data_plus_mac_size is large enough to contain
  * a padding byte and MAC. (If the padding was invalid, it might contain the
  * padding too. ) */
-void ssl3_cbc_digest_record(
+int ssl3_cbc_digest_record(
 	const EVP_MD_CTX *ctx,
 	unsigned char* md_out,
 	size_t* md_out_size,
@@ -409,9 +409,8 @@
 			 * called first to check that the hash function is
 			 * supported. */
 			assert(0);
-			if (md_out_size)
-				*md_out_size = -1;
-			return;
+			*md_out_size = 0;
+			return 0;
 		}
 
 	assert(md_length_size <= MAX_HASH_BIT_COUNT_BYTES);
@@ -591,7 +590,12 @@
 		}
 
 	EVP_MD_CTX_init(&md_ctx);
-	EVP_DigestInit_ex(&md_ctx, ctx->digest, NULL /* engine */);
+	if (!EVP_DigestInit_ex(&md_ctx, ctx->digest, NULL /* engine */))
+		{
+		EVP_MD_CTX_cleanup(&md_ctx);
+		return 0;
+		}
+
 	if (is_sslv3)
 		{
 		/* We repurpose |hmac_pad| to contain the SSLv3 pad2 block. */
@@ -611,7 +615,8 @@
 		EVP_DigestUpdate(&md_ctx, mac_out, md_size);
 		}
 	EVP_DigestFinal(&md_ctx, md_out, &md_out_size_u);
-	if (md_out_size)
-		*md_out_size = md_out_size_u;
+	*md_out_size = md_out_size_u;
 	EVP_MD_CTX_cleanup(&md_ctx);
+
+	return 1;
 	}
diff --git a/ssl/s3_clnt.c b/ssl/s3_clnt.c
index 47ae767..54352d9 100644
--- a/ssl/s3_clnt.c
+++ b/ssl/s3_clnt.c
@@ -276,7 +276,12 @@
 
 			/* don't push the buffering BIO quite yet */
 
-			ssl3_init_finished_mac(s);
+			if (!ssl3_init_finished_mac(s))
+				{
+				OPENSSL_PUT_ERROR(SSL, ssl3_connect, ERR_R_INTERNAL_ERROR);
+				ret = -1;
+				goto end;
+				}
 
 			s->state=SSL3_ST_CW_CLNT_HELLO_A;
 			s->ctx->stats.sess_connect++;
diff --git a/ssl/s3_enc.c b/ssl/s3_enc.c
index fc94a94..2474934 100644
--- a/ssl/s3_enc.c
+++ b/ssl/s3_enc.c
@@ -185,7 +185,11 @@
 		for (j=0; j<k; j++)
 			buf[j]=c;
 		c++;
-		EVP_DigestInit_ex(&s1,EVP_sha1(), NULL);
+		if (!EVP_DigestInit_ex(&s1,EVP_sha1(), NULL))
+			{
+			OPENSSL_PUT_ERROR(SSL, ssl3_generate_key_block, ERR_LIB_EVP);
+			return 0;
+			}
 		EVP_DigestUpdate(&s1,buf,k);
 		EVP_DigestUpdate(&s1,s->session->master_key,
 			s->session->master_key_length);
@@ -193,7 +197,11 @@
 		EVP_DigestUpdate(&s1,s->s3->client_random,SSL3_RANDOM_SIZE);
 		EVP_DigestFinal_ex(&s1,smd,NULL);
 
-		EVP_DigestInit_ex(&m5,EVP_md5(), NULL);
+		if (!EVP_DigestInit_ex(&m5,EVP_md5(), NULL))
+			{
+			OPENSSL_PUT_ERROR(SSL, ssl3_generate_key_block, ERR_LIB_EVP);
+			return 0;
+			}
 		EVP_DigestUpdate(&m5,s->session->master_key,
 			s->session->master_key_length);
 		EVP_DigestUpdate(&m5,smd,SHA_DIGEST_LENGTH);
@@ -470,12 +478,17 @@
 	return(1);
 	}
 
-void ssl3_init_finished_mac(SSL *s)
+int ssl3_init_finished_mac(SSL *s)
 	{
 	if (s->s3->handshake_buffer) BIO_free(s->s3->handshake_buffer);
 	if (s->s3->handshake_dgst) ssl3_free_digest_list(s);
-    s->s3->handshake_buffer=BIO_new(BIO_s_mem());	
+	s->s3->handshake_buffer = BIO_new(BIO_s_mem());
+	if (s->s3->handshake_buffer == NULL) {
+		return 0;
+	}
 	(void)BIO_set_close(s->s3->handshake_buffer,BIO_CLOSE);
+
+	return 1;
 	}
 
 void ssl3_free_digest_list(SSL *s) 
@@ -521,6 +534,12 @@
 	/* Allocate handshake_dgst array */
 	ssl3_free_digest_list(s);
 	s->s3->handshake_dgst = OPENSSL_malloc(SSL_MAX_DIGEST * sizeof(EVP_MD_CTX *));
+	if (s->s3->handshake_dgst == NULL)
+		{
+		OPENSSL_PUT_ERROR(SSL, ssl3_digest_cached_records, ERR_R_MALLOC_FAILURE);
+		return 0;
+		}
+
 	memset(s->s3->handshake_dgst,0,SSL_MAX_DIGEST *sizeof(EVP_MD_CTX *));
 	if (!BIO_mem_contents(s->s3->handshake_buffer, &hdata, &hdatalen))
 		{
@@ -533,8 +552,18 @@
 		{
 		if ((mask & ssl_get_algorithm2(s)) && md) 
 			{
-			s->s3->handshake_dgst[i]=EVP_MD_CTX_create();
-			EVP_DigestInit_ex(s->s3->handshake_dgst[i],md,NULL);
+			s->s3->handshake_dgst[i] = EVP_MD_CTX_create();
+			if (s->s3->handshake_dgst[i] == NULL)
+				{
+				OPENSSL_PUT_ERROR(SSL, ssl3_digest_cached_records, ERR_LIB_EVP);
+				return 0;
+				}
+			if (!EVP_DigestInit_ex(s->s3->handshake_dgst[i], md, NULL))
+				{
+				EVP_MD_CTX_destroy(s->s3->handshake_dgst[i]);
+				OPENSSL_PUT_ERROR(SSL, ssl3_digest_cached_records, ERR_LIB_EVP);
+				return 0;
+				}
 			EVP_DigestUpdate(s->s3->handshake_dgst[i],hdata,hdatalen);
 			} 
 		else 
@@ -603,7 +632,12 @@
 		return 0;
 	}	
 	EVP_MD_CTX_init(&ctx);
-	EVP_MD_CTX_copy_ex(&ctx,d);
+	if (!EVP_MD_CTX_copy_ex(&ctx,d))
+		{
+		EVP_MD_CTX_cleanup(&ctx);
+		OPENSSL_PUT_ERROR(SSL, ssl3_generate_key_block, ERR_LIB_EVP);
+		return 0;
+		}
 	n=EVP_MD_CTX_size(&ctx);
 	if (n < 0)
 		return 0;
@@ -616,7 +650,12 @@
 	EVP_DigestUpdate(&ctx,ssl3_pad_1,npad);
 	EVP_DigestFinal_ex(&ctx,md_buf,&i);
 
-	EVP_DigestInit_ex(&ctx,EVP_MD_CTX_md(&ctx), NULL);
+	if (!EVP_DigestInit_ex(&ctx,EVP_MD_CTX_md(&ctx), NULL))
+		{
+		EVP_MD_CTX_cleanup(&ctx);
+		OPENSSL_PUT_ERROR(SSL, ssl3_generate_key_block, ERR_LIB_EVP);
+		return 0;
+		}
 	EVP_DigestUpdate(&ctx,s->session->master_key,
 		s->session->master_key_length);
 	EVP_DigestUpdate(&ctx,ssl3_pad_2,npad);
@@ -691,13 +730,16 @@
 		header[j++] = rec->length >> 8;
 		header[j++] = rec->length & 0xff;
 
-		ssl3_cbc_digest_record(
+		if (!ssl3_cbc_digest_record(
 			hash,
 			md, &md_size,
 			header, rec->input,
 			rec->length + md_size, orig_len,
 			mac_sec, md_size,
-			1 /* is SSLv3 */);
+			1 /* is SSLv3 */))
+			{
+			return -1;
+			}
 		}
 	else
 		{
@@ -705,7 +747,11 @@
 		/* Chop the digest off the end :-) */
 		EVP_MD_CTX_init(&md_ctx);
 
-		EVP_MD_CTX_copy_ex( &md_ctx,hash);
+		if (!EVP_MD_CTX_copy_ex(&md_ctx,hash))
+			{
+			EVP_MD_CTX_cleanup(&md_ctx);
+			return -1;
+			}
 		EVP_DigestUpdate(&md_ctx,mac_sec,md_size);
 		EVP_DigestUpdate(&md_ctx,ssl3_pad_1,npad);
 		EVP_DigestUpdate(&md_ctx,seq,8);
@@ -717,7 +763,11 @@
 		EVP_DigestUpdate(&md_ctx,rec->input,rec->length);
 		EVP_DigestFinal_ex( &md_ctx,md,NULL);
 
-		EVP_MD_CTX_copy_ex( &md_ctx,hash);
+		if (!EVP_MD_CTX_copy_ex(&md_ctx,hash))
+			{
+			EVP_MD_CTX_cleanup(&md_ctx);
+			return -1;
+			}
 		EVP_DigestUpdate(&md_ctx,mac_sec,md_size);
 		EVP_DigestUpdate(&md_ctx,ssl3_pad_2,npad);
 		EVP_DigestUpdate(&md_ctx,md,md_size);
@@ -761,7 +811,12 @@
 	EVP_MD_CTX_init(&ctx);
 	for (i=0; i<3; i++)
 		{
-		EVP_DigestInit_ex(&ctx, EVP_sha1(), NULL);
+		if (!EVP_DigestInit_ex(&ctx, EVP_sha1(), NULL))
+			{
+			ret = 0;
+			break;
+			}
+
 		EVP_DigestUpdate(&ctx,salt[i],strlen((const char *)salt[i]));
 		EVP_DigestUpdate(&ctx,p,len);
 		EVP_DigestUpdate(&ctx,&(s->s3->client_random[0]),
@@ -770,7 +825,12 @@
 			SSL3_RANDOM_SIZE);
 		EVP_DigestFinal_ex(&ctx,buf,&n);
 
-		EVP_DigestInit_ex(&ctx, EVP_md5(), NULL);
+		if (!EVP_DigestInit_ex(&ctx, EVP_md5(), NULL))
+			{
+			ret = 0;
+			break;
+			}
+
 		EVP_DigestUpdate(&ctx,p,len);
 		EVP_DigestUpdate(&ctx,buf,n);
 		EVP_DigestFinal_ex(&ctx,out,&n);
diff --git a/ssl/s3_srvr.c b/ssl/s3_srvr.c
index 2496783..c7bc58b 100644
--- a/ssl/s3_srvr.c
+++ b/ssl/s3_srvr.c
@@ -294,7 +294,12 @@
 				 */
 				if (!ssl_init_wbio_buffer(s,1)) { ret= -1; goto end; }
 				
-				ssl3_init_finished_mac(s);
+				if (!ssl3_init_finished_mac(s))
+					{
+					OPENSSL_PUT_ERROR(SSL, ssl3_accept, ERR_R_INTERNAL_ERROR);
+					ret = -1;
+					goto end;
+					}
 				s->state=SSL3_ST_SR_CLNT_HELLO_A;
 				s->ctx->stats.sess_accept++;
 				}
@@ -329,7 +334,12 @@
 			s->state=SSL3_ST_SW_FLUSH;
 			s->init_num=0;
 
-			ssl3_init_finished_mac(s);
+			if (!ssl3_init_finished_mac(s))
+				{
+				OPENSSL_PUT_ERROR(SSL, ssl3_accept, ERR_R_INTERNAL_ERROR);
+				ret = -1;
+				goto end;
+				}
 			break;
 
 		case SSL3_ST_SW_HELLO_REQ_C:
@@ -2578,16 +2588,22 @@
 							 &hctx, 1) < 0)
 				{
 				OPENSSL_free(session);
+				EVP_CIPHER_CTX_cleanup(&ctx);
+				HMAC_CTX_cleanup(&hctx);
 				return -1;
 				}
 			}
 		else
 			{
 			RAND_pseudo_bytes(iv, 16);
-			EVP_EncryptInit_ex(&ctx, EVP_aes_128_cbc(), NULL,
-					tctx->tlsext_tick_aes_key, iv);
-			HMAC_Init_ex(&hctx, tctx->tlsext_tick_hmac_key, 16,
-					tlsext_tick_md(), NULL);
+			if (!EVP_EncryptInit_ex(&ctx, EVP_aes_128_cbc(), NULL, tctx->tlsext_tick_aes_key, iv) ||
+			    !HMAC_Init_ex(&hctx, tctx->tlsext_tick_hmac_key, 16, tlsext_tick_md(), NULL))
+				{
+				OPENSSL_free(session);
+				EVP_CIPHER_CTX_cleanup(&ctx);
+				HMAC_CTX_cleanup(&hctx);
+				return -1;
+				}
 			memcpy(key_name, tctx->tlsext_tick_key_name, 16);
 			}
 
diff --git a/ssl/ssl_lib.c b/ssl/ssl_lib.c
index 94383e5..ee66093 100644
--- a/ssl/ssl_lib.c
+++ b/ssl/ssl_lib.c
@@ -3167,7 +3167,10 @@
 {
 	ssl_clear_hash_ctx(hash);
 	*hash = EVP_MD_CTX_create();
-	if (md) EVP_DigestInit_ex(*hash,md,NULL);
+	if (md != NULL && *hash != NULL)
+		{
+		EVP_DigestInit_ex(*hash,md,NULL);
+		}
 	return *hash;
 }
 void ssl_clear_hash_ctx(EVP_MD_CTX **hash) 
diff --git a/ssl/ssl_locl.h b/ssl/ssl_locl.h
index 083e2f4..b1c1899 100644
--- a/ssl/ssl_locl.h
+++ b/ssl/ssl_locl.h
@@ -819,7 +819,7 @@
 
 const SSL_CIPHER *ssl3_get_cipher_by_value(uint16_t value);
 uint16_t ssl3_get_cipher_value(const SSL_CIPHER *c);
-void ssl3_init_finished_mac(SSL *s);
+int ssl3_init_finished_mac(SSL *s);
 int ssl3_send_server_certificate(SSL *s);
 int ssl3_send_new_session_ticket(SSL *s);
 int ssl3_send_cert_status(SSL *s);
@@ -1120,7 +1120,7 @@
 			    unsigned block_size,
 			    unsigned mac_size);
 char ssl3_cbc_record_digest_supported(const EVP_MD_CTX *ctx);
-void ssl3_cbc_digest_record(
+int ssl3_cbc_digest_record(
 	const EVP_MD_CTX *ctx,
 	unsigned char* md_out,
 	size_t* md_out_size,
diff --git a/ssl/t1_enc.c b/ssl/t1_enc.c
index 6803e9b..05f2eca 100644
--- a/ssl/t1_enc.c
+++ b/ssl/t1_enc.c
@@ -447,6 +447,8 @@
 
 		cipher_ctx = s->enc_read_ctx;
 		mac_ctx = ssl_replace_hash(&s->read_hash, NULL);
+		if (mac_ctx == NULL)
+			goto err;
 
 		memcpy(s->s3->read_mac_secret, mac_secret, mac_secret_len);
 		s->s3->read_mac_secret_size = mac_secret_len;
@@ -476,7 +478,11 @@
 			s->write_hash = mac_ctx;
 			}
 		else
+			{
 			mac_ctx = ssl_replace_hash(&s->write_hash, NULL);
+			if (mac_ctx == NULL)
+				goto err;
+			}
 
 		memcpy(s->s3->write_mac_secret, mac_secret, mac_secret_len);
 		s->s3->write_mac_secret_size = mac_secret_len;
@@ -487,10 +493,15 @@
 				     NULL, mac_secret, mac_secret_len);
 	if (!mac_key)
 		return 0;
-	EVP_DigestSignInit(mac_ctx, NULL, s->s3->tmp.new_hash, NULL, mac_key);
+	if (!EVP_DigestSignInit(mac_ctx, NULL, s->s3->tmp.new_hash, NULL, mac_key))
+		{
+		EVP_PKEY_free(mac_key);
+		goto err;
+		}
 	EVP_PKEY_free(mac_key);
 
-	EVP_CipherInit_ex(cipher_ctx, cipher, NULL /* engine */, key, iv, !is_read);
+	if (!EVP_CipherInit_ex(cipher_ctx, cipher, NULL /* engine */, key, iv, !is_read))
+		goto err;
 
 	return 1;
 
@@ -1063,7 +1074,7 @@
 	unsigned char *seq;
 	EVP_MD_CTX *hash;
 	size_t md_size, orig_len;
-	int i;
+	int i, ok;
 	EVP_MD_CTX hmac, *mac_ctx;
 	unsigned char header[13];
 	int t;
@@ -1085,9 +1096,9 @@
 	assert(t >= 0);
 	md_size=t;
 
-	if (!EVP_MD_CTX_copy(&hmac,hash))
-		return -1;
 	mac_ctx = &hmac;
+	if (!EVP_MD_CTX_copy(mac_ctx, hash))
+		return -1;
 
 	if (SSL_IS_DTLS(ssl))
 		{
@@ -1119,7 +1130,7 @@
 		 * timing-side channel information about how many blocks of
 		 * data we are hashing because that gives an attacker a
 		 * timing-oracle. */
-		ssl3_cbc_digest_record(
+		ok = ssl3_cbc_digest_record(
 			mac_ctx,
 			md, &md_size,
 			header, rec->input,
@@ -1132,22 +1143,24 @@
 		{
 		EVP_DigestSignUpdate(mac_ctx,header,sizeof(header));
 		EVP_DigestSignUpdate(mac_ctx,rec->input,rec->length);
-		t=EVP_DigestSignFinal(mac_ctx,md,&md_size);
-		assert(t > 0);
+		ok = EVP_DigestSignFinal(mac_ctx,md,&md_size);
 		}
-		
-	EVP_MD_CTX_cleanup(&hmac);
+
+	EVP_MD_CTX_cleanup(mac_ctx);
+
+	if (!ok)
+		return -1;
 
 	if (!SSL_IS_DTLS(ssl))
 		{
 		for (i=7; i>=0; i--)
 			{
 			++seq[i];
-			if (seq[i] != 0) break; 
+			if (seq[i] != 0) break;
 			}
 		}
 
-	return(md_size);
+	return md_size;
 	}
 
 int tls1_generate_master_secret(SSL *s, unsigned char *out, unsigned char *p,
diff --git a/ssl/test/CMakeLists.txt b/ssl/test/CMakeLists.txt
index 27d9596..f7f0efc 100644
--- a/ssl/test/CMakeLists.txt
+++ b/ssl/test/CMakeLists.txt
@@ -5,8 +5,9 @@
 
 	async_bio.cc
 	bssl_shim.cc
+	malloc.cc
 	packeted_bio.cc
 	test_config.cc
 )
 
-target_link_libraries(bssl_shim ssl crypto)
+target_link_libraries(bssl_shim ssl crypto dl)
diff --git a/ssl/test/bssl_shim.cc b/ssl/test/bssl_shim.cc
index cdd62ff..9f9fac3 100644
--- a/ssl/test/bssl_shim.cc
+++ b/ssl/test/bssl_shim.cc
@@ -41,8 +41,8 @@
 
 static int g_ex_data_index = 0;
 
-static void SetConfigPtr(SSL *ssl, const TestConfig *config) {
-  SSL_set_ex_data(ssl, g_ex_data_index, (void *)config);
+static bool SetConfigPtr(SSL *ssl, const TestConfig *config) {
+  return SSL_set_ex_data(ssl, g_ex_data_index, (void *)config) == 1;
 }
 
 static const TestConfig *GetConfigPtr(SSL *ssl) {
@@ -267,7 +267,8 @@
   }
 
   dh = DH_get_2048_256(NULL);
-  if (!SSL_CTX_set_tmp_dh(ssl_ctx, dh)) {
+  if (dh == NULL ||
+      !SSL_CTX_set_tmp_dh(ssl_ctx, dh)) {
     goto err;
   }
 
@@ -335,7 +336,10 @@
     return 1;
   }
 
-  SetConfigPtr(ssl, config);
+  if (!SetConfigPtr(ssl, config)) {
+    BIO_print_errors_fp(stdout);
+    return 1;
+  }
 
   if (config->fallback_scsv) {
     if (!SSL_enable_fallback_scsv(ssl)) {
diff --git a/ssl/test/malloc.cc b/ssl/test/malloc.cc
new file mode 100644
index 0000000..3e3f8d7
--- /dev/null
+++ b/ssl/test/malloc.cc
@@ -0,0 +1,126 @@
+/* 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. */
+
+#if defined(__linux__)
+
+#include <stdint.h>
+#include <stdlib.h>
+#include <unistd.h>
+#include <unistd.h>
+#include <stdio.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) */
diff --git a/ssl/test/packeted_bio.cc b/ssl/test/packeted_bio.cc
index 629d6c5..3894d85 100644
--- a/ssl/test/packeted_bio.cc
+++ b/ssl/test/packeted_bio.cc
@@ -67,7 +67,9 @@
   uint32_t len = (len_bytes[0] << 24) | (len_bytes[1] << 16) |
       (len_bytes[2] << 8) | len_bytes[3];
   char *buf = (char *)OPENSSL_malloc(len);
-  assert(buf != NULL);
+  if (buf == NULL) {
+    return -1;
+  }
   ret = BIO_read(bio->next_bio, buf, len);
   assert(ret == (int)len);
 
diff --git a/ssl/test/runner/runner.go b/ssl/test/runner/runner.go
index dca0479..cf1b1f9 100644
--- a/ssl/test/runner/runner.go
+++ b/ssl/test/runner/runner.go
@@ -16,14 +16,19 @@
 	"os/exec"
 	"path"
 	"runtime"
+	"strconv"
 	"strings"
 	"sync"
 	"syscall"
 )
 
-var useValgrind = flag.Bool("valgrind", false, "If true, run code under valgrind")
-var useGDB = flag.Bool("gdb", false, "If true, run BoringSSL code under gdb")
-var flagDebug *bool = flag.Bool("debug", false, "Hexdump the contents of the connection")
+var (
+	useValgrind            = flag.Bool("valgrind", false, "If true, run code under valgrind")
+	useGDB                 = flag.Bool("gdb", false, "If true, run BoringSSL code under gdb")
+	flagDebug       *bool  = flag.Bool("debug", false, "Hexdump the contents of the connection")
+	mallocTest      *int64 = flag.Int64("malloc-test", -1, "If non-negative, run each test with each malloc in turn failing from the given number onwards.")
+	mallocTestDebug *bool  = flag.Bool("malloc-test-debug", false, "If true, ask bssl_shim to abort rather than fail a malloc. This can be used with a specific value for --malloc-test to identity the malloc failing that is causing problems.")
+)
 
 const (
 	rsaCertificateFile   = "cert.pem"
@@ -703,7 +708,15 @@
 	return shimEnd, conn
 }
 
-func runTest(test *testCase, buildDir string) error {
+type moreMallocsError struct{}
+
+func (moreMallocsError) Error() string {
+	return "child process did not exhaust all allocation calls"
+}
+
+var errMoreMallocs = moreMallocsError{}
+
+func runTest(test *testCase, buildDir string, mallocNumToFail int64) error {
 	if !test.shouldFail && (len(test.expectedError) > 0 || len(test.expectedLocalError) > 0) {
 		panic("Error expected without shouldFail in " + test.name)
 	}
@@ -758,6 +771,13 @@
 	var stdoutBuf, stderrBuf bytes.Buffer
 	shim.Stdout = &stdoutBuf
 	shim.Stderr = &stderrBuf
+	if mallocNumToFail >= 0 {
+		shim.Env = []string{"MALLOC_NUMBER_TO_FAIL=" + strconv.FormatInt(mallocNumToFail, 10)}
+		if *mallocTestDebug {
+			shim.Env = append(shim.Env, "MALLOC_ABORT_ON_FAIL=1")
+		}
+		shim.Env = append(shim.Env, "_MALLOC_CHECK=1")
+	}
 
 	if err := shim.Start(); err != nil {
 		panic(err)
@@ -805,6 +825,11 @@
 	connResume.Close()
 
 	childErr := shim.Wait()
+	if exitError, ok := childErr.(*exec.ExitError); ok {
+		if exitError.Sys().(syscall.WaitStatus).ExitStatus() == 88 {
+			return errMoreMallocs
+		}
+	}
 
 	stdout := string(stdoutBuf.Bytes())
 	stderr := string(stderrBuf.Bytes())
@@ -2142,8 +2167,22 @@
 	defer wg.Done()
 
 	for test := range c {
-		statusChan <- statusMsg{test: test, started: true}
-		err := runTest(test, buildDir)
+		var err error
+
+		if *mallocTest < 0 {
+			statusChan <- statusMsg{test: test, started: true}
+			err = runTest(test, buildDir, -1)
+		} else {
+			for mallocNumToFail := int64(*mallocTest); ; mallocNumToFail++ {
+				statusChan <- statusMsg{test: test, started: true}
+				if err = runTest(test, buildDir, mallocNumToFail); err != errMoreMallocs {
+					if err != nil {
+						fmt.Printf("\n\nmalloc test failed at %d: %s\n", mallocNumToFail, err)
+					}
+					break
+				}
+			}
+		}
 		statusChan <- statusMsg{test: test, err: err}
 	}
 }