Clear context state if previous operation failed.

Signed-off-by: Mateusz Starzyk <mateusz.starzyk@mobica.com>
diff --git a/library/ccm.c b/library/ccm.c
index a2ca986..34531a4 100644
--- a/library/ccm.c
+++ b/library/ccm.c
@@ -122,7 +122,10 @@
         ctx->y[i] ^= ctx->b[i];                                                               \
                                                                                               \
     if( ( ret = mbedtls_cipher_update( &ctx->cipher_ctx, ctx->y, 16, ctx->y, &olen ) ) != 0 ) \
-        return( ret );
+    {                                                                                         \
+        ctx->state |= CCM_STATE__ERROR;                                                       \
+        return( ret );                                                                        \
+    }                                                                                         \
 
 /*
  * Encrypt or decrypt a partial block with CTR
@@ -135,6 +138,7 @@
         if( ( ret = mbedtls_cipher_update( &ctx->cipher_ctx, ctx->ctr,         \
                                            16, ctx->b, &olen ) ) != 0 )        \
         {                                                                      \
+            ctx->state |= CCM_STATE__ERROR;                                    \
             return( ret );                                                     \
         }                                                                      \
                                                                                \
@@ -145,6 +149,7 @@
 #define CCM_STATE__CLEAR                0
 #define CCM_STATE__STARTED              0x0001
 #define CCM_STATE__LENGHTS_SET          0x0002
+#define CCM_STATE__ERROR                0x0004
 
 static void mbedtls_ccm_clear_state(mbedtls_ccm_context *ctx) {
     ctx->state = CCM_STATE__CLEAR;
@@ -175,7 +180,10 @@
         ctx->b[15-i] = (unsigned char)( len_left & 0xFF );
 
     if( len_left > 0 )
+    {
+        ctx->state |= CCM_STATE__ERROR;
         return( MBEDTLS_ERR_CCM_BAD_INPUT );
+    }
 
     /* Start CBC-MAC with first block*/
     UPDATE_CBC_MAC;
@@ -188,12 +196,15 @@
                         const unsigned char *iv,
                         size_t iv_len )
 {
-    int ret;
-
     /* Also implies q is within bounds */
     if( iv_len < 7 || iv_len > 13 )
         return( MBEDTLS_ERR_CCM_BAD_INPUT );
 
+    if( ctx->state & CCM_STATE__ERROR )
+    {
+        mbedtls_ccm_clear_state(ctx);
+    }
+
     ctx->mode = mode;
     ctx->q = 16 - 1 - (unsigned char) iv_len;
 
@@ -230,9 +241,7 @@
     memcpy( ctx->b + 1, iv, iv_len );
 
     ctx->state |= CCM_STATE__STARTED;
-    ret = mbedtls_ccm_calculate_first_block(ctx);
-
-    return ret;
+    return mbedtls_ccm_calculate_first_block(ctx);
 }
 
 int mbedtls_ccm_set_lengths( mbedtls_ccm_context *ctx,
@@ -240,8 +249,6 @@
                              size_t plaintext_len,
                              size_t tag_len )
 {
-    int ret;
-
     /*
      * Check length requirements: SP800-38C A.1
      * Additional requirement: a < 2^16 - 2^8 to simplify the code.
@@ -255,6 +262,11 @@
     if( total_ad_len >= 0xFF00 )
         return( MBEDTLS_ERR_CCM_BAD_INPUT );
 
+    if( ctx->state & CCM_STATE__ERROR )
+    {
+        mbedtls_ccm_clear_state(ctx);
+    }
+
     /*
      * First block B_0:
      * 0        .. 0        flags           - set by: mbedtls_ccm_starts() and mbedtls_ccm_set_lenghts()
@@ -273,9 +285,7 @@
     ctx->plaintext_len = plaintext_len;
 
     ctx->state |= CCM_STATE__LENGHTS_SET;
-    ret = mbedtls_ccm_calculate_first_block(ctx);
-
-    return ret;
+    return mbedtls_ccm_calculate_first_block(ctx);
 }
 
 int mbedtls_ccm_update_ad( mbedtls_ccm_context *ctx,