Ensure mbedtls_sha3_finish zeroizes the context

Signed-off-by: Dave Rodgman <dave.rodgman@arm.com>
diff --git a/library/sha3.c b/library/sha3.c
index 4b97a85..9825504 100644
--- a/library/sha3.c
+++ b/library/sha3.c
@@ -259,10 +259,13 @@
 int mbedtls_sha3_finish(mbedtls_sha3_context *ctx,
                         uint8_t *output, size_t olen)
 {
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+
     /* Catch SHA-3 families, with fixed output length */
     if (ctx->olen > 0) {
         if (ctx->olen > olen) {
-            return MBEDTLS_ERR_SHA3_BAD_INPUT_DATA;
+            ret = MBEDTLS_ERR_SHA3_BAD_INPUT_DATA;
+            goto exit;
         }
         olen = ctx->olen;
     }
@@ -280,7 +283,11 @@
         }
     }
 
-    return 0;
+    ret = 0;
+
+exit:
+    mbedtls_platform_zeroize(ctx, sizeof(mbedtls_sha3_context));
+    return ret;
 }
 
 /*
diff --git a/tests/suites/test_suite_shax.function b/tests/suites/test_suite_shax.function
index 7dd9166..629e281 100644
--- a/tests/suites/test_suite_shax.function
+++ b/tests/suites/test_suite_shax.function
@@ -176,9 +176,12 @@
     TEST_EQUAL(mbedtls_sha3_starts(&ctx, MBEDTLS_SHA3_NONE), MBEDTLS_ERR_SHA3_BAD_INPUT_DATA);
 
     TEST_EQUAL(mbedtls_sha3_starts(&ctx, MBEDTLS_SHA3_256), 0);
-
     TEST_EQUAL(mbedtls_sha3_finish(&ctx, output, 0), MBEDTLS_ERR_SHA3_BAD_INPUT_DATA);
+
+    TEST_EQUAL(mbedtls_sha3_starts(&ctx, MBEDTLS_SHA3_256), 0);
     TEST_EQUAL(mbedtls_sha3_finish(&ctx, output, 31), MBEDTLS_ERR_SHA3_BAD_INPUT_DATA);
+
+    TEST_EQUAL(mbedtls_sha3_starts(&ctx, MBEDTLS_SHA3_256), 0);
     TEST_EQUAL(mbedtls_sha3_finish(&ctx, output, 32), 0);
 
 exit: