diff --git a/include/mbedtls/ccm.h b/include/mbedtls/ccm.h
index f9f8000fba..813959be0b 100644
--- a/include/mbedtls/ccm.h
+++ b/include/mbedtls/ccm.h
@@ -81,7 +81,10 @@ typedef struct mbedtls_ccm_context
     unsigned char MBEDTLS_PRIVATE(y)[16];    /*!< The Y working buffer */
     unsigned char MBEDTLS_PRIVATE(ctr)[16];  /*!< The counter buffer */
     unsigned char MBEDTLS_PRIVATE(q);        /*!< The Q working value */
-    size_t MBEDTLS_PRIVATE(plaintext_len);   /*!< The counter buffer */
+    size_t MBEDTLS_PRIVATE(plaintext_len);   /*!< Total plaintext length */
+    size_t MBEDTLS_PRIVATE(add_len);         /*!< Total authentication data length */
+    size_t MBEDTLS_PRIVATE(tag_len);         /*!< Total tag length */
+    size_t MBEDTLS_PRIVATE(processed);       /*!< How many bytes of input data were processed (chunked input) */
     int MBEDTLS_PRIVATE(mode);               /*!< The operation to perform:
                                                 #MBEDTLS_CCM_ENCRYPT or
                                                 #MBEDTLS_CCM_DECRYPT or
diff --git a/library/ccm.c b/library/ccm.c
index 34531a4162..36b1e91583 100644
--- a/library/ccm.c
+++ b/library/ccm.c
@@ -172,10 +172,20 @@ static int mbedtls_ccm_calculate_first_block(mbedtls_ccm_context *ctx)
 
     /*
      * First block B_0:
-     * 0        .. 0        flags           - set by: mbedtls_ccm_starts() and mbedtls_ccm_set_lenghts()
+     * 0        .. 0        flags
      * 1        .. iv_len   nonce (aka iv)  - set by: mbedtls_ccm_starts()
-     * iv_len+1 .. 15       length          - set by: mbedtls_ccm_calculate_first_block()
+     * iv_len+1 .. 15       length
+     *
+     * With flags as (bits):
+     * 7        0
+     * 6        add present?
+     * 5 .. 3   (t - 2) / 2
+     * 2 .. 0   q - 1
      */
+    ctx->b[0] |= ( ctx->add_len > 0 ) << 6;
+    ctx->b[0] |= ( ( ctx->tag_len - 2 ) / 2 ) << 3;
+    ctx->b[0] |= ctx->q - 1;
+
     for( i = 0, len_left = ctx->plaintext_len; i < ctx->q; i++, len_left >>= 8 )
         ctx->b[15-i] = (unsigned char)( len_left & 0xFF );
 
@@ -225,19 +235,8 @@ int mbedtls_ccm_starts( mbedtls_ccm_context *ctx,
     ctx->ctr[15] = 1;
 
     /*
-     * First block B_0:
-     * 0        .. 0        flags           - set by: mbedtls_ccm_starts() and mbedtls_ccm_set_lenghts()
-     * 1        .. iv_len   nonce (aka iv)  - set by: mbedtls_ccm_starts()
-     * iv_len+1 .. 15       length          - set by: mbedtls_ccm_calculate_first_block()
-     *
-     * With flags as (bits):
-     * 7        0
-     * 6        add present?                - set by: mbedtls_ccm_set_lengths()
-     * 5 .. 3   (t - 2) / 2                 - set by: mbedtls_ccm_set_lengths()
-     * 2 .. 0   q - 1                       - set by: mbedtls_ccm_starts()
+     * See mbedtls_ccm_calculate_first_block() for B block layout description
      */
-    ctx->b[0] |= ctx->q - 1;
-
     memcpy( ctx->b + 1, iv, iv_len );
 
     ctx->state |= CCM_STATE__STARTED;
@@ -267,22 +266,10 @@ int mbedtls_ccm_set_lengths( mbedtls_ccm_context *ctx,
         mbedtls_ccm_clear_state(ctx);
     }
 
-    /*
-     * First block B_0:
-     * 0        .. 0        flags           - set by: mbedtls_ccm_starts() and mbedtls_ccm_set_lenghts()
-     * 1        .. iv_len   nonce (aka iv)  - set by: mbedtls_ccm_starts()
-     * iv_len+1 .. 15       length          - set by: mbedtls_ccm_calculate_first_block()
-     *
-     * With flags as (bits):
-     * 7        0
-     * 6        add present?                - set by: mbedtls_ccm_set_lengths()
-     * 5 .. 3   (t - 2) / 2                 - set by: mbedtls_ccm_set_lengths()
-     * 2 .. 0   q - 1                       - set by: mbedtls_ccm_starts()
-     */
-    ctx->b[0] |= ( total_ad_len > 0 ) << 6;
-    ctx->b[0] |= ( ( tag_len - 2 ) / 2 ) << 3;
-
     ctx->plaintext_len = plaintext_len;
+    ctx->add_len = total_ad_len;
+    ctx->tag_len = tag_len;
+    ctx->processed = 0;
 
     ctx->state |= CCM_STATE__LENGHTS_SET;
     return mbedtls_ccm_calculate_first_block(ctx);