From aa79ba108a83e6f004cab4d75dea485ccec91154 Mon Sep 17 00:00:00 2001
From: Dong Heng <dongheng@espressif.com>
Date: Mon, 30 Dec 2019 19:31:17 +0800
Subject: [PATCH] feat(util): support SHA224/SHA384 calculation although call
 sha256/sha512 APIs

---
 .../crypto_hash_sha256_mbedtls.c                   |  1 +
 .../crypto_hash_sha512_mbedtls.c                   |  1 +
 components/util/include/esp_sha.h                  |  4 ++++
 components/util/src/sha256.c                       | 12 ++++++++----
 components/util/src/sha512.c                       | 14 +++++++++-----
 5 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/components/libsodium/port/crypto_hash_mbedtls/crypto_hash_sha256_mbedtls.c b/components/libsodium/port/crypto_hash_mbedtls/crypto_hash_sha256_mbedtls.c
index 9ffd28c6..960e2bda 100644
--- a/components/libsodium/port/crypto_hash_mbedtls/crypto_hash_sha256_mbedtls.c
+++ b/components/libsodium/port/crypto_hash_mbedtls/crypto_hash_sha256_mbedtls.c
@@ -46,6 +46,7 @@ static void sha256_libsodium_to_mbedtls(mbedtls_sha256_context *mb_ctx, crypto_h
     memcpy(mb_ctx->total, &ls_state->count, sizeof(mb_ctx->total));
     memcpy(mb_ctx->state, ls_state->state, sizeof(mb_ctx->state));
     memcpy(mb_ctx->buffer, ls_state->buf, sizeof(mb_ctx->buffer));
+    mb_ctx->is224 = 0;
 }
 
 int
diff --git a/components/libsodium/port/crypto_hash_mbedtls/crypto_hash_sha512_mbedtls.c b/components/libsodium/port/crypto_hash_mbedtls/crypto_hash_sha512_mbedtls.c
index acb03849..b005f0ce 100644
--- a/components/libsodium/port/crypto_hash_mbedtls/crypto_hash_sha512_mbedtls.c
+++ b/components/libsodium/port/crypto_hash_mbedtls/crypto_hash_sha512_mbedtls.c
@@ -47,6 +47,7 @@ static void sha512_libsodium_to_mbedtls(mbedtls_sha512_context *mb_ctx, crypto_h
     memcpy(mb_ctx->total, ls_state->count, sizeof(mb_ctx->total));
     memcpy(mb_ctx->state, ls_state->state, sizeof(mb_ctx->state));
     memcpy(mb_ctx->buffer, ls_state->buf, sizeof(mb_ctx->buffer));
+    mb_ctx->is384 = 0;
 }
 
 int
diff --git a/components/util/include/esp_sha.h b/components/util/include/esp_sha.h
index 52a81a10..238d1be1 100644
--- a/components/util/include/esp_sha.h
+++ b/components/util/include/esp_sha.h
@@ -31,12 +31,16 @@ typedef struct {
     uint32_t        total[2];
     uint32_t        state[8];
     uint8_t         buffer[64];
+
+    int             is224;
 } esp_sha256_t;
 
 typedef struct {
     uint64_t        total[2];
     uint64_t        state[8];
     uint8_t         buffer[128];
+
+    int             is384;
 } esp_sha512_t;
 
 typedef esp_sha256_t esp_sha224_t;
diff --git a/components/util/src/sha256.c b/components/util/src/sha256.c
index cc6de2bc..3b658beb 100644
--- a/components/util/src/sha256.c
+++ b/components/util/src/sha256.c
@@ -108,6 +108,8 @@ int esp_sha256_init(esp_sha256_t *ctx)
 {
     util_assert(ctx);
 
+    ctx->is224 = 0;
+
     ctx->total[0] = 0;
     ctx->total[1] = 0;
 
@@ -127,6 +129,8 @@ int esp_sha224_init(esp_sha224_t *ctx)
 {
     util_assert(ctx);
 
+    ctx->is224 = 1;
+
     ctx->total[0] = 0;
     ctx->total[1] = 0;
 
@@ -230,19 +234,19 @@ int esp_sha224_finish(esp_sha224_t *ctx, void *dest)
     ESP_PUT_BE32(out + 20, ctx->state[5]);
     ESP_PUT_BE32(out + 24, ctx->state[6]);
 
+    if (!ctx->is224) {
+        ESP_PUT_BE32(out + 28, ctx->state[7]);
+    }
+
     return( 0 );
 }
 
 int esp_sha256_finish(esp_sha256_t *ctx, void *dest)
 {
-    uint8_t *out = (uint8_t *)dest;
-
     util_assert(ctx);
     util_assert(dest);
 
     esp_sha224_finish(ctx, dest);
 
-    ESP_PUT_BE32(out + 28, ctx->state[7]);
-
     return 0;
 }
diff --git a/components/util/src/sha512.c b/components/util/src/sha512.c
index 173a7c7e..61faf2c1 100644
--- a/components/util/src/sha512.c
+++ b/components/util/src/sha512.c
@@ -159,6 +159,8 @@ int esp_sha512_init(esp_sha512_t *ctx)
 {
     util_assert(ctx);
 
+    ctx->is384 = 0;
+
     ctx->total[0] = 0;
     ctx->total[1] = 0;
 
@@ -178,6 +180,8 @@ int esp_sha384_init(esp_sha384_t *ctx)
 {
     util_assert(ctx);
 
+    ctx->is384 = 1;
+
     ctx->total[0] = 0;
     ctx->total[1] = 0;
 
@@ -279,20 +283,20 @@ int esp_sha384_finish(esp_sha384_t *ctx, void *dest)
     PUT_UINT64_BE(ctx->state[4], output, 32);
     PUT_UINT64_BE(ctx->state[5], output, 40);
 
+    if (!ctx->is384) {
+        PUT_UINT64_BE(ctx->state[6], output, 48);
+        PUT_UINT64_BE(ctx->state[7], output, 56);
+    }
+
     return 0;
 }
 
 int esp_sha512_finish(esp_sha512_t *ctx, void *dest)
 {
-    uint8_t *output = (uint8_t *)dest;
-
     util_assert(ctx);
     util_assert(dest);
 
     esp_sha384_finish(ctx, dest);
 
-    PUT_UINT64_BE(ctx->state[6], output, 48);
-    PUT_UINT64_BE(ctx->state[7], output, 56);
-
     return 0;
 }