Combining Zephyr's Socket API with MbedTLS for DTLS and Custom PSK Callback

Question

I am currently working on a project where I need to use Zephyr’s socket implementation in combination with MbedTLS for DTLS. Specifically, I want to integrate MbedTLS so that I can manipulate the handshake process, particularly to modify the PSK (Pre-Shared Key) ID during the DTLS handshake.

So far, in order to create a DTLS connection I used this:

sock_fd = zsock_socket(af, SOCK_DGRAM, IPPROTO_DTLS_1_2);

which to my understanding uses Zephyr's Socket API and MbedTLS internally for security.

However, I cannot use this now because I need to manually define a callback function. I have tried creating a UDP socket and manually use MbedTLS for security but it does not seem to work (I’m probably doing something wrong). Example:

    mbedtls_ssl_context ssl;
    mbedtls_ssl_config conf;

    mbedtls_ssl_init(&ssl);
    mbedtls_ssl_config_init(&conf);

    // Configure the DTLS parameters
    if (mbedtls_ssl_config_defaults(&conf, 
                                MBEDTLS_SSL_IS_SERVER, // Role: Server
                                MBEDTLS_SSL_TRANSPORT_DATAGRAM, // DTLS
                                MBEDTLS_SSL_PRESET_DEFAULT) != 0) {
        printf("Failed to configure SSL\n");
        return -1;
    }

    service->data->sock_fd = zsock_socket(af, SOCK_DGRAM, IPPROTO_UDP);
    
    const unsigned char psk[] = {0x01, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f};
    const unsigned char psk_id[] = "PSK_Identity";

    if (mbedtls_ssl_conf_psk(&conf, psk, sizeof(psk), psk_id, sizeof(psk_id)-1) != 0) {
        printf("Failed to configure PSK\n");
        goto end;
    }

    // Link socket with mbedTLS
    mbedtls_ssl_set_bio(&ssl, &service->data->sock_fd, mbedtls_net_send, mbedtls_net_recv, NULL);

    if (mbedtls_ssl_setup(&ssl, &conf) != 0) {
        printf("Failed to set up SSL\n");
        goto end;
    }

    ret = mbedtls_ssl_handshake(&ssl);
    if (ret != 0) {
        printf("Handshake failed: %d\n", ret);
        goto end;
    }

When I run this code I get the error: Handshake failed: -28928 which translates to this define:

#define MBEDTLS_ERR_SSL_BAD_INPUT_DATA -0x7100

Could you kindly provide guidance or resources on how to combine the native socket API of Zephyr with MbedTLS for DTLS? More specifically, I am interested in understanding how to set up a callback function that allows me to modify the PSK ID in the handshake.

NOTE: I use Zephyr version 3.5.99

************************************ EDIT ************************************

I have now come to the realization that I have defined mbedtls_net_send and mbedtls_net_recv as such:

int mbedtls_net_send(void *ctx, const unsigned char *buf, size_t len) {
    int sockfd = *(int *)ctx;  // 'ctx' is the socket descriptor
    return zsock_send(sockfd, buf, len, 0);  // Send data using Zephyr's socket API
}

// Implement mbedtls_net_recv to receive data from UDP socket
int mbedtls_net_recv(void *ctx, unsigned char *buf, size_t len) {
    int sockfd = *(int *)ctx;  // 'ctx' is the socket descriptor
    return zsock_recv(sockfd, buf, len, 0);  // Receive data using Zephyr's socket API
}

I'm guessing that I need to use sendto and recvfrom for my project to work properly since it is a DTLS connection. However, I am struggling with how to pass the extra arguments required which are basically the dest_addr and dest_addr_len. Can I include them in the context? What is the "right" way to do this?