From 0f87fffc02b72e2487e0fd24be472ae08fb10231 Mon Sep 17 00:00:00 2001
From: Robert Lukotka <robert.lukotka@vacuumlabs.com>
Date: Sat, 7 Oct 2023 12:29:06 +0200
Subject: [PATCH] Review changes

---
 src/diffieHellman.c | 21 ++++++++++++---------
 src/hash.h          | 13 +++++--------
 2 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/src/diffieHellman.c b/src/diffieHellman.c
index 5e4ea8ef..44c28537 100644
--- a/src/diffieHellman.c
+++ b/src/diffieHellman.c
@@ -97,20 +97,20 @@ static WARN_UNUSED_RESULT uint16_t processDHOneBlockFromCache(dh_context_t *ctx,
 
 __noinline_due_to_stack__ WARN_UNUSED_RESULT uint16_t
 dh_init_aes_key(dh_aes_key_t *dhKey, const bip44_path_t *pathSpec, const public_key_t *publicKey) {
-    TRACE_STACK_USAGE();
+    // Crypto assets to be cleared
     private_key_t privateKey;
-    unsigned char basicSecret[32];
-    unsigned char secret[SHA_512_SIZE];
-    uint8_t K[SHA_512_SIZE];
-
-    TRACE("dh_init_aesKey");
-
     explicit_bzero(&privateKey, SIZEOF(privateKey));
+    unsigned char basicSecret[32];
     explicit_bzero(basicSecret, SIZEOF(basicSecret));
+    unsigned char secret[SHA_512_SIZE];
     explicit_bzero(secret, SIZEOF(secret));
+    uint8_t K[SHA_512_SIZE];
     explicit_bzero(K, SIZEOF(K));
     explicit_bzero(dhKey, SIZEOF(*dhKey));
 
+    TRACE_STACK_USAGE();
+    TRACE("dh_init_aesKey");
+
     {
         uint16_t err = derivePrivateKey(pathSpec, &privateKey);
         if (err != SUCCESS) {
@@ -157,6 +157,7 @@ dh_init_aes_key(dh_aes_key_t *dhKey, const bip44_path_t *pathSpec, const public_
 
     STATIC_ASSERT(SIZEOF(dhKey->km) == DH_KM_SIZE, "Incompatible types");
     memmove(dhKey->km, K + DH_AES_SECRET_SIZE, DH_KM_SIZE);
+    explicit_bzero(K, SIZEOF(K));
     return SUCCESS;
 }
 
@@ -280,6 +281,7 @@ dh_encode_append(dh_context_t *ctx,
         written += restLength;  // processDHOneBlockFromCache returns number of bytes written
     }
 
+    explicit_bzero(&aesKey, SIZEOF(aesKey));
     TRACE("Leaving dh_encode_append, written: %d", (int) written);
     *outSize = written;
     return SUCCESS;
@@ -294,6 +296,7 @@ dh_encode_finalize(dh_context_t *ctx,
     // Crypto assets to be cleared
     dh_aes_key_t aesKey;
     explicit_bzero(&aesKey, SIZEOF(aesKey));
+    // ctx->base64EncodingCache will contain final HMAC at one point
     // Crypto assets to be cleared in case of failure
     explicit_bzero(outBuffer, *outSize);  // It suffices to clear this after final hmac
 
@@ -404,9 +407,9 @@ dh_encode_finalize(dh_context_t *ctx,
             explicit_bzero(outBuffer, *outSize);
             return ERR_ASSERT;
     }
+    explicit_bzero(ctx->base64EncodingCache, SIZEOF(ctx->base64EncodingCache));
     *outSize = written;
     return SUCCESS;
-    ;
 }
 
 #ifdef DEVEL
@@ -493,7 +496,7 @@ __noinline_due_to_stack__ WARN_UNUSED_RESULT uint16_t dh_decode(bip44_path_t *pa
     dh_aes_key_t aesKey;
     explicit_bzero(&aesKey, SIZEOF(aesKey));
     // Crypto assets to be cleared in case of failure
-    TRACE_BUFFER(buffer, *size);  // buffer after decoding
+    TRACE_BUFFER(buffer, *size);  // only after decoding
 
     if (*size < DH_AES_IV_SIZE + CX_AES_BLOCK_SIZE + DH_HMAC_SIZE) {
         return ERR_INVALID_DATA;
diff --git a/src/hash.h b/src/hash.h
index cfbe08eb..ffff32f0 100644
--- a/src/hash.h
+++ b/src/hash.h
@@ -33,7 +33,8 @@ static __attribute__((always_inline, unused)) cx_err_t sha_256_init(sha_256_cont
 static __attribute__((always_inline, unused)) cx_err_t sha_256_append(sha_256_context_t* ctx,
                                                                       const uint8_t* inBuffer,
                                                                       size_t inSize) {
-    if (ctx->initialized_magic != HASH_CONTEXT_INITIALIZED_MAGIC) {
+    if (ctx->initialized_magic != HASH_CONTEXT_INITIALIZED_MAGIC ||
+        inSize >= BUFFER_SIZE_PARANOIA) {
         return CX_INVALID_PARAMETER;
     }
     TRACE_BUFFER(inBuffer, inSize);
@@ -64,9 +65,6 @@ static __attribute__((always_inline, unused)) cx_err_t sha_256_hash(const uint8_
                                                                     size_t inSize,
                                                                     uint8_t* outBuffer,
                                                                     size_t outSize) {
-    if (inSize >= BUFFER_SIZE_PARANOIA) {
-        return CX_INVALID_PARAMETER;
-    }
     sha_256_context_t ctx;
     cx_err_t err = sha_256_init(&ctx);
     if (err != CX_OK) {
@@ -95,9 +93,11 @@ static __attribute__((always_inline, unused)) cx_err_t sha_512_init(sha_512_cont
 static __attribute__((always_inline, unused)) cx_err_t sha_512_append(sha_512_context_t* ctx,
                                                                       const uint8_t* inBuffer,
                                                                       size_t inSize) {
-    if (ctx->initialized_magic != HASH_CONTEXT_INITIALIZED_MAGIC) {
+    if (ctx->initialized_magic != HASH_CONTEXT_INITIALIZED_MAGIC ||
+        inSize >= BUFFER_SIZE_PARANOIA) {
         return CX_INVALID_PARAMETER;
     }
+
     TRACE_BUFFER(inBuffer, inSize);
     return cx_hash_no_throw(&ctx->cx_ctx.header,
                             0, /* Do not output the hash, yet */
@@ -126,9 +126,6 @@ static __attribute__((always_inline, unused)) cx_err_t sha_512_hash(const uint8_
                                                                     size_t inSize,
                                                                     uint8_t* outBuffer,
                                                                     size_t outSize) {
-    if (inSize >= BUFFER_SIZE_PARANOIA) {
-        return CX_INVALID_PARAMETER;
-    }
     sha_512_context_t ctx;
     cx_err_t err = sha_512_init(&ctx);
     if (err != CX_OK) {