From e495a8ae35dae75b103c6c1b70d332a12c7a82c3 Mon Sep 17 00:00:00 2001
From: Gilles Peskine <Gilles.Peskine@arm.com>
Date: Tue, 11 Sep 2018 15:34:17 +0200
Subject: [PATCH] CTR_DRBG: clean stack buffers

Wipe stack buffers that may contain sensitive data (data that
contributes to the DRBG state.
---
 library/ctr_drbg.c | 44 ++++++++++++++++++--------------------------
 1 file changed, 18 insertions(+), 26 deletions(-)

diff --git a/library/ctr_drbg.c b/library/ctr_drbg.c
index fead18f72..559538121 100644
--- a/library/ctr_drbg.c
+++ b/library/ctr_drbg.c
@@ -299,9 +299,7 @@ static int ctr_drbg_update_internal( mbedtls_ctr_drbg_context *ctx,
          * Crypt counter block
          */
         if( ( ret = mbedtls_aes_crypt_ecb( &ctx->aes_ctx, MBEDTLS_AES_ENCRYPT, ctx->counter, p ) ) != 0 )
-        {
-            return( ret );
-        }
+            goto exit;
 
         p += MBEDTLS_CTR_DRBG_BLOCKSIZE;
     }
@@ -313,12 +311,12 @@ static int ctr_drbg_update_internal( mbedtls_ctr_drbg_context *ctx,
      * Update key and counter
      */
     if( ( ret = mbedtls_aes_setkey_enc( &ctx->aes_ctx, tmp, MBEDTLS_CTR_DRBG_KEYBITS ) ) != 0 )
-    {
-        return( ret );
-    }
+        goto exit;
     memcpy( ctx->counter, tmp + MBEDTLS_CTR_DRBG_KEYSIZE, MBEDTLS_CTR_DRBG_BLOCKSIZE );
 
-    return( 0 );
+exit:
+    mbedtls_platform_zeroize( tmp, sizeof( tmp ) );
+    return( ret );
 }
 
 /* CTR_DRBG_Instantiate with derivation function (SP 800-90A &sect;10.2.1.3.2)
@@ -347,6 +345,7 @@ void mbedtls_ctr_drbg_update( mbedtls_ctr_drbg_context *ctx,
 
         block_cipher_df( add_input, additional, add_len );
         ctr_drbg_update_internal( ctx, add_input );
+        mbedtls_platform_zeroize( add_input, sizeof( add_input ) );
     }
 }
 
@@ -399,20 +398,18 @@ int mbedtls_ctr_drbg_reseed( mbedtls_ctr_drbg_context *ctx,
      * Reduce to 384 bits
      */
     if( ( ret = block_cipher_df( seed, seed, seedlen ) ) != 0 )
-    {
-        return( ret );
-    }
+        goto exit;
 
     /*
      * Update state
      */
     if( ( ret = ctr_drbg_update_internal( ctx, seed ) ) != 0 )
-    {
-        return( ret );
-    }
+        goto exit;
     ctx->reseed_counter = 1;
 
-    return( 0 );
+exit:
+    mbedtls_platform_zeroize( seed, sizeof( seed ) );
+    return( ret );
 }
 
 /* CTR_DRBG_Generate with derivation function (SP 800-90A &sect;10.2.1.5.2)
@@ -467,13 +464,9 @@ int mbedtls_ctr_drbg_random_with_add( void *p_rng,
     if( add_len > 0 )
     {
         if( ( ret = block_cipher_df( add_input, additional, add_len ) ) != 0 )
-        {
-            return( ret );
-        }
+            goto exit;
         if( ( ret = ctr_drbg_update_internal( ctx, add_input ) ) != 0 )
-        {
-            return( ret );
-        }
+            goto exit;
     }
 
     while( output_len > 0 )
@@ -489,9 +482,7 @@ int mbedtls_ctr_drbg_random_with_add( void *p_rng,
          * Crypt counter block
          */
         if( ( ret = mbedtls_aes_crypt_ecb( &ctx->aes_ctx, MBEDTLS_AES_ENCRYPT, ctx->counter, tmp ) ) != 0 )
-        {
-            return( ret );
-        }
+            goto exit;
 
         use_len = ( output_len > MBEDTLS_CTR_DRBG_BLOCKSIZE ) ? MBEDTLS_CTR_DRBG_BLOCKSIZE :
                                                        output_len;
@@ -504,12 +495,13 @@ int mbedtls_ctr_drbg_random_with_add( void *p_rng,
     }
 
     if( ( ret = ctr_drbg_update_internal( ctx, add_input ) ) != 0 )
-    {
-        return( ret );
-    }
+        goto exit;
 
     ctx->reseed_counter++;
 
+exit:
+    mbedtls_platform_zeroize( add_input, sizeof( add_input ) );
+    mbedtls_platform_zeroize( tmp, sizeof( tmp ) );
     return( 0 );
 }