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>

3 files changed, 163 insertions, 0 deletions

diff --git a/extmod/extmod.mk b/extmod/extmod.mk
index c132fd89c..514197d6b 100644
--- a/extmod/extmod.mk
+++ b/extmod/extmod.mk
@@ -42,6 +42,7 @@ SRC_EXTMOD_C += \
 	extmod/modsocket.c \
 	extmod/modtls_axtls.c \
 	extmod/modtls_mbedtls.c \
+	extmod/mbedtls/mbedtls_alt.c \
 	extmod/modtime.c \
 	extmod/moductypes.c \
 	extmod/modvfs.c \
@@ -242,6 +243,10 @@ MBEDTLS_CONFIG_FILE ?= \"mbedtls/mbedtls_config_port.h\"
 GIT_SUBMODULES += $(MBEDTLS_DIR)
 CFLAGS_EXTMOD += -DMBEDTLS_CONFIG_FILE=$(MBEDTLS_CONFIG_FILE)
 CFLAGS_EXTMOD += -DMICROPY_SSL_MBEDTLS=1 -I$(TOP)/$(MBEDTLS_DIR)/include
+ifeq ($(MICROPY_PY_SSL_ECDSA_SIGN_ALT),1)
+CFLAGS_EXTMOD += -DMICROPY_PY_SSL_ECDSA_SIGN_ALT=1
+LDFLAGS_EXTMOD += -Wl,--wrap=mbedtls_ecdsa_write_signature
+endif
 SRC_THIRDPARTY_C += lib/mbedtls_errors/mp_mbedtls_errors.c
 SRC_THIRDPARTY_C += $(addprefix $(MBEDTLS_DIR)/library/,\
 	aes.c \
diff --git a/extmod/mbedtls/mbedtls_alt.c b/extmod/mbedtls/mbedtls_alt.c
new file mode 100644
index 000000000..ccfb37348
--- /dev/null
+++ b/extmod/mbedtls/mbedtls_alt.c
@@ -0,0 +1,88 @@
+/*
+ * This file is part of the MicroPython project, http://micropython.org/
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Copyright The Mbed TLS Contributors
+ * Copyright (c) 2024 Damien P. George
+ *
+ * This file provides default fallback functions for use with alternate
+ * cryptography functions implemented in Python.
+ */
+#if MICROPY_PY_SSL_ECDSA_SIGN_ALT
+#if defined(MBEDTLS_ECP_RESTARTABLE) || defined(MBEDTLS_ECDSA_DETERMINISTIC)
+#error "MICROPY_PY_SSL_ECDSA_SIGN_ALT cannot be used with MBEDTLS_ECP_RESTARTABLE or MBEDTLS_ECDSA_DETERMINISTIC"
+#endif
+
+#include <string.h>
+#define MBEDTLS_ALLOW_PRIVATE_ACCESS
+#include "mbedtls/platform.h"
+#include "mbedtls/ssl.h"
+#include "mbedtls/error.h"
+#include "mbedtls/ecdsa.h"
+#include "mbedtls/asn1write.h"
+
+extern 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);
+
+
+// Compute and write signature
+// See lib/mbedtls/library/ecdsa.c:688
+//
+// Note: To avoid duplicating a lot of code, MBEDTLS_ECDSA_SIGN_ALT is not defined,
+// which allows the default mbedtls_ecdsa_sign to be used as a fallback function.
+// However, mbedtls_ecdsa_sign cannot be wrapped because it is called internally
+// within its object file, so we wrap mbedtls_ecdsa_read/write_signature instead.
+int __wrap_mbedtls_ecdsa_write_signature(mbedtls_ecdsa_context *ctx,
+    mbedtls_md_type_t md_alg,
+    const unsigned char *hash, size_t hlen,
+    unsigned char *sig, size_t sig_size, size_t *slen,
+    int (*f_rng)(void *, unsigned char *, size_t),
+    void *p_rng) {
+
+    (void)md_alg;
+
+    if (f_rng == NULL) {
+        return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
+    }
+
+    // Check if curve is supported for ECDSA.
+    if (!mbedtls_ecdsa_can_do(ctx->grp.id) || ctx->grp.N.p == NULL) {
+        return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
+    }
+
+    // Try signing with the alternative function first.
+    int ret = micropy_mbedtls_ecdsa_sign_alt(&ctx->d, hash, hlen, sig, sig_size, slen);
+
+    // Fallback to the default mbedtls implementation if needed.
+    if (ret == MBEDTLS_ERR_PLATFORM_FEATURE_UNSUPPORTED) {
+        mbedtls_mpi r, s;
+        mbedtls_mpi_init(&r);
+        mbedtls_mpi_init(&s);
+
+        size_t len = 0;
+        unsigned char buf[MBEDTLS_ECDSA_MAX_LEN] = { 0 };
+        unsigned char *p = buf + sizeof(buf);
+
+        MBEDTLS_MPI_CHK(mbedtls_ecdsa_sign(&ctx->grp, &r, &s, &ctx->d, hash, hlen, f_rng, p_rng));
+        MBEDTLS_ASN1_CHK_CLEANUP_ADD(len, mbedtls_asn1_write_mpi(&p, buf, &s));
+        MBEDTLS_ASN1_CHK_CLEANUP_ADD(len, mbedtls_asn1_write_mpi(&p, buf, &r));
+        MBEDTLS_ASN1_CHK_CLEANUP_ADD(len, mbedtls_asn1_write_len(&p, buf, len));
+        MBEDTLS_ASN1_CHK_CLEANUP_ADD(len, mbedtls_asn1_write_tag(&p, buf, MBEDTLS_ASN1_CONSTRUCTED | MBEDTLS_ASN1_SEQUENCE));
+
+        if (len > sig_size) {
+            ret = MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL;
+        } else {
+            ret = 0;
+            *slen = len;
+            memcpy(sig, p, len);
+        }
+
+    cleanup:
+        mbedtls_mpi_free(&r);
+        mbedtls_mpi_free(&s);
+    }
+
+    return ret;
+}
+#endif
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) },