extmod/modussl_mbedtls: Add support for PSK authentication

damz · damz · commit f1d2b811be3b · 2022-12-14T20:52:12.000-08:00
diff --git a/extmod/modussl_mbedtls.c b/extmod/modussl_mbedtls.c
@@ -75,6 +75,8 @@ struct ssl_args {
     mp_arg_val_t cadata;
     mp_arg_val_t do_handshake;
     mp_arg_val_t dtls;
+    mp_arg_val_t psk_identity;
+    mp_arg_val_t psk_key;
 };
 
 STATIC const mp_obj_type_t ussl_socket_type;
@@ -242,6 +244,22 @@ STATIC mp_obj_ssl_socket_t *socket_new(mp_obj_t sock, struct ssl_args *args) {
     mbedtls_ssl_conf_dbg(&o->conf, mbedtls_debug, NULL);
     #endif
 
+    #if defined(MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED) || defined(MBEDTLS_KEY_EXCHANGE__SOME__PSK_ENABLED)
+    // banana();
+    if (args->psk_identity.u_obj != mp_const_none && args->psk_key.u_obj != mp_const_none) {
+        size_t psk_identity_len;
+        size_t psk_key_len;
+        const byte *psk_identity = (const byte *)mp_obj_str_get_data(args->psk_identity.u_obj, &psk_identity_len);
+        const byte *psk_key = (const byte *)mp_obj_str_get_data(args->psk_key.u_obj, &psk_key_len);
+        // len should include terminating null
+        ret = mbedtls_ssl_conf_psk(&o->conf, (const unsigned char *) psk_key, psk_key_len, (const unsigned char *) psk_identity, psk_identity_len);
+        if (ret != 0) {
+            ret = MBEDTLS_ERR_PK_BAD_INPUT_DATA; // use general error for all key errors
+            goto cleanup;
+        }
+    }
+    #endif
+
     ret = mbedtls_ssl_setup(&o->ssl, &o->conf);
     if (ret != 0) {
         goto cleanup;
@@ -466,6 +484,8 @@ STATIC mp_obj_t mod_ssl_wrap_socket(size_t n_args, const mp_obj_t *pos_args, mp_
         { 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} },
+        { MP_QSTR_psk_identity, MP_ARG_KW_ONLY | MP_ARG_OBJ, {.u_rom_obj = MP_ROM_NONE} },
+        { MP_QSTR_psk_key, MP_ARG_KW_ONLY | MP_ARG_OBJ, {.u_rom_obj = MP_ROM_NONE} },
     };
 
     // TODO: Check that sock implements stream protocol
diff --git a/ports/esp32/boards/sdkconfig.base b/ports/esp32/boards/sdkconfig.base
@@ -47,6 +47,8 @@ CONFIG_LWIP_PPP_CHAP_SUPPORT=y
 # 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
+CONFIG_MBEDTLS_PSK_MODES=y
+CONFIG_MBEDTLS_KEY_EXCHANGE_PSK=y
 
 # ULP coprocessor support
 CONFIG_ESP32_ULP_COPROC_ENABLED=y
