extmod/modtls_mbedtls: Support alternate sign callbacks in Python.

This commit enables the implementation of alternative mbedTLS cryptography
functions, such as ECDSA sign and verify, in pure Python.  Alternative
functions are implemented in Python callbacks, that get invoked from
wrapper functions when needed.  The callback can return None to fall back
to the default mbedTLS function.

A common use case for this feature is with secure elements that have
drivers implemented in Python.  Currently, only the ECDSA alternate sign
function wrapper is implemented.

Tested signing with a private EC key stored on an NXP SE05x secure element.

Signed-off-by: iabdalkader <i.abdalkader@gmail.com>

1 files changed, 70 insertions, 0 deletions

diff --git a/extmod/modtls_mbedtls.c b/extmod/modtls_mbedtls.c
index b261e7a70..0a1b8828a 100644
--- a/extmod/modtls_mbedtls.c
+++ b/extmod/modtls_mbedtls.c
@@ -53,6 +53,10 @@
 #else
 #include "mbedtls/version.h"
 #endif
+#if MICROPY_PY_SSL_ECDSA_SIGN_ALT
+#include "mbedtls/ecdsa.h"
+#include "mbedtls/asn1.h"
+#endif
 
 #define MP_STREAM_POLL_RDWR (MP_STREAM_POLL_RD | MP_STREAM_POLL_WR)
 
@@ -68,6 +72,9 @@ typedef struct _mp_obj_ssl_context_t {
     int authmode;
     int *ciphersuites;
     mp_obj_t handler;
+    #if MICROPY_PY_SSL_ECDSA_SIGN_ALT
+    mp_obj_t ecdsa_sign_callback;
+    #endif
 } mp_obj_ssl_context_t;
 
 // This corresponds to an SSLSocket object.
@@ -248,6 +255,9 @@ static mp_obj_t ssl_context_make_new(const mp_obj_type_t *type_in, size_t n_args
     mbedtls_pk_init(&self->pkey);
     self->ciphersuites = NULL;
     self->handler = mp_const_none;
+    #if MICROPY_PY_SSL_ECDSA_SIGN_ALT
+    self->ecdsa_sign_callback = mp_const_none;
+    #endif
 
     #ifdef MBEDTLS_DEBUG_C
     // Debug level (0-4) 1=warning, 2=info, 3=debug, 4=verbose
@@ -295,6 +305,10 @@ static void ssl_context_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
             dest[0] = MP_OBJ_NEW_SMALL_INT(self->authmode);
         } else if (attr == MP_QSTR_verify_callback) {
             dest[0] = self->handler;
+        #if MICROPY_PY_SSL_ECDSA_SIGN_ALT
+        } else if (attr == MP_QSTR_ecdsa_sign_callback) {
+            dest[0] = self->ecdsa_sign_callback;
+        #endif
         } else {
             // Continue lookup in locals_dict.
             dest[1] = MP_OBJ_SENTINEL;
@@ -305,6 +319,11 @@ static void ssl_context_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
             self->authmode = mp_obj_get_int(dest[1]);
             dest[0] = MP_OBJ_NULL;
             mbedtls_ssl_conf_authmode(&self->conf, self->authmode);
+        #if MICROPY_PY_SSL_ECDSA_SIGN_ALT
+        } else if (attr == MP_QSTR_ecdsa_sign_callback) {
+            dest[0] = MP_OBJ_NULL;
+            self->ecdsa_sign_callback = dest[1];
+        #endif
         } else if (attr == MP_QSTR_verify_callback) {
             dest[0] = MP_OBJ_NULL;
             self->handler = dest[1];
@@ -786,6 +805,57 @@ static MP_DEFINE_CONST_OBJ_TYPE(
 /******************************************************************************/
 // ssl module.
 
+#if MICROPY_PY_SSL_ECDSA_SIGN_ALT
+int micropy_mbedtls_ecdsa_sign_alt(const mbedtls_mpi *d, const unsigned char *hash, size_t hlen, unsigned char *sig, size_t sig_size, size_t *slen) {
+    uint8_t key[256];
+
+    // Check if the current context has an alternative sign function.
+    mp_obj_ssl_context_t *ssl_ctx = MP_STATE_THREAD(tls_ssl_context);
+    if (ssl_ctx == NULL || ssl_ctx->ecdsa_sign_callback == mp_const_none) {
+        return MBEDTLS_ERR_PLATFORM_FEATURE_UNSUPPORTED;
+    }
+
+    size_t klen = mbedtls_mpi_size(d);
+    if (klen > sizeof(key)) {
+        return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
+    }
+
+    // Convert the MPI private key (d) to a binary array
+    if (mbedtls_mpi_write_binary(d, key, klen) != 0) {
+        return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
+    }
+
+    nlr_buf_t nlr;
+    mp_buffer_info_t sig_buf;
+    if (nlr_push(&nlr) == 0) {
+        mp_obj_t ret = mp_call_function_2(ssl_ctx->ecdsa_sign_callback,
+            mp_obj_new_bytearray_by_ref(klen, (void *)key),
+            mp_obj_new_bytearray_by_ref(hlen, (void *)hash));
+        if (ret == mp_const_none) {
+            // key couldn't be used by the alternative implementation.
+            nlr_pop();
+            return MBEDTLS_ERR_PLATFORM_FEATURE_UNSUPPORTED;
+        }
+        mp_get_buffer_raise(ret, &sig_buf, MP_BUFFER_READ);
+        nlr_pop();
+    } else {
+        // The alternative implementation failed to sign.
+        mp_obj_print_exception(&mp_plat_print, MP_OBJ_FROM_PTR(nlr.ret_val));
+        return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
+    }
+
+    // Check if the buffer fits.
+    if (sig_buf.len > sig_size) {
+        return MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL;
+    }
+
+    // Copy ASN.1 signature to buffer.
+    *slen = sig_buf.len;
+    memcpy(sig, sig_buf.buf, sig_buf.len);
+    return 0;
+}
+#endif
+
 static const mp_rom_map_elem_t mp_module_tls_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_tls) },