extmod/modussl_mbedtls: Wire in support for DTLS.

damz · damz · commit 003377c6c514 · 2023-01-18T18:25:18.000-08:00
A new boolean argument `dtls` is added to `ussl.wrap_socket()`, if true the library assumes that the underlying socket is a datagram socket (i.e. UDP or similar). Implement our own timer callbacks as the out of the box implementation relies on `gettimeofday()`. To fully support asyncio for DTLS socket, we will need to return a readable or writable event in `poll(MP_STREAM_POLL, ...)` if `_mbedtls_timing_get_delay(self) >= 1`. This is left for future work so as not to interfere with micropython#9871.
diff --git a/extmod/modussl_mbedtls.c b/extmod/modussl_mbedtls.c
@@ -35,6 +35,8 @@
 #include "py/runtime.h"
 #include "py/stream.h"
 #include "py/objstr.h"
+#include "py/smallint.h"
+#include "py/mphal.h"
 
 // mbedtls_time_t
 #include "mbedtls/platform.h"
@@ -43,6 +45,7 @@
 #include "mbedtls/pk.h"
 #include "mbedtls/entropy.h"
 #include "mbedtls/ctr_drbg.h"
+#include "mbedtls/timing.h"
 #include "mbedtls/debug.h"
 #include "mbedtls/error.h"
 
@@ -61,6 +64,12 @@ typedef struct _mp_obj_ssl_socket_t {
 
     uintptr_t poll_mask; // Indicates which read or write operations the protocol needs next
     int last_error; // The last error code, if any
+
+    #ifdef MBEDTLS_SSL_PROTO_DTLS
+    mp_uint_t timer_start_ms;
+    mp_int_t timer_fin_ms;
+    mp_int_t timer_int_ms;
+    #endif
 } mp_obj_ssl_socket_t;
 
 struct ssl_args {
@@ -71,6 +80,7 @@ struct ssl_args {
     mp_arg_val_t cert_reqs;
     mp_arg_val_t cadata;
     mp_arg_val_t do_handshake;
+    mp_arg_val_t dtls;
 };
 
 STATIC const mp_obj_type_t ussl_socket_type;
@@ -158,6 +168,40 @@ STATIC int _mbedtls_ssl_recv(void *ctx, byte *buf, size_t len) {
     }
 }
 
+#ifdef MBEDTLS_SSL_PROTO_DTLS
+STATIC void _mbedtls_timing_set_delay(void *ctx, uint32_t int_ms, uint32_t fin_ms) {
+    mp_obj_ssl_socket_t *o = (mp_obj_ssl_socket_t *)ctx;
+
+    o->timer_int_ms = int_ms;
+    o->timer_fin_ms = fin_ms;
+
+    if (fin_ms != 0) {
+        o->timer_start_ms = mp_hal_ticks_ms() & (MICROPY_PY_UTIME_TICKS_PERIOD - 1);
+    }
+}
+
+STATIC int _mbedtls_timing_get_delay(void *ctx) {
+    mp_obj_ssl_socket_t *o = (mp_obj_ssl_socket_t *)ctx;
+
+    if (o->timer_fin_ms == 0) {
+        return -1;
+    }
+
+    mp_uint_t now = mp_hal_ticks_ms() & (MICROPY_PY_UTIME_TICKS_PERIOD - 1);
+    mp_int_t elapsed_ms = ((now - o->timer_start_ms + MICROPY_PY_UTIME_TICKS_PERIOD / 2) & (MICROPY_PY_UTIME_TICKS_PERIOD - 1))
+        - MICROPY_PY_UTIME_TICKS_PERIOD / 2;
+
+    if (elapsed_ms >= o->timer_fin_ms) {
+        return 2;
+    }
+
+    if (elapsed_ms >= o->timer_int_ms) {
+        return 1;
+    }
+
+    return 0;
+}
+#endif
 
 STATIC mp_obj_ssl_socket_t *socket_new(mp_obj_t sock, struct ssl_args *args) {
     // Verify the socket object has the full stream protocol
@@ -194,7 +238,7 @@ STATIC mp_obj_ssl_socket_t *socket_new(mp_obj_t sock, struct ssl_args *args) {
 
     ret = mbedtls_ssl_config_defaults(&o->conf,
         args->server_side.u_bool ? MBEDTLS_SSL_IS_SERVER : MBEDTLS_SSL_IS_CLIENT,
-        MBEDTLS_SSL_TRANSPORT_STREAM,
+        args->dtls.u_bool ? MBEDTLS_SSL_TRANSPORT_DATAGRAM : MBEDTLS_SSL_TRANSPORT_STREAM,
         MBEDTLS_SSL_PRESET_DEFAULT);
     if (ret != 0) {
         goto cleanup;
@@ -219,6 +263,12 @@ STATIC mp_obj_ssl_socket_t *socket_new(mp_obj_t sock, struct ssl_args *args) {
         }
     }
 
+    #ifdef MBEDTLS_SSL_PROTO_DTLS
+    mbedtls_ssl_set_timer_cb(&o->ssl, o,
+        _mbedtls_timing_set_delay,
+        _mbedtls_timing_get_delay);
+    #endif
+
     mbedtls_ssl_set_bio(&o->ssl, &o->sock, _mbedtls_ssl_send, _mbedtls_ssl_recv, NULL);
 
     if (args->key.u_obj != mp_const_none) {
@@ -483,6 +533,7 @@ STATIC mp_obj_t mod_ssl_wrap_socket(size_t n_args, const mp_obj_t *pos_args, mp_
         { MP_QSTR_cert_reqs, MP_ARG_KW_ONLY | MP_ARG_INT, {.u_int = MBEDTLS_SSL_VERIFY_NONE}},
         { MP_QSTR_cadata, MP_ARG_KW_ONLY | MP_ARG_OBJ, {.u_rom_obj = MP_ROM_NONE} },
         { MP_QSTR_do_handshake, MP_ARG_KW_ONLY | MP_ARG_BOOL, {.u_bool = true} },
+        { MP_QSTR_dtls, MP_ARG_KW_ONLY | MP_ARG_BOOL, {.u_bool = false} },
     };
 
     // TODO: Check that sock implements stream protocol
diff --git a/ports/esp32/boards/sdkconfig.base b/ports/esp32/boards/sdkconfig.base
@@ -46,6 +46,7 @@ CONFIG_LWIP_PPP_CHAP_SUPPORT=y
 # SSL
 # Use 4kiB output buffer instead of default 16kiB (because IDF heap is fragmented in 4.0)
 CONFIG_MBEDTLS_ASYMMETRIC_CONTENT_LEN=y
+CONFIG_MBEDTLS_SSL_PROTO_DTLS=y
 
 # ULP coprocessor support
 CONFIG_ESP32_ULP_COPROC_ENABLED=y
