From f89c9bec9aea680839ae10aef2d9ac627a9c08c3 Mon Sep 17 00:00:00 2001
From: manastasova <44320407+manastasova@users.noreply.github.com>
Date: Fri, 20 Sep 2024 15:21:24 -0700
Subject: [PATCH] Add return checks on SHA3 functions in ML-KEM (#1859)

### Issues:
Resolves #P155314914

### Description of changes:
Add comments on omitted return code checks in ML-KEM

### Testing:
`./crypto/crypto_test`

By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 license and the ISC license.
---
 .../ml_kem/ml_kem_ref/symmetric-shake.c         | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/crypto/fipsmodule/ml_kem/ml_kem_ref/symmetric-shake.c b/crypto/fipsmodule/ml_kem/ml_kem_ref/symmetric-shake.c
index 859e92fd98..02f2f0f70a 100644
--- a/crypto/fipsmodule/ml_kem/ml_kem_ref/symmetric-shake.c
+++ b/crypto/fipsmodule/ml_kem/ml_kem_ref/symmetric-shake.c
@@ -25,7 +25,11 @@ void kyber_shake128_absorb(KECCAK1600_CTX *ctx,
   extseed[KYBER_SYMBYTES+0] = x;
   extseed[KYBER_SYMBYTES+1] = y;
 
+  // Return code checks can be omitted
+  // SHAKE_Init always returns 1 when called with correct block size value
   SHAKE_Init(ctx, SHAKE128_BLOCKSIZE);
+
+  // SHA3_Update always returns 1 on first call of sizeof(extseed) (34 bytes)
   SHA3_Update(ctx, extseed, sizeof(extseed));
 }
 
@@ -43,6 +47,8 @@ void kyber_shake128_absorb(KECCAK1600_CTX *ctx,
 **************************************************/
 void kyber_shake128_squeeze(KECCAK1600_CTX *ctx, uint8_t *out, int nblocks)
 {
+  // Return code checks can be omitted
+  // SHAKE_Final always returns 1
   SHAKE_Final(out, ctx, nblocks * SHAKE128_BLOCKSIZE);
 }
 
@@ -64,6 +70,8 @@ void kyber_shake256_prf(uint8_t *out, size_t outlen, const uint8_t key[KYBER_SYM
   memcpy(extkey, key, KYBER_SYMBYTES);
   extkey[KYBER_SYMBYTES] = nonce;
 
+  // Return code checks can be omitted
+  // SHAKE256 never returns NULL when the internal SHAKE_Init is called with correct block size value
   SHAKE256(extkey, sizeof(extkey), out, outlen);
 }
 
@@ -81,8 +89,17 @@ void kyber_shake256_prf(uint8_t *out, size_t outlen, const uint8_t key[KYBER_SYM
 void kyber_shake256_rkprf(ml_kem_params *params, uint8_t out[KYBER_SSBYTES], const uint8_t key[KYBER_SYMBYTES], const uint8_t *input)
 {
   KECCAK1600_CTX ctx;
+
+  // Return code checks can be omitted
+  // SHAKE_Init always returns 1 when called with correct block size value
   SHAKE_Init(&ctx, SHAKE256_BLOCKSIZE);
+
+  // SHA3_Update always returns 1 on first call of KYBER_SYMBYTES (32 bytes)
   SHA3_Update(&ctx, key, KYBER_SYMBYTES);
+
+  // SHA3_Update always returns 1 processing all data blocks that don't need pad
   SHA3_Update(&ctx, input, params->ciphertext_bytes);
+
+  // SHAKE_Final always returns 1
   SHAKE_Final(out, &ctx, KYBER_SSBYTES);
 }