Add buffer and context clearing upon suspected FI

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
diff --git a/library/aes.c b/library/aes.c
index 8cfb4ba..e49f74f 100644
--- a/library/aes.c
+++ b/library/aes.c
@@ -822,6 +822,7 @@
         }
     }
 
+    mbedtls_platform_memset( RK, 0, ( keybits >> 5 ) * 4 );
     return( MBEDTLS_ERR_PLATFORM_FAULT_DETECTED );
 }
 #endif /* !MBEDTLS_AES_SETKEY_ENC_ALT */
@@ -1176,6 +1177,8 @@
         }
     }
 
+    // Clear the output in case of a FI
+    mbedtls_platform_memset( output, 0, 16 );
     return( MBEDTLS_ERR_PLATFORM_FAULT_DETECTED );
 }
 
@@ -1460,6 +1463,8 @@
         }
     }
 
+    // Clear the output in case of a FI
+    mbedtls_platform_memset( output, 0, 16 );
     return( MBEDTLS_ERR_PLATFORM_FAULT_DETECTED );
 }
 
diff --git a/library/ccm.c b/library/ccm.c
index ab0540b..54d051e 100644
--- a/library/ccm.c
+++ b/library/ccm.c
@@ -101,12 +101,14 @@
         return( ret );
     }
 
-    if( keybits_dup != keybits || key_dup != key )
+    if( keybits_dup == keybits && key_dup == key )
     {
-        return MBEDTLS_ERR_PLATFORM_FAULT_DETECTED;
+        return( ret );
     }
 
-    return( ret );
+    // In case of a FI - clear the context
+    mbedtls_cipher_free( &ctx->cipher_ctx );
+    return( MBEDTLS_ERR_PLATFORM_FAULT_DETECTED );
 }
 
 /*
@@ -336,6 +338,9 @@
         add_dup != add || add_len_dup != add_len || input_dup != input ||
         output_dup != output || tag_dup != tag || tag_len_dup != tag_len)
     {
+
+        // In case of a FI - clear the output
+        mbedtls_platform_memset( output, 0, length );
         return MBEDTLS_ERR_PLATFORM_FAULT_DETECTED;
     }
 
diff --git a/library/hmac_drbg.c b/library/hmac_drbg.c
index ecca880..58750c8 100644
--- a/library/hmac_drbg.c
+++ b/library/hmac_drbg.c
@@ -212,6 +212,7 @@
     int ret = MBEDTLS_ERR_PLATFORM_FAULT_DETECTED;
     volatile const unsigned char *additional_dup = additional;
     volatile size_t len_dup = len;
+    int reseed_counter_backup = -1;
 
     if( use_nonce == HMAC_NONCE_NO )
         total_entropy_len = ctx->entropy_len;
@@ -269,6 +270,7 @@
         goto exit;
 
     /* 3. Reset reseed_counter */
+    reseed_counter_backup = ctx->reseed_counter;
     ctx->reseed_counter = 1;
 
 exit:
@@ -278,6 +280,10 @@
 
     if( additional_dup != additional || len_dup != len )
     {
+        /* Rollback the reseed_counter in case of FI */
+        if( reseed_counter_backup != -1 )
+            ctx->reseed_counter = reseed_counter_backup;
+
         return MBEDTLS_ERR_PLATFORM_FAULT_DETECTED;
     }
 
@@ -290,6 +296,9 @@
         return ret;
     }
 
+    /* Rollback the reseed_counter in case of FI */
+    if( reseed_counter_backup != -1 )
+        ctx->reseed_counter = reseed_counter_backup;
     return( MBEDTLS_ERR_PLATFORM_FAULT_DETECTED );
 }
 
diff --git a/library/sha256.c b/library/sha256.c
index bf52eae..493e88e 100644
--- a/library/sha256.c
+++ b/library/sha256.c
@@ -287,7 +287,8 @@
     {
         return( 0 );
     }
-
+    /* Free the ctx upon suspected FI */
+    mbedtls_sha256_free( ctx );
     return( MBEDTLS_ERR_PLATFORM_FAULT_DETECTED );
 }
 
@@ -362,6 +363,8 @@
             return( 0 );
         }
     }
+    /* Free the ctx upon suspected FI */
+    mbedtls_sha256_free( ctx );
     return( MBEDTLS_ERR_PLATFORM_FAULT_DETECTED );
 }
 
@@ -458,6 +461,9 @@
     {
         return( 0 );
     }
+    /* Free the ctx and clear output upon suspected FI */
+    mbedtls_sha256_free( ctx );
+    mbedtls_platform_memset( output, 0, 32 );
     return( MBEDTLS_ERR_PLATFORM_FAULT_DETECTED );
 }
 
@@ -506,6 +512,7 @@
     {
         return( ret );
     }
+    mbedtls_platform_memset( output, 0, 32 );
     return( MBEDTLS_ERR_PLATFORM_FAULT_DETECTED );
 }