Improve CTR_DRBG error handling and cleanup

Check AES return values and return error when needed. Propagate the
underlying AES return code.

Perform more memory cleanup.
diff --git a/library/ctr_drbg.c b/library/ctr_drbg.c
index 55612c7..2d2da24 100644
--- a/library/ctr_drbg.c
+++ b/library/ctr_drbg.c
@@ -94,11 +94,15 @@
     /*
      * Initialize with an empty key
      */
-    mbedtls_aes_setkey_enc( &ctx->aes_ctx, key, MBEDTLS_CTR_DRBG_KEYBITS );
+    if( ( ret = mbedtls_aes_setkey_enc( &ctx->aes_ctx, key, MBEDTLS_CTR_DRBG_KEYBITS ) ) != 0 )
+    {
+        return( ret );
+    }
 
     if( ( ret = mbedtls_ctr_drbg_reseed( ctx, custom, len ) ) != 0 )
+    {
         return( ret );
-
+    }
     return( 0 );
 }
 
@@ -148,6 +152,7 @@
     unsigned char chain[MBEDTLS_CTR_DRBG_BLOCKSIZE];
     unsigned char *p, *iv;
     mbedtls_aes_context aes_ctx;
+    int ret = 0;
 
     int i, j;
     size_t buf_len, use_len;
@@ -180,7 +185,10 @@
     for( i = 0; i < MBEDTLS_CTR_DRBG_KEYSIZE; i++ )
         key[i] = i;
 
-    mbedtls_aes_setkey_enc( &aes_ctx, key, MBEDTLS_CTR_DRBG_KEYBITS );
+    if( ( ret = mbedtls_aes_setkey_enc( &aes_ctx, key, MBEDTLS_CTR_DRBG_KEYBITS ) ) != 0 )
+    {
+        goto exit;
+    }
 
     /*
      * Reduce data to MBEDTLS_CTR_DRBG_SEEDLEN bytes of data
@@ -199,7 +207,10 @@
             use_len -= ( use_len >= MBEDTLS_CTR_DRBG_BLOCKSIZE ) ?
                        MBEDTLS_CTR_DRBG_BLOCKSIZE : use_len;
 
-            mbedtls_aes_crypt_ecb( &aes_ctx, MBEDTLS_AES_ENCRYPT, chain, chain );
+            if( ( ret = mbedtls_aes_crypt_ecb( &aes_ctx, MBEDTLS_AES_ENCRYPT, chain, chain ) ) != 0 )
+            {
+                goto exit;
+            }
         }
 
         memcpy( tmp + j, chain, MBEDTLS_CTR_DRBG_BLOCKSIZE );
@@ -213,20 +224,40 @@
     /*
      * Do final encryption with reduced data
      */
-    mbedtls_aes_setkey_enc( &aes_ctx, tmp, MBEDTLS_CTR_DRBG_KEYBITS );
+    if( ( ret = mbedtls_aes_setkey_enc( &aes_ctx, tmp, MBEDTLS_CTR_DRBG_KEYBITS ) ) != 0 )
+    {
+        goto exit;
+    }
     iv = tmp + MBEDTLS_CTR_DRBG_KEYSIZE;
     p = output;
 
     for( j = 0; j < MBEDTLS_CTR_DRBG_SEEDLEN; j += MBEDTLS_CTR_DRBG_BLOCKSIZE )
     {
-        mbedtls_aes_crypt_ecb( &aes_ctx, MBEDTLS_AES_ENCRYPT, iv, iv );
+        if( ( ret = mbedtls_aes_crypt_ecb( &aes_ctx, MBEDTLS_AES_ENCRYPT, iv, iv ) ) != 0 )
+        {
+            goto exit;
+        }
         memcpy( p, iv, MBEDTLS_CTR_DRBG_BLOCKSIZE );
         p += MBEDTLS_CTR_DRBG_BLOCKSIZE;
     }
-
+exit:
     mbedtls_aes_free( &aes_ctx );
+    /*
+    * tidy up the stack
+    */
+    mbedtls_zeroize( buf, sizeof( buf ) );
+    mbedtls_zeroize( tmp, sizeof( tmp ) );
+    mbedtls_zeroize( key, sizeof( key ) );
+    mbedtls_zeroize( chain, sizeof( chain ) );
+    if( 0 != ret )
+    {
+        /*
+        * wipe partial seed from memory
+        */
+        mbedtls_zeroize( output, MBEDTLS_CTR_DRBG_SEEDLEN );
+    }
 
-    return( 0 );
+    return( ret );
 }
 
 static int ctr_drbg_update_internal( mbedtls_ctr_drbg_context *ctx,
@@ -235,6 +266,7 @@
     unsigned char tmp[MBEDTLS_CTR_DRBG_SEEDLEN];
     unsigned char *p = tmp;
     int i, j;
+    int ret = 0;
 
     memset( tmp, 0, MBEDTLS_CTR_DRBG_SEEDLEN );
 
@@ -250,7 +282,10 @@
         /*
          * Crypt counter block
          */
-        mbedtls_aes_crypt_ecb( &ctx->aes_ctx, MBEDTLS_AES_ENCRYPT, ctx->counter, p );
+        if( ( ret = mbedtls_aes_crypt_ecb( &ctx->aes_ctx, MBEDTLS_AES_ENCRYPT, ctx->counter, p ) ) != 0 )
+        {
+            return( ret );
+        }
 
         p += MBEDTLS_CTR_DRBG_BLOCKSIZE;
     }
@@ -261,7 +296,10 @@
     /*
      * Update key and counter
      */
-    mbedtls_aes_setkey_enc( &ctx->aes_ctx, tmp, MBEDTLS_CTR_DRBG_KEYBITS );
+    if( ( ret = mbedtls_aes_setkey_enc( &ctx->aes_ctx, tmp, MBEDTLS_CTR_DRBG_KEYBITS ) ) != 0 )
+    {
+        return( ret );
+    }
     memcpy( ctx->counter, tmp + MBEDTLS_CTR_DRBG_KEYSIZE, MBEDTLS_CTR_DRBG_BLOCKSIZE );
 
     return( 0 );
@@ -289,6 +327,7 @@
 {
     unsigned char seed[MBEDTLS_CTR_DRBG_MAX_SEED_INPUT];
     size_t seedlen = 0;
+    int ret;
 
     if( ctx->entropy_len > MBEDTLS_CTR_DRBG_MAX_SEED_INPUT ||
         len > MBEDTLS_CTR_DRBG_MAX_SEED_INPUT - ctx->entropy_len )
@@ -319,12 +358,18 @@
     /*
      * Reduce to 384 bits
      */
-    block_cipher_df( seed, seed, seedlen );
+    if( ( ret = block_cipher_df( seed, seed, seedlen ) ) != 0 )
+    {
+        return( ret );
+    }
 
     /*
      * Update state
      */
-    ctr_drbg_update_internal( ctx, seed );
+    if( ( ret = ctr_drbg_update_internal( ctx, seed ) ) != 0 )
+    {
+        return( ret );
+    }
     ctx->reseed_counter = 1;
 
     return( 0 );
@@ -354,15 +399,22 @@
         ctx->prediction_resistance )
     {
         if( ( ret = mbedtls_ctr_drbg_reseed( ctx, additional, add_len ) ) != 0 )
+        {
             return( ret );
-
+        }
         add_len = 0;
     }
 
     if( add_len > 0 )
     {
-        block_cipher_df( add_input, additional, add_len );
-        ctr_drbg_update_internal( ctx, add_input );
+        if( ( ret = block_cipher_df( add_input, additional, add_len ) ) != 0 )
+        {
+            return( ret );
+        }
+        if( ( ret = ctr_drbg_update_internal( ctx, add_input ) ) != 0 )
+        {
+            return( ret );
+        }
     }
 
     while( output_len > 0 )
@@ -377,7 +429,10 @@
         /*
          * Crypt counter block
          */
-        mbedtls_aes_crypt_ecb( &ctx->aes_ctx, MBEDTLS_AES_ENCRYPT, ctx->counter, tmp );
+        if( ( ret = mbedtls_aes_crypt_ecb( &ctx->aes_ctx, MBEDTLS_AES_ENCRYPT, ctx->counter, tmp ) ) != 0 )
+        {
+            return( ret );
+        }
 
         use_len = ( output_len > MBEDTLS_CTR_DRBG_BLOCKSIZE ) ? MBEDTLS_CTR_DRBG_BLOCKSIZE :
                                                        output_len;
@@ -389,7 +444,10 @@
         output_len -= use_len;
     }
 
-    ctr_drbg_update_internal( ctx, add_input );
+    if( ( ret = ctr_drbg_update_internal( ctx, add_input ) ) != 0 )
+    {
+        return( ret );
+    }
 
     ctx->reseed_counter++;