From 14219feb278508d71e592b133f3f8dd54864d9b7 Mon Sep 17 00:00:00 2001
From: Hanno Becker <hanno.becker@arm.com>
Date: Wed, 3 Jul 2019 17:02:43 +0100
Subject: [PATCH] Add IO wrappers to ssl_client2 as interm's between NET and
 SSL layer

---
 programs/ssl/ssl_client2.c | 89 +++++++++++++++++++++++++++++++-------
 1 file changed, 74 insertions(+), 15 deletions(-)

diff --git a/programs/ssl/ssl_client2.c b/programs/ssl/ssl_client2.c
index f0de38153..ad1936f43 100644
--- a/programs/ssl/ssl_client2.c
+++ b/programs/ssl/ssl_client2.c
@@ -490,7 +490,8 @@ static void my_debug( void *ctx, int level,
  * Test recv/send functions that make sure each try returns
  * WANT_READ/WANT_WRITE at least once before sucesseding
  */
-static int my_recv( void *ctx, unsigned char *buf, size_t len )
+
+static int delayed_recv( void *ctx, unsigned char *buf, size_t len )
 {
     static int first_try = 1;
     int ret;
@@ -507,7 +508,7 @@ static int my_recv( void *ctx, unsigned char *buf, size_t len )
     return( ret );
 }
 
-static int my_send( void *ctx, const unsigned char *buf, size_t len )
+static int delayed_send( void *ctx, const unsigned char *buf, size_t len )
 {
     static int first_try = 1;
     int ret;
@@ -527,6 +528,70 @@ static int my_send( void *ctx, const unsigned char *buf, size_t len )
           MBEDTLS_SSL_CONF_SEND &&
           MBEDTLS_SSL_CONF_RECV_TIMEOUT */
 
+typedef struct
+{
+    mbedtls_ssl_context *ssl;
+    mbedtls_net_context *net;
+} io_ctx_t;
+
+static int recv_cb( void *ctx, unsigned char *buf, size_t len )
+{
+    io_ctx_t *io_ctx = (io_ctx_t*) ctx;
+    size_t recv_len;
+    int ret;
+
+    if( opt.nbio == 2 )
+        ret = delayed_recv( io_ctx->net, buf, len );
+    else
+        ret = mbedtls_net_recv( io_ctx->net, buf, len );
+    if( ret < 0 )
+        return( ret );
+    recv_len = (size_t) ret;
+
+    if( opt.transport == MBEDTLS_SSL_TRANSPORT_DATAGRAM )
+    {
+        /* Here's the place to do any datagram/record checking
+         * in between receiving the packet from the underlying
+         * transport and passing it on to the TLS stack. */
+        mbedtls_printf( "[RECV] Datagram of size %u\n", (unsigned) recv_len );
+    }
+
+    return( (int) recv_len );
+}
+
+static int recv_timeout_cb( void *ctx, unsigned char *buf, size_t len,
+                            uint32_t timeout )
+{
+    io_ctx_t *io_ctx = (io_ctx_t*) ctx;
+    int ret;
+    size_t recv_len;
+
+    ret = mbedtls_net_recv_timeout( io_ctx->net, buf, len, timeout );
+    if( ret < 0 )
+        return( ret );
+    recv_len = (size_t) ret;
+
+    if( opt.transport == MBEDTLS_SSL_TRANSPORT_DATAGRAM )
+    {
+        /* Here's the place to do any datagram/record checking
+         * in between receiving the packet from the underlying
+         * transport and passing it on to the TLS stack. */
+        mbedtls_printf( "[RECV] Datagram of size %u\n", (unsigned) recv_len );
+    }
+
+    return( (int) recv_len );
+}
+
+static int send_cb( void *ctx, unsigned char const *buf, size_t len )
+{
+    io_ctx_t *io_ctx = (io_ctx_t*) ctx;
+
+    if( opt.nbio == 2 )
+        return( delayed_send( io_ctx->net, buf, len ) );
+
+    return( mbedtls_net_send( io_ctx->net, buf, len ) );
+}
+
 #if defined(MBEDTLS_X509_CRT_PARSE_C)
 static unsigned char peer_crt_info[1024];
 
@@ -761,6 +826,7 @@ int main( int argc, char *argv[] )
 {
     int ret = 0, len, tail_len, i, written, frags, retry_left;
     mbedtls_net_context server_fd;
+    io_ctx_t io_ctx;
 
     unsigned char buf[MAX_REQUEST_SIZE + 1];
 
@@ -1928,12 +1994,10 @@ int main( int argc, char *argv[] )
 #if !defined(MBEDTLS_SSL_CONF_RECV) && \
     !defined(MBEDTLS_SSL_CONF_SEND) && \
     !defined(MBEDTLS_SSL_CONF_RECV_TIMEOUT)
-    if( opt.nbio == 2 )
-        mbedtls_ssl_set_bio( &ssl, &server_fd, my_send, my_recv, NULL );
-    else
-        mbedtls_ssl_set_bio( &ssl, &server_fd,
-                             mbedtls_net_send, mbedtls_net_recv,
-                             opt.nbio == 0 ? mbedtls_net_recv_timeout : NULL );
+    io_ctx.ssl = &ssl;
+    io_ctx.net = &server_fd;
+    mbedtls_ssl_set_bio( &ssl, &io_ctx, send_cb, recv_cb,
+                         opt.nbio == 0 ? recv_timeout_cb : NULL );
 #else
      mbedtls_ssl_set_bio_ctx( &ssl, &server_fd );
 #endif
@@ -2550,13 +2614,8 @@ send_request:
                 goto exit;
             }
 
-            if( opt.nbio == 2 )
-                mbedtls_ssl_set_bio( &ssl, &server_fd, my_send, my_recv,
-                                     NULL );
-            else
-                mbedtls_ssl_set_bio( &ssl, &server_fd, mbedtls_net_send,
-                            mbedtls_net_recv,
-                            opt.nbio == 0 ? mbedtls_net_recv_timeout : NULL );
+            mbedtls_ssl_set_bio( &ssl, &io_ctx, send_cb, recv_cb,
+                                 opt.nbio == 0 ? recv_timeout_cb : NULL );
 
 #if defined(MBEDTLS_TIMING_C)
 #if !defined(MBEDTLS_SSL_CONF_SET_TIMER) && \