extmod/moducryptolib: Add an mbedTLS implementation for this module.

Jongy · dpgeorge · commit eacb233b8f27 · 2018-06-27T16:44:59.000+10:00
diff --git a/extmod/moducryptolib.c b/extmod/moducryptolib.c
@@ -54,6 +54,24 @@ enum {
 #define AES_CTX_IMPL AES_CTX
 #endif
 
+#if MICROPY_SSL_MBEDTLS
+#include <mbedtls/aes.h>
+
+// we can't run mbedtls AES key schedule until we know whether we're used for encrypt or decrypt.
+// therefore, we store the key & keysize and on the first call to encrypt/decrypt we override them
+// with the mbedtls_aes_context, as they are not longer required. (this is done to save space)
+struct mbedtls_aes_ctx_with_key {
+    union {
+        mbedtls_aes_context mbedtls_ctx;
+        struct {
+            uint8_t key[32];
+            uint8_t keysize;
+        } init_data;
+    } u;
+    unsigned char iv[16];
+};
+#define AES_CTX_IMPL struct mbedtls_aes_ctx_with_key
+#endif
 
 typedef struct _mp_obj_aes_t {
     mp_obj_base_t base;
@@ -104,6 +122,42 @@ STATIC void aes_process_cbc_impl(AES_CTX_IMPL *ctx, const uint8_t *in, uint8_t *
 }
 #endif
 
+#if MICROPY_SSL_MBEDTLS
+STATIC void aes_initial_set_key_impl(AES_CTX_IMPL *ctx, const uint8_t *key, size_t keysize, const uint8_t iv[16]) {
+    ctx->u.init_data.keysize = keysize;
+    memcpy(ctx->u.init_data.key, key, keysize);
+
+    if (NULL != iv) {
+        memcpy(ctx->iv, iv, sizeof(ctx->iv));
+    }
+}
+
+STATIC void aes_final_set_key_impl(AES_CTX_IMPL *ctx, bool encrypt) {
+    // first, copy key aside
+    uint8_t key[32];
+    uint8_t keysize = ctx->u.init_data.keysize;
+    memcpy(key, ctx->u.init_data.key, keysize);
+    // now, override key with the mbedtls context object
+    mbedtls_aes_init(&ctx->u.mbedtls_ctx);
+
+    // setkey call will succeed, we've already checked the keysize earlier.
+    assert(16 == keysize || 32 == keysize);
+    if (encrypt) {
+        mbedtls_aes_setkey_enc(&ctx->u.mbedtls_ctx, key, keysize * 8);
+    } else {
+        mbedtls_aes_setkey_dec(&ctx->u.mbedtls_ctx, key, keysize * 8);
+    }
+}
+
+STATIC void aes_process_ecb_impl(AES_CTX_IMPL *ctx, const uint8_t in[16], uint8_t out[16], bool encrypt) {
+    mbedtls_aes_crypt_ecb(&ctx->u.mbedtls_ctx, encrypt ? MBEDTLS_AES_ENCRYPT : MBEDTLS_AES_DECRYPT, in, out);
+}
+
+STATIC void aes_process_cbc_impl(AES_CTX_IMPL *ctx, const uint8_t *in, uint8_t *out, size_t in_len, bool encrypt) {
+    mbedtls_aes_crypt_cbc(&ctx->u.mbedtls_ctx, encrypt ? MBEDTLS_AES_ENCRYPT : MBEDTLS_AES_DECRYPT, in_len, ctx->iv, in, out);
+}
+#endif
+
 STATIC mp_obj_t aes_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) {
     mp_arg_check_num(n_args, n_kw, 2, 3, false);
     mp_obj_aes_t *o = m_new_obj(mp_obj_aes_t);
