diff --git a/extmod/mbedtls/mbedtls_config_common.h b/extmod/mbedtls/mbedtls_config_common.h
index 4028bdf56d29b..c261f0cb5885d 100644
--- 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
index 5a9487ab89323..b79615bbf4d92 100644
--- 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
index fef25fe3087f5..51793350cda46 100644
--- 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
new file mode 100644
index 0000000000000..3af4ac97cf1fe
--- /dev/null
+++ 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
new file mode 100644
index 0000000000000..909c496d019e1
--- /dev/null
+++ 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'
pFad - Phonifier reborn
Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.
Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.
Alternative Proxies:
Alternative Proxy
pFad Proxy
pFad v3 Proxy
pFad v4 Proxy