block_cipher: add setkey()

Signed-off-by: Manuel Pégourié-Gonnard <manuel.pegourie-gonnard@arm.com>
diff --git a/library/block_cipher.c b/library/block_cipher.c
index 08364d7..2146c6c 100644
--- a/library/block_cipher.c
+++ b/library/block_cipher.c
@@ -66,4 +66,25 @@
     }
 }
 
+int mbedtls_block_cipher_setkey(mbedtls_block_cipher_context_t *ctx,
+                                const unsigned char *key,
+                                unsigned key_bitlen)
+{
+    switch (ctx->id) {
+#if defined(MBEDTLS_AES_C)
+        case MBEDTLS_BLOCK_CIPHER_ID_AES:
+            return mbedtls_aes_setkey_enc(&ctx->ctx.aes, key, key_bitlen);
+#endif
+#if defined(MBEDTLS_ARIA_C)
+        case MBEDTLS_BLOCK_CIPHER_ID_ARIA:
+            return mbedtls_aria_setkey_enc(&ctx->ctx.aria, key, key_bitlen);
+#endif
+#if defined(MBEDTLS_CAMELLIA_C)
+        case MBEDTLS_BLOCK_CIPHER_ID_CAMELLIA:
+            return mbedtls_camellia_setkey_enc(&ctx->ctx.camellia, key, key_bitlen);
+#endif
+        default:
+            return MBEDTLS_ERR_CIPHER_INVALID_CONTEXT;
+    }
+}
 #endif /* MBEDTLS_BLOCK_CIPHER_C */
diff --git a/library/block_cipher_internal.h b/library/block_cipher_internal.h
index 4f25e26..47326fa 100644
--- a/library/block_cipher_internal.h
+++ b/library/block_cipher_internal.h
@@ -50,6 +50,24 @@
                                mbedtls_cipher_id_t cipher_id);
 
 /**
+ * \brief           Set the key into the context.
+ *
+ * \param ctx       The context to configure.
+ * \param key       The buffer holding the key material.
+ * \param key_bitlen    The size of the key in bits.
+ *
+ * \retval          \c 0 on success.
+ * \retval          #MBEDTLS_ERR_CIPHER_INVALID_CONTEXT if the context was not
+ *                  properly set up before calling this function.
+ * \retval          One of #MBEDTLS_ERR_AES_INVALID_KEY_LENGTH,
+ *                  #MBEDTLS_ERR_ARIA_BAD_INPUT_DATA,
+ *                  #MBEDTLS_ERR_CAMELLIA_BAD_INPUT_DATA if \p key_bitlen is
+ *                  invalid.
+ */
+int mbedtls_block_cipher_setkey(mbedtls_block_cipher_context_t *ctx,
+                                const unsigned char *key,
+                                unsigned key_bitlen);
+/**
  * \brief           Clear the context.
  *
  * \param ctx       The context to clear.
diff --git a/tests/suites/test_suite_block_cipher.function b/tests/suites/test_suite_block_cipher.function
index 7133653..dcafdc8 100644
--- a/tests/suites/test_suite_block_cipher.function
+++ b/tests/suites/test_suite_block_cipher.function
@@ -2,6 +2,19 @@
 #include "block_cipher_internal.h"
 
 #define BLOCK_SIZE 16
+
+#if defined(MBEDTLS_AES_C)
+#define VALID_CIPHER_ID MBEDTLS_CIPHER_ID_AES
+#define BADKEY_ERROR MBEDTLS_ERR_AES_INVALID_KEY_LENGTH
+#elif defined(MBEDTLS_ARIA_C)
+#define VALID_CIPHER_ID MBEDTLS_CIPHER_ID_ARIA
+#define BADKEY_ERROR MBEDTLS_ERR_ARIA_BAD_INPUT_DATA
+#elif defined(MBEDTLS_CAMELLIA_C)
+#define VALID_CIPHER_ID MBEDTLS_CIPHER_ID_CAMELLIA
+#define BADKEY_ERROR MBEDTLS_ERR_CAMELLIA_BAD_INPUT_DATA
+#else
+#undef VALID_CIPHER_ID
+#endif
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -9,19 +22,33 @@
  * END_DEPENDENCIES
  */
 
-/* BEGIN_CASE */
+/* BEGIN_CASE depends_on:VALID_CIPHER_ID */
 void invalid()
 {
+    /* That size is valid for a key or an input/output block. */
+    unsigned char buf[16] = { 0 };
+
     mbedtls_block_cipher_context_t ctx;
 
     mbedtls_block_cipher_init(&ctx);
 
+    /* Bad parameters to setup */
     TEST_EQUAL(MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA,
                mbedtls_block_cipher_setup(&ctx, MBEDTLS_CIPHER_ID_NONE));
-
     TEST_EQUAL(MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA,
                mbedtls_block_cipher_setup(&ctx, MBEDTLS_CIPHER_ID_DES));
 
+    /* setkey() before successful setup() */
+    TEST_EQUAL(MBEDTLS_ERR_CIPHER_INVALID_CONTEXT,
+               mbedtls_block_cipher_setkey(&ctx, buf, 128));
+
+    /* Now properly setup the context */
+    TEST_EQUAL(0, mbedtls_block_cipher_setup(&ctx, VALID_CIPHER_ID));
+
+    /* Bad parameters to setkey() */
+    TEST_EQUAL(BADKEY_ERROR,
+               mbedtls_block_cipher_setkey(&ctx, buf, 42));
+
 exit:
     mbedtls_block_cipher_free(&ctx);
 }