extmod/mbedtls: Implement recommended DTLS features, make optional.

- DTLS spec recommends HelloVerify and Anti Replay protection be enabled,
  and these are enabled in the default mbedTLS config. Implement them here.

- To help compensate for the possible increase in code size, add a
  MICROPY_PY_SSL_DTLS build config macro that's enabled for EXTRA and
  above by default.

This allows bare metal mbedTLS ports to use DTLS with HelloVerify support.

This work was funded through GitHub Sponsors.

Signed-off-by: Angus Gratton <angus@redyak.com.au>

6 files changed, 84 insertions, 21 deletions

diff --git a/docs/library/ssl.rst b/docs/library/ssl.rst
index 4327c74ba..c86101872 100644
--- a/docs/library/ssl.rst
+++ b/docs/library/ssl.rst
@@ -66,7 +66,7 @@ class SSLContext
    Set the available ciphers for sockets created with this context.  *ciphers* should be
    a list of strings in the `IANA cipher suite format <https://wiki.mozilla.org/Security/Cipher_Suites>`_ .
 
-.. method:: SSLContext.wrap_socket(sock, *, server_side=False, do_handshake_on_connect=True, server_hostname=None)
+.. method:: SSLContext.wrap_socket(sock, *, server_side=False, do_handshake_on_connect=True, server_hostname=None, client_id=None)
 
    Takes a `stream` *sock* (usually socket.socket instance of ``SOCK_STREAM`` type),
    and returns an instance of ssl.SSLSocket, wrapping the underlying stream.
@@ -89,6 +89,9 @@ class SSLContext
      server certificate.  It also sets the name for Server Name Indication (SNI), allowing the server
      to present the proper certificate.
 
+   - *client_id* is a MicroPython-specific extension argument used only when implementing a DTLS
+     Server. See :ref:`dtls` for details.
+
 .. warning::
 
    Some implementations of ``ssl`` module do NOT validate server certificates,
@@ -117,6 +120,8 @@ Exceptions
 
    This exception does NOT exist. Instead its base class, OSError, is used.
 
+.. _dtls:
+
 DTLS support
 ------------
 
@@ -125,16 +130,47 @@ DTLS support
 
    This is a MicroPython extension.
 
-This module supports DTLS in client and server mode via the `PROTOCOL_DTLS_CLIENT`
-and `PROTOCOL_DTLS_SERVER` constants that can be used as the ``protocol`` argument
-of `SSLContext`.
+On most ports, this module supports DTLS in client and server mode via the
+`PROTOCOL_DTLS_CLIENT` and `PROTOCOL_DTLS_SERVER` constants that can be used as
+the ``protocol`` argument of `SSLContext`.
 
 In this case the underlying socket is expected to behave as a datagram socket (i.e.
 like the socket opened with ``socket.socket`` with ``socket.AF_INET`` as ``af`` and
 ``socket.SOCK_DGRAM`` as ``type``).
 
-DTLS is only supported on ports that use mbed TLS, and it is not enabled by default:
-it requires enabling ``MBEDTLS_SSL_PROTO_DTLS`` in the specific port configuration.
+DTLS is only supported on ports that use mbedTLS, and it is enabled by default
+in most configurations but can be manually disabled by defining
+``MICROPY_PY_SSL_DTLS`` to 0.
+
+DTLS server support
+^^^^^^^^^^^^^^^^^^^
+
+MicroPython's DTLS server support is configured with "Hello Verify" as required
+for DTLS 1.2. This is transparent for DTLS clients, but there are relevant
+considerations when implementing a DTLS server in MicroPython:
+
+- The server should pass an additional argument *client_id* when calling
+  `SSLContext.wrap_socket()`. This ID must be a `bytes` object (or similar) with
+  a transport-specific identifier representing the client.
+
+  The simplest approach is to convert the tuple of ``(client_ip, client_port)``
+  returned from ``socket.recv_from()`` into a byte string, i.e.::
+
+        _, client_addr = sock.recvfrom(1, socket.MSG_PEEK)
+        sock.connect(client_addr)  # Connect back to the client
+        sock = ssl_ctx.wrap_socket(sock, server_side=True,
+                                   client_id=repr(client_addr).encode())
+
+- The first time a client connects, the server call to ``wrap_socket`` will fail
+  with a `OSError` error "Hello Verify Required". This is because the DTLS
+  "Hello Verify" cookie is not yet known by the client. If the same client
+  connects a second time then ``wrap_socket`` will succeed.
+
+- DTLS cookies for "Hello Verify" are associated with the `SSLContext` object,
+  so the same `SSLContext` object should be used to wrap a subsequent connection
+  from the same client. The cookie implementation includes a timeout and has
+  constant memory use regardless of how many clients connect, so it's OK to
+  reuse the same `SSLContext` object for the lifetime of the server.
 
 Constants
 ---------
diff --git a/extmod/mbedtls/mbedtls_config_common.h b/extmod/mbedtls/mbedtls_config_common.h
index 6cd14befc..1f7ac8818 100644
--- a/extmod/mbedtls/mbedtls_config_common.h
+++ b/extmod/mbedtls/mbedtls_config_common.h
@@ -26,6 +26,8 @@
 #ifndef MICROPY_INCLUDED_MBEDTLS_CONFIG_COMMON_H
 #define MICROPY_INCLUDED_MBEDTLS_CONFIG_COMMON_H
 
+#include "py/mpconfig.h"
+
 // If you want to debug MBEDTLS uncomment the following and
 // pass "3" to mbedtls_debug_set_threshold in socket_new.
 // #define MBEDTLS_DEBUG_C
@@ -89,12 +91,18 @@
 #define MBEDTLS_SHA384_C
 #define MBEDTLS_SHA512_C
 #define MBEDTLS_SSL_CLI_C
-#define MBEDTLS_SSL_PROTO_DTLS
 #define MBEDTLS_SSL_SRV_C
 #define MBEDTLS_SSL_TLS_C
 #define MBEDTLS_X509_CRT_PARSE_C
 #define MBEDTLS_X509_USE_C
 
+#if MICROPY_PY_SSL_DTLS
+#define MBEDTLS_SSL_PROTO_DTLS
+#define MBEDTLS_SSL_DTLS_ANTI_REPLAY
+#define MBEDTLS_SSL_DTLS_HELLO_VERIFY
+#define MBEDTLS_SSL_COOKIE_C
+#endif
+
 // A port may enable this option to select additional bare-metal configuration.
 #if MICROPY_MBEDTLS_CONFIG_BARE_METAL
 
diff --git a/extmod/modtls_mbedtls.c b/extmod/modtls_mbedtls.c
index 4bd0aea9a..418275440 100644
--- a/extmod/modtls_mbedtls.c
+++ b/extmod/modtls_mbedtls.c
@@ -78,7 +78,7 @@
 #define MP_PROTOCOL_TLS_CLIENT  0
 #define MP_PROTOCOL_TLS_SERVER  MP_ENDPOINT_IS_SERVER
 #define MP_PROTOCOL_DTLS_CLIENT MP_TRANSPORT_IS_DTLS
-#define MP_PROTOCOL_DTLS_SERVER MP_ENDPOINT_IS_SERVER | MP_TRANSPORT_IS_DTLS
+#define MP_PROTOCOL_DTLS_SERVER (MP_ENDPOINT_IS_SERVER | MP_TRANSPORT_IS_DTLS)
 
 // This corresponds to an SSLContext object.
 typedef struct _mp_obj_ssl_context_t {
@@ -96,6 +96,7 @@ typedef struct _mp_obj_ssl_context_t {
     mp_obj_t ecdsa_sign_callback;
     #endif
     #ifdef MBEDTLS_SSL_DTLS_HELLO_VERIFY
+    bool is_dtls_server;
     mbedtls_ssl_cookie_ctx cookie_ctx;
     #endif
 } mp_obj_ssl_context_t;
@@ -328,14 +329,17 @@ static mp_obj_t ssl_context_make_new(const mp_obj_type_t *type_in, size_t n_args
     #endif
 
     #ifdef MBEDTLS_SSL_DTLS_HELLO_VERIFY
-    mbedtls_ssl_cookie_init(&self->cookie_ctx);
-    ret = mbedtls_ssl_cookie_setup(&self->cookie_ctx, mbedtls_ctr_drbg_random, &self->ctr_drbg);
-    if (ret != 0) {
-        mbedtls_raise_error(ret);
+    self->is_dtls_server = (protocol == MP_PROTOCOL_DTLS_SERVER);
+    if (self->is_dtls_server) {
+        mbedtls_ssl_cookie_init(&self->cookie_ctx);
+        ret = mbedtls_ssl_cookie_setup(&self->cookie_ctx, mbedtls_ctr_drbg_random, &self->ctr_drbg);
+        if (ret != 0) {
+            mbedtls_raise_error(ret);
+        }
+        mbedtls_ssl_conf_dtls_cookies(&self->conf, mbedtls_ssl_cookie_write, mbedtls_ssl_cookie_check,
+            &self->cookie_ctx);
     }
-    mbedtls_ssl_conf_dtls_cookies(&self->conf, mbedtls_ssl_cookie_write, mbedtls_ssl_cookie_check,
-        &self->cookie_ctx);
-    #endif
+    #endif // MBEDTLS_SSL_DTLS_HELLO_VERIFY
 
     return MP_OBJ_FROM_PTR(self);
 }
@@ -664,19 +668,18 @@ static mp_obj_t ssl_socket_make_new(mp_obj_ssl_context_t *ssl_context, mp_obj_t
     #ifdef MBEDTLS_SSL_PROTO_DTLS
     mbedtls_ssl_set_timer_cb(&o->ssl, o, _mbedtls_timing_set_delay, _mbedtls_timing_get_delay);
     #endif
+
     #ifdef MBEDTLS_SSL_DTLS_HELLO_VERIFY
-    if (client_id != mp_const_none) {
+    if (ssl_context->is_dtls_server) {
+        // require the client_id parameter for DTLS (as per mbedTLS requirement)
+        ret = MBEDTLS_ERR_SSL_BAD_INPUT_DATA;
         mp_buffer_info_t buf;
         if (mp_get_buffer(client_id, &buf, MP_BUFFER_READ)) {
             ret = mbedtls_ssl_set_client_transport_id(&o->ssl, buf.buf, buf.len);
-        } else {
-            ret = MBEDTLS_ERR_SSL_BAD_INPUT_DATA;
         }
         if (ret != 0) {
             goto cleanup;
         }
-    } else {
-        // TODO: should it be an error not to provide this argument for DTLS server?
     }
     #endif
 
diff --git a/py/mpconfig.h b/py/mpconfig.h
index 4c1276275..a1025fe5e 100644
--- a/py/mpconfig.h
+++ b/py/mpconfig.h
@@ -1941,6 +1941,11 @@ typedef time_t mp_timestamp_t;
 #define MICROPY_PY_SSL_MBEDTLS_NEED_ACTIVE_CONTEXT (MICROPY_PY_SSL_ECDSA_SIGN_ALT)
 #endif
 
+// Whether to support DTLS protocol (non-CPython feature)
+#ifndef MICROPY_PY_SSL_DTLS
+#define MICROPY_PY_SSL_DTLS (MICROPY_SSL_MBEDTLS && MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EXTRA_FEATURES)
+#endif
+
 // Whether to provide the "vfs" module
 #ifndef MICROPY_PY_VFS
 #define MICROPY_PY_VFS (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_CORE_FEATURES && MICROPY_VFS)
diff --git a/tests/extmod/tls_dtls.py b/tests/extmod/tls_dtls.py
index b2d716769..a475cce8c 100644
--- a/tests/extmod/tls_dtls.py
+++ b/tests/extmod/tls_dtls.py
@@ -34,9 +34,19 @@ client_socket = DummySocket()
 # Wrap the DTLS Server
 dtls_server_ctx = SSLContext(PROTOCOL_DTLS_SERVER)
 dtls_server_ctx.verify_mode = CERT_NONE
-dtls_server = dtls_server_ctx.wrap_socket(server_socket, do_handshake_on_connect=False)
+dtls_server = dtls_server_ctx.wrap_socket(
+    server_socket, do_handshake_on_connect=False, client_id=b'dummy_client_id'
+)
 print("Wrapped DTLS Server")
 
+# wrap DTLS server with invalid client_id
+try:
+    dtls_server = dtls_server_ctx.wrap_socket(
+        server_socket, do_handshake_on_connect=False, client_id=4
+    )
+except OSError:
+    print("Failed to wrap DTLS Server with invalid client_id")
+
 # Wrap the DTLS Client
 dtls_client_ctx = SSLContext(PROTOCOL_DTLS_CLIENT)
 dtls_client_ctx.verify_mode = CERT_NONE
diff --git a/tests/extmod/tls_dtls.py.exp b/tests/extmod/tls_dtls.py.exp
index 78d72bff1..dbd005d0e 100644
--- a/tests/extmod/tls_dtls.py.exp
+++ b/tests/extmod/tls_dtls.py.exp
@@ -1,3 +1,4 @@
 Wrapped DTLS Server
+Failed to wrap DTLS Server with invalid client_id
 Wrapped DTLS Client
 OK