micropython · aminop1us · Apr 30, 2024 · Apr 30, 2024 · Apr 30, 2024 · Apr 30, 2024
diff --git a/extmod/mbedtls/mbedtls_config_common.h b/extmod/mbedtls/mbedtls_config_common.h
@@ -47,6 +47,7 @@
 #define MBEDTLS_KEY_EXCHANGE_RSA_ENABLED
 #define MBEDTLS_KEY_EXCHANGE_ECDHE_ECDSA_ENABLED
 #define MBEDTLS_KEY_EXCHANGE_ECDHE_RSA_ENABLED
+#define MBEDTLS_KEY_EXCHANGE_PSK_ENABLED
 #define MBEDTLS_CAN_ECDH
 #define MBEDTLS_PK_CAN_ECDSA_SIGN
 #define MBEDTLS_PKCS1_V15

diff --git a/extmod/modtls_mbedtls.c b/extmod/modtls_mbedtls.c
@@ -199,6 +199,24 @@ static int ssl_sock_cert_verify(void *ptr, mbedtls_x509_crt *crt, int depth, uin
     return mp_obj_get_int(mp_call_function_2(o->handler, MP_OBJ_FROM_PTR(&cert), MP_OBJ_NEW_SMALL_INT(depth)));
 }
 
+static int ssl_conf_psk_cb(void *parameter, mbedtls_ssl_context *ssl,
+    const unsigned char *psk_identity, size_t psk_identity_len) {
+    mp_obj_t callback = MP_OBJ_FROM_PTR(parameter);
+    mp_obj_t psk_identity_obj = mp_obj_new_bytes(psk_identity, psk_identity_len);
+
+    // Call the callback
+    mp_obj_t psk_key_obj = mp_call_function_1(callback, psk_identity_obj);
+
+    // Check if psk key object is supplied, set handshake psk
+    int ret = -1;
+    if (psk_key_obj != mp_const_none) {
+        size_t psk_key_len;
+        const unsigned char *psk_key = (const unsigned char *)mp_obj_str_get_data(psk_key_obj, &psk_key_len);
+        ret = mbedtls_ssl_set_hs_psk(ssl, psk_key, psk_key_len);
+    }
+    return ret;
+}
+
 /******************************************************************************/
 // SSLContext type.
 
@@ -401,6 +419,36 @@ static mp_obj_t ssl_context_load_verify_locations(mp_obj_t self_in, mp_obj_t cad
 }
 static MP_DEFINE_CONST_FUN_OBJ_2(ssl_context_load_verify_locations_obj, ssl_context_load_verify_locations);
 
+// SSLContext.set_psk_server_callback(callback)
+static mp_obj_t ssl_context_set_psk_server_callback(mp_obj_t self_in, mp_obj_t callback) {
+    mp_obj_ssl_context_t *self = MP_OBJ_TO_PTR(self_in);
+
+    // Check if callback is supplied, config psk callback
+    if (callback != mp_const_none) {
+        mbedtls_ssl_conf_psk_cb(&self->conf, ssl_conf_psk_cb, MP_OBJ_TO_PTR(callback));
+    }
+    return mp_const_none;
+}
+static MP_DEFINE_CONST_FUN_OBJ_2(ssl_context_set_psk_server_callback_obj, ssl_context_set_psk_server_callback);
+
+// SSLContext.set_psk_client_callback(callback)
+static mp_obj_t ssl_context_set_psk_client_callback(mp_obj_t self_in, mp_obj_t callback) {
+    mp_obj_ssl_context_t *self = MP_OBJ_TO_PTR(self_in);
+
+    // Check if callback is supplied, config psk
+    if (callback != mp_const_none) {
+        mp_obj_tuple_t *psk_pairwise_obj = MP_OBJ_TO_PTR(mp_call_function_1(callback, mp_const_none));
+        if (psk_pairwise_obj->len == 2) {
+            size_t psk_identity_len, psk_key_len;
+            const unsigned char *psk_identity = (const unsigned char *)mp_obj_str_get_data(psk_pairwise_obj->items[0], &psk_identity_len);
+            const unsigned char *psk_key = (const unsigned char *)mp_obj_str_get_data(psk_pairwise_obj->items[1], &psk_key_len);
+            mbedtls_ssl_conf_psk(&self->conf, psk_key, psk_key_len, psk_identity, psk_identity_len);
+        }
+    }
+    return mp_const_none;
+}
+static MP_DEFINE_CONST_FUN_OBJ_2(ssl_context_set_psk_client_callback_obj, ssl_context_set_psk_client_callback);
+
 static mp_obj_t ssl_context_wrap_socket(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
     enum { ARG_server_side, ARG_do_handshake_on_connect, ARG_server_hostname };
     static const mp_arg_t allowed_args[] = {
@@ -429,6 +477,8 @@ static const mp_rom_map_elem_t ssl_context_locals_dict_table[] = {
     { MP_ROM_QSTR(MP_QSTR_set_ciphers), MP_ROM_PTR(&ssl_context_set_ciphers_obj)},
     { MP_ROM_QSTR(MP_QSTR_load_cert_chain), MP_ROM_PTR(&ssl_context_load_cert_chain_obj)},
     { MP_ROM_QSTR(MP_QSTR_load_verify_locations), MP_ROM_PTR(&ssl_context_load_verify_locations_obj)},
+    { MP_ROM_QSTR(MP_QSTR_set_psk_server_callback), MP_ROM_PTR(&ssl_context_set_psk_server_callback_obj)},
+    { MP_ROM_QSTR(MP_QSTR_set_psk_client_callback), MP_ROM_PTR(&ssl_context_set_psk_client_callback_obj)},
     { MP_ROM_QSTR(MP_QSTR_wrap_socket), MP_ROM_PTR(&ssl_context_wrap_socket_obj) },
 };
 static MP_DEFINE_CONST_DICT(ssl_context_locals_dict, ssl_context_locals_dict_table);

diff --git a/ports/esp32/boards/sdkconfig.base b/ports/esp32/boards/sdkconfig.base
@@ -85,6 +85,13 @@ CONFIG_MBEDTLS_CERTIFICATE_BUNDLE_DEFAULT_FULL=n
 # significant factor in overall physical security.)
 CONFIG_MBEDTLS_DEFAULT_MEM_ALLOC=y
 
+# Enable MBEDTLS to support TLS-PSK
+CONFIG_MBEDTLS_PSK_MODES=y
+CONFIG_MBEDTLS_KEY_EXCHANGE_PSK=y
+CONFIG_MBEDTLS_KEY_EXCHANGE_DHE_PSK=y
+CONFIG_MBEDTLS_KEY_EXCHANGE_ECDHE_PSK=y
+CONFIG_MBEDTLS_KEY_EXCHANGE_RSA_PSK=y
+
 # ULP coprocessor support
 # Only on: ESP32, ESP32S2, ESP32S3
 CONFIG_ULP_COPROC_ENABLED=y

diff --git a/tests/multi_net/sslcontext_server_client_with_psk.py b/tests/multi_net/sslcontext_server_client_with_psk.py
@@ -0,0 +1,53 @@
+# Test asyncio TCP server and client with TLS-PSK, using set_psk_sever_callback() and set_psk_client_callback().
+
+try:
+    import socket
+    import tls
+except ImportError:
+    print("SKIP")
+    raise SystemExit
+
+PORT = 8000
+
+
+# Server
+def instance0():
+    multitest.globals(IP=multitest.get_network_ip())
+    s = socket.socket()
+    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    s.bind(socket.getaddrinfo("0.0.0.0", PORT)[0][-1])
+    s.listen(1)
+    multitest.next()
+    s2, _ = s.accept()
+    server_ctx = tls.SSLContext(tls.PROTOCOL_TLS_SERVER)
+
+    def psk_server_callback(identity):
+        psk_dict = {
+            b"PSK-Identity-1": bytes.fromhex("c0ffee"),
+            b"PSK-Identity-2": bytes.fromhex("facade"),
+        }
+        return psk_dict[identity]
+
+    server_ctx.set_psk_server_callback(psk_server_callback)
+    s2 = server_ctx.wrap_socket(s2, server_side=True)
+    print(s2.read(16))
+    s2.write(b"server to client")
+    s2.close()
+    s.close()
+
+
+# Client
+def instance1():
+    multitest.next()
+    s = socket.socket()
+    s.connect(socket.getaddrinfo(IP, PORT)[0][-1])
+    client_ctx = tls.SSLContext(tls.PROTOCOL_TLS_CLIENT)
+
+    def psk_client_callback(identity):
+        return (b"PSK-Identity-1", bytes.fromhex("c0ffee"))
+
+    client_ctx.set_psk_client_callback(psk_client_callback)
+    s = client_ctx.wrap_socket(s, server_hostname="micropython.local")
+    s.write(b"client to server")
+    print(s.read(16))
+    s.close()
diff --git a/tests/multi_net/sslcontext_server_client_with_psk.py.exp b/tests/multi_net/sslcontext_server_client_with_psk.py.exp
@@ -0,0 +1,4 @@
+--- instance0 ---
+b'client to server'
+--- instance1 ---
+b'server to client'