1 files changed, 79 insertions, 20 deletions

diff --git a/extmod/moducryptolib.c b/extmod/moducryptolib.c
index af9eec624..6c45c2fde 100644
--- a/extmod/moducryptolib.c
+++ b/extmod/moducryptolib.c
@@ -41,10 +41,17 @@
 
 // values follow PEP 272
 enum {
-    UCRYPTOLIB_MODE_MIN = 0,
-    UCRYPTOLIB_MODE_ECB,
-    UCRYPTOLIB_MODE_CBC,
-    UCRYPTOLIB_MODE_MAX,
+    UCRYPTOLIB_MODE_ECB = 1,
+    UCRYPTOLIB_MODE_CBC = 2,
+    UCRYPTOLIB_MODE_CTR = 6,
+};
+
+struct ctr_params {
+    // counter is the IV of the AES context.
+
+    size_t offset; // in encrypted_counter
+    // encrypted counter
+    uint8_t encrypted_counter[16];
 };
 
 #if MICROPY_SSL_AXTLS
@@ -82,6 +89,19 @@ typedef struct _mp_obj_aes_t {
     uint8_t key_type: 2;
 } mp_obj_aes_t;
 
+STATIC inline bool is_ctr_mode(int block_mode) {
+    #if MICROPY_PY_UCRYPTOLIB_CTR
+    return block_mode == UCRYPTOLIB_MODE_CTR;
+    #else
+    return false;
+    #endif
+}
+
+STATIC inline struct ctr_params *ctr_params_from_aes(mp_obj_aes_t *o) {
+    // ctr_params follows aes object struct
+    return (struct ctr_params*)&o[1];
+}
+
 #if MICROPY_SSL_AXTLS
 STATIC void aes_initial_set_key_impl(AES_CTX_IMPL *ctx, const uint8_t *key, size_t keysize, const uint8_t iv[16]) {
     assert(16 == keysize || 32 == keysize);
@@ -155,20 +175,38 @@ STATIC void aes_process_ecb_impl(AES_CTX_IMPL *ctx, const uint8_t in[16], uint8_
 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);
 }
+
+#if MICROPY_PY_UCRYPTOLIB_CTR
+STATIC void aes_process_ctr_impl(AES_CTX_IMPL *ctx, const uint8_t *in, uint8_t *out, size_t in_len, struct ctr_params *ctr_params) {
+    mbedtls_aes_crypt_ctr(&ctx->u.mbedtls_ctx, in_len, &ctr_params->offset, ctx->iv, ctr_params->encrypted_counter, in, out);
+}
+#endif
+
 #endif
 
 STATIC mp_obj_t ucryptolib_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);
-    o->base.type = type;
 
-    o->block_mode = mp_obj_get_int(args[1]);
-    o->key_type = AES_KEYTYPE_NONE;
+    const mp_int_t block_mode = mp_obj_get_int(args[1]);
+
+    switch (block_mode) {
+        case UCRYPTOLIB_MODE_ECB:
+        case UCRYPTOLIB_MODE_CBC:
+        #if MICROPY_PY_UCRYPTOLIB_CTR
+        case UCRYPTOLIB_MODE_CTR:
+        #endif
+            break;
 
-    if (o->block_mode <= UCRYPTOLIB_MODE_MIN || o->block_mode >= UCRYPTOLIB_MODE_MAX) {
-        mp_raise_ValueError("mode");
+        default:
+            mp_raise_ValueError("mode");
     }
 
+    mp_obj_aes_t *o = m_new_obj_var(mp_obj_aes_t, struct ctr_params, !!is_ctr_mode(block_mode));
+    o->base.type = type;
+
+    o->block_mode = block_mode;
+    o->key_type = AES_KEYTYPE_NONE;
+
     mp_buffer_info_t keyinfo;
     mp_get_buffer_raise(args[0], &keyinfo, MP_BUFFER_READ);
     if (32 != keyinfo.len && 16 != keyinfo.len) {
@@ -183,10 +221,14 @@ STATIC mp_obj_t ucryptolib_aes_make_new(const mp_obj_type_t *type, size_t n_args
         if (16 != ivinfo.len) {
             mp_raise_ValueError("IV");
         }
-    } else if (o->block_mode == UCRYPTOLIB_MODE_CBC) {
+    } else if (o->block_mode == UCRYPTOLIB_MODE_CBC || is_ctr_mode(o->block_mode)) {
         mp_raise_ValueError("IV");
     }
 
+    if (is_ctr_mode(block_mode)) {
+        ctr_params_from_aes(o)->offset = 0;
+    }
+
     aes_initial_set_key_impl(&o->ctx, keyinfo.buf, keyinfo.len, ivinfo.buf);
 
     return MP_OBJ_FROM_PTR(o);
@@ -204,7 +246,7 @@ STATIC mp_obj_t aes_process(size_t n_args, const mp_obj_t *args, bool encrypt) {
     mp_buffer_info_t in_bufinfo;
     mp_get_buffer_raise(in_buf, &in_bufinfo, MP_BUFFER_READ);
 
-    if (in_bufinfo.len % 16 != 0) {
+    if (!is_ctr_mode(self->block_mode) && in_bufinfo.len % 16 != 0) {
         mp_raise_ValueError("blksize % 16");
     }
 
@@ -224,7 +266,9 @@ STATIC mp_obj_t aes_process(size_t n_args, const mp_obj_t *args, bool encrypt) {
     }
 
     if (AES_KEYTYPE_NONE == self->key_type) {
-        aes_final_set_key_impl(&self->ctx, encrypt);
+        // always set key for encryption if CTR mode.
+        const bool encrypt_mode = encrypt || is_ctr_mode(self->block_mode);
+        aes_final_set_key_impl(&self->ctx, encrypt_mode);
         self->key_type = encrypt ? AES_KEYTYPE_ENC : AES_KEYTYPE_DEC;
     } else {
         if ((encrypt && self->key_type == AES_KEYTYPE_DEC) ||
@@ -234,14 +278,26 @@ STATIC mp_obj_t aes_process(size_t n_args, const mp_obj_t *args, bool encrypt) {
         }
     }
 
-    if (self->block_mode == UCRYPTOLIB_MODE_ECB) {
-        uint8_t *in = in_bufinfo.buf, *out = out_buf_ptr;
-        uint8_t *top = in + in_bufinfo.len;
-        for (; in < top; in += 16, out += 16) {
-            aes_process_ecb_impl(&self->ctx, in, out, encrypt);
+    switch (self->block_mode) {
+        case UCRYPTOLIB_MODE_ECB: {
+            uint8_t *in = in_bufinfo.buf, *out = out_buf_ptr;
+            uint8_t *top = in + in_bufinfo.len;
+            for (; in < top; in += 16, out += 16) {
+                aes_process_ecb_impl(&self->ctx, in, out, encrypt);
+            }
+            break;
         }
-    } else {
-        aes_process_cbc_impl(&self->ctx, in_bufinfo.buf, out_buf_ptr, in_bufinfo.len, encrypt);
+
+        case UCRYPTOLIB_MODE_CBC:
+            aes_process_cbc_impl(&self->ctx, in_bufinfo.buf, out_buf_ptr, in_bufinfo.len, encrypt);
+            break;
+
+        #if MICROPY_PY_UCRYPTOLIB_CTR
+        case UCRYPTOLIB_MODE_CTR:
+            aes_process_ctr_impl(&self->ctx, in_bufinfo.buf, out_buf_ptr, in_bufinfo.len,
+                ctr_params_from_aes(self));
+            break;
+        #endif
     }
 
     if (out_buf != MP_OBJ_NULL) {
@@ -279,6 +335,9 @@ STATIC const mp_rom_map_elem_t mp_module_ucryptolib_globals_table[] = {
 #if MICROPY_PY_UCRYPTOLIB_CONSTS
     { MP_ROM_QSTR(MP_QSTR_MODE_ECB), MP_ROM_INT(UCRYPTOLIB_MODE_ECB) },
     { MP_ROM_QSTR(MP_QSTR_MODE_CBC), MP_ROM_INT(UCRYPTOLIB_MODE_CBC) },
+    #if MICROPY_PY_UCRYPTOLIB_CTR
+    { MP_ROM_QSTR(MP_QSTR_MODE_CTR), MP_ROM_INT(UCRYPTOLIB_MODE_CTR) },
+    #endif
 #endif
 };