From 20bac2fbe4ac0e002f4e622634cd858925c258c7 Mon Sep 17 00:00:00 2001
From: Mateusz Starzyk <mateusz.starzyk@mobica.com>
Date: Mon, 12 Jul 2021 14:52:44 +0200
Subject: [PATCH] Fix chunked ccm update.

Signed-off-by: Mateusz Starzyk <mateusz.starzyk@mobica.com>
---
 library/ccm.c | 71 ++++++++++++++++++++++++++++-----------------------
 1 file changed, 39 insertions(+), 32 deletions(-)

diff --git a/library/ccm.c b/library/ccm.c
index 4f7ebfa82..5450e408c 100644
--- a/library/ccm.c
+++ b/library/ccm.c
@@ -373,42 +373,49 @@ int mbedtls_ccm_update( mbedtls_ccm_context *ctx,
         ctx->processed += use_len;
         memcpy( ctx->b + offset, input, use_len );
 
+        if( ctx->mode == MBEDTLS_CCM_ENCRYPT || \
+            ctx->mode == MBEDTLS_CCM_STAR_ENCRYPT )
+        {
+            if( use_len + offset == 16 || ctx->processed == ctx->plaintext_len )
+            {
+                UPDATE_CBC_MAC;
+            }
+            ret = mbedtls_ccm_crypt( ctx, offset, use_len, ctx->b + offset, output );
+            if( ret != 0 )
+                return ret;
+        }
+
+        if( ctx->mode == MBEDTLS_CCM_DECRYPT || \
+            ctx->mode == MBEDTLS_CCM_STAR_DECRYPT )
+        {
+            ret = mbedtls_ccm_crypt( ctx, offset, use_len, ctx->b + offset, output );
+            if( ret != 0 )
+                return ret;
+
+            for( i = 0; i < use_len; i++ )
+                ctx->y[i + offset] ^= output[i];
+
+            if( use_len + offset == 16 || ctx->processed == ctx->plaintext_len )
+            {
+                if( ( ret = mbedtls_cipher_update( &ctx->cipher_ctx, ctx->y, 16, ctx->y, &olen ) ) != 0 )
+                {
+                    ctx->state |= CCM_STATE__ERROR;
+                    return( ret );
+                }
+            }
+        }
+
         if( use_len + offset == 16 || ctx->processed == ctx->plaintext_len )
         {
-            if( ctx->mode == MBEDTLS_CCM_ENCRYPT || \
-                ctx->mode == MBEDTLS_CCM_STAR_ENCRYPT )
-            {
-                UPDATE_CBC_MAC;
-                ret = mbedtls_ccm_crypt( ctx, 0, use_len, ctx->b, output );
-                if( ret != 0 )
-                    return ret;
-                memset( ctx->b, 0, 16 );
-            }
-
-            if( ctx->mode == MBEDTLS_CCM_DECRYPT || \
-                ctx->mode == MBEDTLS_CCM_STAR_DECRYPT )
-            {
-                ret = mbedtls_ccm_crypt( ctx, 0, use_len, ctx->b, output );
-                if( ret != 0 )
-                    return ret;
-                memset( ctx->b, 0, 16 );
-                memcpy( ctx->b, output, use_len );
-                UPDATE_CBC_MAC;
-                memset( ctx->b, 0, 16 );
-            }
-
-            input_len -= use_len;
-            input += use_len;
-            output += use_len;
-
-            /*
-            * Increment counter.
-            * No need to check for overflow thanks to the length check above.
-            */
             for( i = 0; i < ctx->q; i++ )
-                if( ++(ctx->ctr)[15-i] != 0 )
-                    break;
+            if( ++(ctx->ctr)[15-i] != 0 )
+                break;
+            memset( ctx->b, 0, 16 );
         }
+
+        input_len -= use_len;
+        input += use_len;
+        output += use_len;
     }
 
     return 0;