diff --git a/library/cmac.c b/library/cmac.c
index 39ebb8709..4c25a67d0 100644
--- a/library/cmac.c
+++ b/library/cmac.c
@@ -173,12 +173,6 @@ int mbedtls_cmac_setkey( mbedtls_cmac_context *ctx,
     if( cipher_info == NULL )
         return( MBEDTLS_ERR_CMAC_BAD_INPUT );
 
-    ctx->K1 = mbedtls_calloc( cipher_info->block_size, sizeof( unsigned char ) );
-    ctx->K2 = mbedtls_calloc( cipher_info->block_size, sizeof( unsigned char ) );
-
-    if(ctx->K1 == NULL || ctx->K2 == NULL )
-        return MBEDTLS_ERR_CMAC_ALLOC_FAILED;
-
     mbedtls_cipher_free( &ctx->cipher_ctx );
 
     if( ( ret = mbedtls_cipher_setup( &ctx->cipher_ctx, cipher_info ) ) != 0 )
@@ -190,6 +184,16 @@ int mbedtls_cmac_setkey( mbedtls_cmac_context *ctx,
         return( ret );
     }
 
+    ctx->K1 = mbedtls_calloc( cipher_info->block_size, sizeof( unsigned char ) );
+    ctx->K2 = mbedtls_calloc( cipher_info->block_size, sizeof( unsigned char ) );
+
+    if( ctx->K1 == NULL || ctx->K2 == NULL )
+    {
+        mbedtls_free(ctx->K1);
+        mbedtls_free(ctx->K2);
+        return( MBEDTLS_ERR_CMAC_ALLOC_FAILED );
+    }
+
     return( cmac_generate_subkeys( ctx ) );
 }