From 12f3c0f51f1d84b83c7d78bcfdba60ede036d792 Mon Sep 17 00:00:00 2001
From: Damien Tournoud <damien@platform.sh>
Date: Sat, 6 Jan 2024 16:10:44 -0800
Subject: [PATCH] extmod/modussl_mbedtls: Wire in support for PSK modes.

Signed-off-by: Damien Tournoud <damien@platform.sh>
---
 extmod/modssl_mbedtls.c             | 25 +++++++++++++++++++++++++
 ports/esp32/boards/sdkconfig.base   |  3 +++
 ports/unix/mbedtls/mbedtls_config.h |  4 ++++
 3 files changed, 32 insertions(+)

diff --git a/extmod/modssl_mbedtls.c b/extmod/modssl_mbedtls.c
index feb6d04306aa..9688c6869e47 100644
--- a/extmod/modssl_mbedtls.c
+++ b/extmod/modssl_mbedtls.c
@@ -373,6 +373,28 @@ STATIC mp_obj_t ssl_context_set_ciphers(mp_obj_t self_in, mp_obj_t ciphersuite)
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(ssl_context_set_ciphers_obj, ssl_context_set_ciphers);
 
+#ifdef MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED
+// SSLContext.set_psk(username, password)
+STATIC mp_obj_t ssl_context_set_psk(mp_obj_t self_in, mp_obj_t username, mp_obj_t password) {
+    mp_obj_ssl_context_t *ssl_context = MP_OBJ_TO_PTR(self_in);
+    int ret;
+
+    size_t psk_identity_len;
+    const byte *psk_identity = (const byte *)mp_obj_str_get_data(username, &psk_identity_len);
+
+    size_t psk_key_len;
+    const byte *psk_key = (const byte *)mp_obj_str_get_data(password, &psk_key_len);
+
+    ret = mbedtls_ssl_conf_psk(&ssl_context->conf, (const unsigned char *) psk_key, psk_key_len, (const unsigned char *) psk_identity, psk_identity_len);
+    if (ret != 0) {
+        mbedtls_raise_error(ret);
+    }
+
+    return mp_const_none;
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_3(ssl_context_set_psk_obj, ssl_context_set_psk);
+#endif
+
 STATIC void ssl_context_load_key(mp_obj_ssl_context_t *self, mp_obj_t key_obj, mp_obj_t cert_obj) {
     size_t key_len;
     const byte *key = (const byte *)mp_obj_str_get_data(key_obj, &key_len);
@@ -487,6 +509,9 @@ STATIC const mp_rom_map_elem_t ssl_context_locals_dict_table[] = {
     #endif
     { MP_ROM_QSTR(MP_QSTR_get_ciphers), MP_ROM_PTR(&ssl_context_get_ciphers_obj)},
     { MP_ROM_QSTR(MP_QSTR_set_ciphers), MP_ROM_PTR(&ssl_context_set_ciphers_obj)},
+    #ifdef MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED
+    { MP_ROM_QSTR(MP_QSTR_set_psk), MP_ROM_PTR(&ssl_context_set_psk_obj)},
+    #endif
     { 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_wrap_socket), MP_ROM_PTR(&ssl_context_wrap_socket_obj) },
diff --git a/ports/esp32/boards/sdkconfig.base b/ports/esp32/boards/sdkconfig.base
index 4eb7e165fc33..18e0ab81d65d 100644
--- a/ports/esp32/boards/sdkconfig.base
+++ b/ports/esp32/boards/sdkconfig.base
@@ -64,6 +64,9 @@ CONFIG_MBEDTLS_PLATFORM_TIME_ALT=y
 CONFIG_MBEDTLS_HAVE_TIME=y
 # Enable DTLS
 CONFIG_MBEDTLS_SSL_PROTO_DTLS=y
+# Enable PSK support
+CONFIG_MBEDTLS_PSK_MODES=y
+CONFIG_MBEDTLS_KEY_EXCHANGE_PSK=y
 
 # Disable ALPN support as it's not implemented in MicroPython
 CONFIG_MBEDTLS_SSL_ALPN=n
diff --git a/ports/unix/mbedtls/mbedtls_config.h b/ports/unix/mbedtls/mbedtls_config.h
index 88f31469c121..91f98627581b 100644
--- a/ports/unix/mbedtls/mbedtls_config.h
+++ b/ports/unix/mbedtls/mbedtls_config.h
@@ -32,6 +32,10 @@
 // Enable DTLS
 #define MBEDTLS_SSL_PROTO_DTLS
 
+// Enable PSK modes and ciphers
+#define MBEDTLS_PSK_MODES
+#define MBEDTLS_KEY_EXCHANGE_PSK_ENABLED
+
 // Enable mbedtls modules
 #define MBEDTLS_HAVEGE_C
 #define MBEDTLS_TIMING_C