libmbedtls: add interfaces in mbedtls for context memory operation

For integrating into OPTEE_OS, it needs add some interfaces:
1. add mbedtls_cipher_clone() for cipher to copy context between two
operations.
2. add mbedtls_cipher_setup_info() for cipher. cipher need to get its
"cipher_info" according the key length, while the key length is not an
input in allocate function. So, use a default key len in the beginning.
It need to reset the cipher info again in init function.
3. add mbedtls_cipher_cmac_setup() for cmac. This function is separate
from mbedtls_cipher_cmac_starts().
4. copy hmac context in md.

Acked-by: Etienne Carriere <etienne.carriere@linaro.org>
Signed-off-by: Edison Ai <edison.ai@arm.com>
Signed-off-by: Jens Wiklander <jens.wiklander@linaro.org>
diff --git a/lib/libmbedtls/mbedtls/include/mbedtls/cipher.h b/lib/libmbedtls/mbedtls/include/mbedtls/cipher.h
index 2083547..48670a4 100644
--- a/lib/libmbedtls/mbedtls/include/mbedtls/cipher.h
+++ b/lib/libmbedtls/mbedtls/include/mbedtls/cipher.h
@@ -394,6 +394,20 @@
  */
 void mbedtls_cipher_free( mbedtls_cipher_context_t *ctx );
 
+/**
+ * \brief           Clone the state of an cipher context
+ *
+ * \note            The two contexts must have been setup to the same type
+ *                  (cloning from AES to DES make no sense).
+ *
+ * \param dst       The destination context
+ * \param src       The context to be cloned
+ *
+ * \return          \c 0 on success,
+ *                  \c MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA on parameter failure.
+ */
+int mbedtls_cipher_clone( mbedtls_cipher_context_t *dst,
+                          const mbedtls_cipher_context_t *src );
 
 /**
  * \brief               This function initializes and fills the cipher-context
@@ -417,6 +431,17 @@
                           const mbedtls_cipher_info_t *cipher_info );
 
 /**
+ * \brief               setup the cipher info structure.
+ *
+ * \param ctx           cipher's context. Must have been initialised.
+ * \param cipher_info   cipher to use.
+ *
+ * \return              0 on success,
+ *                      MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA on parameter failure
+ */
+int mbedtls_cipher_setup_info( mbedtls_cipher_context_t *ctx, const mbedtls_cipher_info_t *cipher_info );
+
+/**
  * \brief        This function returns the block size of the given cipher.
  *
  * \param ctx    The context of the cipher. This must be initialized.
diff --git a/lib/libmbedtls/mbedtls/include/mbedtls/cipher_internal.h b/lib/libmbedtls/mbedtls/include/mbedtls/cipher_internal.h
index 57430a3..dc08f20 100644
--- a/lib/libmbedtls/mbedtls/include/mbedtls/cipher_internal.h
+++ b/lib/libmbedtls/mbedtls/include/mbedtls/cipher_internal.h
@@ -103,6 +103,9 @@
     /** Allocate a new context */
     void * (*ctx_alloc_func)( void );
 
+    /** Clone context **/
+    void (*ctx_clone_func)( void *dst, const void *src );
+
     /** Free the given context */
     void (*ctx_free_func)( void *ctx );
 
diff --git a/lib/libmbedtls/mbedtls/include/mbedtls/cmac.h b/lib/libmbedtls/mbedtls/include/mbedtls/cmac.h
index 0cffa1e..e2a4c44 100644
--- a/lib/libmbedtls/mbedtls/include/mbedtls/cmac.h
+++ b/lib/libmbedtls/mbedtls/include/mbedtls/cmac.h
@@ -69,6 +69,19 @@
 #endif /* !MBEDTLS_CMAC_ALT */
 
 /**
+ * \brief               Initialises and allocat cmac context memory
+ *                      Must be called with an initialized cipher context.
+ *
+ * \param ctx           The cipher context used for the CMAC operation, initialized
+ *                      as one of the following types: MBEDTLS_CIPHER_AES_128_ECB,
+ *                      MBEDTLS_CIPHER_AES_192_ECB, MBEDTLS_CIPHER_AES_256_ECB,
+ *                      or MBEDTLS_CIPHER_DES_EDE3_ECB.
+ * \return              \c 0 on success.
+ * \return              A cipher-specific error code on failure.
+ */
+int mbedtls_cipher_cmac_setup(mbedtls_cipher_context_t *ctx);
+
+/**
  * \brief               This function sets the CMAC key, and prepares to authenticate
  *                      the input data.
  *                      Must be called with an initialized cipher context.
diff --git a/lib/libmbedtls/mbedtls/library/cipher.c b/lib/libmbedtls/mbedtls/library/cipher.c
index 599cd36..e722313 100644
--- a/lib/libmbedtls/mbedtls/library/cipher.c
+++ b/lib/libmbedtls/mbedtls/library/cipher.c
@@ -179,6 +179,36 @@
     mbedtls_platform_zeroize( ctx, sizeof(mbedtls_cipher_context_t) );
 }
 
+int mbedtls_cipher_clone( mbedtls_cipher_context_t *dst,
+                          const mbedtls_cipher_context_t *src )
+{
+    if( dst == NULL || dst->cipher_info == NULL ||
+        src == NULL || src->cipher_info == NULL)
+    {
+        return( MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA );
+    }
+
+    dst->cipher_info = src->cipher_info;
+    dst->key_bitlen = src->key_bitlen;
+    dst->operation = src->operation;
+#if defined(MBEDTLS_CIPHER_MODE_WITH_PADDING)
+    dst->add_padding = src->add_padding;
+    dst->get_padding = src->get_padding;
+#endif
+    memcpy( dst->unprocessed_data, src->unprocessed_data, MBEDTLS_MAX_BLOCK_LENGTH );
+    dst->unprocessed_len = src->unprocessed_len;
+    memcpy( dst->iv, src->iv, MBEDTLS_MAX_IV_LENGTH );
+    dst->iv_size = src->iv_size;
+    if( dst->cipher_info->base->ctx_clone_func )
+        dst->cipher_info->base->ctx_clone_func( dst->cipher_ctx, src->cipher_ctx );
+
+#if defined(MBEDTLS_CMAC_C)
+    if( dst->cmac_ctx != NULL && src->cmac_ctx != NULL )
+        memcpy( dst->cmac_ctx, src->cmac_ctx, sizeof( mbedtls_cmac_context_t ) );
+#endif
+    return( 0 );
+}
+
 int mbedtls_cipher_setup( mbedtls_cipher_context_t *ctx, const mbedtls_cipher_info_t *cipher_info )
 {
     CIPHER_VALIDATE_RET( ctx != NULL );
@@ -206,6 +236,15 @@
     return( 0 );
 }
 
+int mbedtls_cipher_setup_info( mbedtls_cipher_context_t *ctx, const mbedtls_cipher_info_t *cipher_info )
+{
+    if( NULL == cipher_info || NULL == ctx )
+        return( MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA );
+
+    ctx->cipher_info = cipher_info;
+    return( 0 );
+}
+
 int mbedtls_cipher_setkey( mbedtls_cipher_context_t *ctx,
                            const unsigned char *key,
                            int key_bitlen,
diff --git a/lib/libmbedtls/mbedtls/library/cipher_wrap.c b/lib/libmbedtls/mbedtls/library/cipher_wrap.c
index b846f97..df55bbd 100644
--- a/lib/libmbedtls/mbedtls/library/cipher_wrap.c
+++ b/lib/libmbedtls/mbedtls/library/cipher_wrap.c
@@ -29,6 +29,8 @@
 #include MBEDTLS_CONFIG_FILE
 #endif
 
+#include <string.h>
+
 #if defined(MBEDTLS_CIPHER_C)
 
 #include "mbedtls/cipher_internal.h"
@@ -97,6 +99,11 @@
     return( ctx );
 }
 
+static void gcm_ctx_clone( void *dst, const void *src )
+{
+    memcpy( dst, src, sizeof( mbedtls_gcm_context ) );
+}
+
 static void gcm_ctx_free( void *ctx )
 {
     mbedtls_gcm_free( ctx );
@@ -116,6 +123,11 @@
     return( ctx );
 }
 
+static void ccm_ctx_clone( void *dst, const void *src )
+{
+    memcpy( dst, src, sizeof( mbedtls_ccm_context ) );
+}
+
 static void ccm_ctx_free( void *ctx )
 {
     mbedtls_ccm_free( ctx );
@@ -220,6 +232,11 @@
     return( aes );
 }
 
+static void aes_ctx_clone( void *dst, const void *src )
+{
+    memcpy( dst, src, sizeof( mbedtls_aes_context ) );
+}
+
 static void aes_ctx_free( void *ctx )
 {
     mbedtls_aes_free( (mbedtls_aes_context *) ctx );
@@ -250,6 +267,7 @@
     aes_setkey_enc_wrap,
     aes_setkey_dec_wrap,
     aes_ctx_alloc,
+    aes_ctx_clone,
     aes_ctx_free
 };
 
@@ -544,6 +562,7 @@
     gcm_aes_setkey_wrap,
     gcm_aes_setkey_wrap,
     gcm_ctx_alloc,
+    gcm_ctx_clone,
     gcm_ctx_free,
 };
 
@@ -613,6 +632,7 @@
     ccm_aes_setkey_wrap,
     ccm_aes_setkey_wrap,
     ccm_ctx_alloc,
+    ccm_ctx_clone,
     ccm_ctx_free,
 };
 
@@ -716,6 +736,11 @@
     return( ctx );
 }
 
+static void camellia_ctx_clone( void *dst, const void *src )
+{
+    memcpy( dst, src, sizeof( mbedtls_camellia_context ) );
+}
+
 static void camellia_ctx_free( void *ctx )
 {
     mbedtls_camellia_free( (mbedtls_camellia_context *) ctx );
@@ -746,6 +771,7 @@
     camellia_setkey_enc_wrap,
     camellia_setkey_dec_wrap,
     camellia_ctx_alloc,
+    camellia_ctx_clone,
     camellia_ctx_free
 };
 
@@ -919,6 +945,7 @@
     gcm_camellia_setkey_wrap,
     gcm_camellia_setkey_wrap,
     gcm_ctx_alloc,
+    gcm_ctx_clone,
     gcm_ctx_free,
 };
 
@@ -988,6 +1015,7 @@
     ccm_camellia_setkey_wrap,
     ccm_camellia_setkey_wrap,
     ccm_ctx_alloc,
+    ccm_ctx_clone,
     ccm_ctx_free,
 };
 
@@ -1497,6 +1525,11 @@
     return( des );
 }
 
+static void des_ctx_clone( void *dst, const void *src )
+{
+    memcpy( dst, src, sizeof( mbedtls_des_context ) );
+}
+
 static void des_ctx_free( void *ctx )
 {
     mbedtls_des_free( (mbedtls_des_context *) ctx );
@@ -1516,6 +1549,11 @@
     return( des3 );
 }
 
+static void des3_ctx_clone( void *dst, const void *src )
+{
+    memcpy( dst, src, sizeof( mbedtls_des3_context ) );
+}
+
 static void des3_ctx_free( void *ctx )
 {
     mbedtls_des3_free( (mbedtls_des3_context *) ctx );
@@ -1546,6 +1584,7 @@
     des_setkey_enc_wrap,
     des_setkey_dec_wrap,
     des_ctx_alloc,
+    des_ctx_clone,
     des_ctx_free
 };
 
@@ -1597,6 +1636,7 @@
     des3_set2key_enc_wrap,
     des3_set2key_dec_wrap,
     des3_ctx_alloc,
+    des3_ctx_clone,
     des3_ctx_free
 };
 
@@ -1648,6 +1688,7 @@
     des3_set3key_enc_wrap,
     des3_set3key_dec_wrap,
     des3_ctx_alloc,
+    des3_ctx_clone,
     des3_ctx_free
 };
 
@@ -1733,6 +1774,11 @@
     return( ctx );
 }
 
+static void blowfish_ctx_clone( void *dst, const void *src )
+{
+    memcpy( dst, src, sizeof( mbedtls_blowfish_context ) );
+}
+
 static void blowfish_ctx_free( void *ctx )
 {
     mbedtls_blowfish_free( (mbedtls_blowfish_context *) ctx );
@@ -1763,6 +1809,7 @@
     blowfish_setkey_wrap,
     blowfish_setkey_wrap,
     blowfish_ctx_alloc,
+    blowfish_ctx_clone,
     blowfish_ctx_free
 };
 
@@ -1849,6 +1896,11 @@
     return( ctx );
 }
 
+static void arc4_ctx_clone( void *dst, const void *src )
+{
+    memcpy( dst, src, sizeof( mbedtls_arc4_context ) );
+}
+
 static void arc4_ctx_free( void *ctx )
 {
     mbedtls_arc4_free( (mbedtls_arc4_context *) ctx );
@@ -1879,6 +1931,7 @@
     arc4_setkey_wrap,
     arc4_setkey_wrap,
     arc4_ctx_alloc,
+    arc4_ctx_clone,
     arc4_ctx_free
 };
 
@@ -2075,6 +2128,12 @@
     return( (void *) 1 );
 }
 
+static void null_ctx_clone( void *dst, const void *src )
+{
+    ((void) dst);
+    ((void) src);
+}
+
 static void null_ctx_free( void *ctx )
 {
     ((void) ctx);
@@ -2104,6 +2163,7 @@
     null_setkey,
     null_setkey,
     null_ctx_alloc,
+    null_ctx_clone,
     null_ctx_free
 };
 
diff --git a/lib/libmbedtls/mbedtls/library/cmac.c b/lib/libmbedtls/mbedtls/library/cmac.c
index 37b66b5..9a9b6a2 100644
--- a/lib/libmbedtls/mbedtls/library/cmac.c
+++ b/lib/libmbedtls/mbedtls/library/cmac.c
@@ -53,7 +53,6 @@
 
 #include <string.h>
 
-
 #if defined(MBEDTLS_PLATFORM_C)
 #include "mbedtls/platform.h"
 #else
@@ -199,11 +198,26 @@
     }
 }
 
+int mbedtls_cipher_cmac_setup(mbedtls_cipher_context_t *ctx)
+{
+    mbedtls_cmac_context_t *cmac_ctx;
+
+    /* Allocated and initialise in the cipher context memory for the CMAC
+     * context */
+    cmac_ctx = mbedtls_calloc( 1, sizeof( mbedtls_cmac_context_t ) );
+    if( cmac_ctx == NULL )
+        return( MBEDTLS_ERR_CIPHER_ALLOC_FAILED );
+
+    ctx->cmac_ctx = cmac_ctx;
+
+    mbedtls_platform_zeroize( cmac_ctx->state, sizeof( cmac_ctx->state ) );
+    return 0;
+}
+
 int mbedtls_cipher_cmac_starts( mbedtls_cipher_context_t *ctx,
                                 const unsigned char *key, size_t keybits )
 {
     mbedtls_cipher_type_t type;
-    mbedtls_cmac_context_t *cmac_ctx;
     int retval;
 
     if( ctx == NULL || ctx->cipher_info == NULL || key == NULL )
@@ -226,17 +240,11 @@
             return( MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA );
     }
 
-    /* Allocated and initialise in the cipher context memory for the CMAC
-     * context */
-    cmac_ctx = mbedtls_calloc( 1, sizeof( mbedtls_cmac_context_t ) );
-    if( cmac_ctx == NULL )
-        return( MBEDTLS_ERR_CIPHER_ALLOC_FAILED );
+    /* Check if cmac ctx had been allocated by mbedtls_cipher_cmac_setup() */
+    if( ctx->cmac_ctx != NULL )
+        return 0;
 
-    ctx->cmac_ctx = cmac_ctx;
-
-    mbedtls_platform_zeroize( cmac_ctx->state, sizeof( cmac_ctx->state ) );
-
-    return 0;
+    return mbedtls_cipher_cmac_setup( ctx );
 }
 
 int mbedtls_cipher_cmac_update( mbedtls_cipher_context_t *ctx,
diff --git a/lib/libmbedtls/mbedtls/library/md.c b/lib/libmbedtls/mbedtls/library/md.c
index 7d1208c..f0e85f0 100644
--- a/lib/libmbedtls/mbedtls/library/md.c
+++ b/lib/libmbedtls/mbedtls/library/md.c
@@ -209,6 +209,9 @@
 
     dst->md_info->clone_func( dst->md_ctx, src->md_ctx );
 
+    if( dst->hmac_ctx != NULL && src->hmac_ctx != NULL )
+        memcpy( dst->hmac_ctx, src->hmac_ctx, 2 * src->md_info->block_size );
+
     return( 0 );
 }