Add checks to crypto operation calls in MCUboot

Signed-off-by: Karambite <01sudesh.e@gmail.com>

Author: Karambite

boot/bootutil/src/encrypted.c

@@ -726,7 +726,9 @@ boot_enc_encrypt(struct enc_key_data *enc_state, int slot, uint32_t off,
     nonce[15] = (uint8_t)off;
 
     assert(enc->valid == 1);
-    bootutil_aes_ctr_encrypt(&enc->aes_ctr, nonce, buf, sz, blk_off, buf);
+    if (bootutil_aes_ctr_encrypt(&enc->aes_ctr, nonce, buf, sz, blk_off, buf)) {
+        bootutil_aes_ctr_drop(&enc->aes_ctr);
+    }
 }
 
 void
@@ -749,7 +751,9 @@ boot_enc_decrypt(struct enc_key_data *enc_state, int slot, uint32_t off,
     nonce[15] = (uint8_t)off;
 
     assert(enc->valid == 1);
-    bootutil_aes_ctr_decrypt(&enc->aes_ctr, nonce, buf, sz, blk_off, buf);
+    if (bootutil_aes_ctr_decrypt(&enc->aes_ctr, nonce, buf, sz, blk_off, buf)) {
+        bootutil_aes_ctr_drop(&enc->aes_ctr);
+    }
 }
 
 /**

boot/bootutil/src/image_rsa.c

@@ -94,9 +94,15 @@ pss_mgf1(uint8_t *mask, const uint8_t *hash)
 
     while (count > 0) {
         bootutil_sha_init(&ctx);
-        bootutil_sha_update(&ctx, hash, PSS_HLEN);
-        bootutil_sha_update(&ctx, counter, 4);
-        bootutil_sha_finish(&ctx, htmp);
+        if (bootutil_sha_update(&ctx, hash, PSS_HLEN)) {
+            goto out;
+        }
+        if (bootutil_sha_update(&ctx, counter, 4)) {
+            goto out;
+        }
+        if(bootutil_sha_finish(&ctx, htmp)){
+            goto out;
+        }
 
         counter[3]++;
 
@@ -109,6 +115,7 @@ pss_mgf1(uint8_t *mask, const uint8_t *hash)
         count -= bytes;
     }
 
+out:
     bootutil_sha_drop(&ctx);
 }
 
@@ -222,17 +229,25 @@ bootutil_cmp_rsasig(bootutil_rsa_context *ctx, uint8_t *hash, uint32_t hlen,
 
     /* Step 13.  Let H' = Hash(M') */
     bootutil_sha_init(&shactx);
-    bootutil_sha_update(&shactx, pss_zeros, 8);
-    bootutil_sha_update(&shactx, hash, PSS_HLEN);
-    bootutil_sha_update(&shactx, &db_mask[PSS_MASK_SALT_POS], PSS_SLEN);
-    bootutil_sha_finish(&shactx, h2);
-    bootutil_sha_drop(&shactx);
+    if (bootutil_sha_update(&shactx, pss_zeros, 8)) {
+        goto out;
+    }
+    if (bootutil_sha_update(&shactx, hash, PSS_HLEN)) {
+        goto out;
+    }
+    if (bootutil_sha_update(&shactx, &db_mask[PSS_MASK_SALT_POS], PSS_SLEN)) {
+        goto out;
+    }
+    if (bootutil_sha_finish(&shactx, h2)) {
+        goto out;
+    }
 
     /* Step 14.  If H = H', output "consistent".  Otherwise, output
      * "inconsistent". */
     FIH_CALL(boot_fih_memequal, fih_rc, h2, &em[PSS_HASH_OFFSET], PSS_HLEN);
 
 out:
+    bootutil_sha_drop(&shactx);
     FIH_RET(fih_rc);
 }
 

boot/bootutil/src/image_validate.c

@@ -140,7 +140,11 @@ bootutil_img_hash(struct boot_loader_state *state,
     /* in some cases (split image) the hash is seeded with data from
      * the loader image */
     if (seed && (seed_len > 0)) {
-        bootutil_sha_update(&sha_ctx, seed, seed_len);
+        rc = bootutil_sha_update(&sha_ctx, seed, seed_len);
+        if (rc){
+            bootutil_sha_drop(&sha_ctx);
+            return rc;
+        }
     }
 
     /* Hash is computed over image header and image itself. */
@@ -155,12 +159,21 @@ bootutil_img_hash(struct boot_loader_state *state,
     /* No chunk loading, storage is mapped to address space and can
      * be directly given to hashing function.
      */
-    bootutil_sha_update(&sha_ctx, (void *)flash_area_get_off(fap), size);
+    rc = bootutil_sha_update(&sha_ctx, (void *)flash_area_get_off(fap), size);
+    if (rc){
+        bootutil_sha_drop(&sha_ctx);
+        return rc;
+    }
 #else /* MCUBOOT_HASH_STORAGE_DIRECTLY */
 #ifdef MCUBOOT_RAM_LOAD
-    bootutil_sha_update(&sha_ctx,
+    rc = bootutil_sha_update(&sha_ctx,
                         (void*)(IMAGE_RAM_BASE + hdr->ih_load_addr),
                         size);
+    if (rc){
+        bootutil_sha_drop(&sha_ctx);
+        return rc;
+    }
+    
 #else
     for (off = 0; off < size; off += blk_sz) {
         blk_sz = size - off;
@@ -202,14 +215,18 @@ bootutil_img_hash(struct boot_loader_state *state,
             }
         }
 #endif
-        bootutil_sha_update(&sha_ctx, tmp_buf, blk_sz);
+        rc = bootutil_sha_update(&sha_ctx, tmp_buf, blk_sz);
+        if (rc){
+            bootutil_sha_drop(&sha_ctx);
+            return rc;
+        }
     }
 #endif /* MCUBOOT_RAM_LOAD */
 #endif /* MCUBOOT_HASH_STORAGE_DIRECTLY */
-    bootutil_sha_finish(&sha_ctx, hash_result);
+    rc = bootutil_sha_finish(&sha_ctx, hash_result);
     bootutil_sha_drop(&sha_ctx);
 
-    return 0;
+    return rc;
 }
 #endif
 
@@ -287,8 +304,12 @@ bootutil_find_key(uint8_t *keyhash, uint8_t keyhash_len)
     for (i = 0; i < bootutil_key_cnt; i++) {
         key = &bootutil_keys[i];
         bootutil_sha_init(&sha_ctx);
-        bootutil_sha_update(&sha_ctx, key->key, *key->len);
-        bootutil_sha_finish(&sha_ctx, hash);
+        if (bootutil_sha_update(&sha_ctx, key->key, *key->len)){
+            break;
+        }
+        if (bootutil_sha_finish(&sha_ctx, hash)){
+            break;
+        }
         if (!memcmp(hash, keyhash, keyhash_len)) {
             bootutil_sha_drop(&sha_ctx);
             return i;
@@ -310,9 +331,16 @@ bootutil_find_key(uint8_t image_index, uint8_t *key, uint16_t key_len)
     FIH_DECLARE(fih_rc, FIH_FAILURE);
 
     bootutil_sha_init(&sha_ctx);
-    bootutil_sha_update(&sha_ctx, key, key_len);
-    bootutil_sha_finish(&sha_ctx, hash);
+    rc = bootutil_sha_update(&sha_ctx, key, key_len);
+    if (rc){
+        bootutil_sha_drop(&sha_ctx);
+        return rc;
+    }
+    rc = bootutil_sha_finish(&sha_ctx, hash);
     bootutil_sha_drop(&sha_ctx);
+    if (rc){   
+        return rc;
+    }
 
     rc = boot_retrieve_public_key_hash(image_index, key_hash, &key_hash_size);
     if (rc) {