From 304d61cede3df47cc3a8b67d929c85541b9e2465 Mon Sep 17 00:00:00 2001
From: Jarno Lamsa <jarno.lamsa@arm.com>
Date: Thu, 6 Jun 2019 10:40:52 +0300
Subject: [PATCH] Add option for ssl-context re-initialization flow

---
 programs/ssl/ssl_client2.c | 30 +++++++++++++++++++++++++++---
 programs/ssl/ssl_server2.c | 29 ++++++++++++++++++++++++++---
 2 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/programs/ssl/ssl_client2.c b/programs/ssl/ssl_client2.c
index 37aecc8b0..47d7d73a6 100644
--- a/programs/ssl/ssl_client2.c
+++ b/programs/ssl/ssl_client2.c
@@ -346,7 +346,9 @@ int main( void )
 
 #if defined(MBEDTLS_SSL_CONTEXT_SERIALIZATION)
 #define USAGE_SERIALIZATION \
-    "    serialize=%%d        default: 0 (do not serialize/deserialize)\n"
+    "    serialize=%%d        default: 0 (do not serialize/deserialize)\n" \
+    "                         options: 1 (serialize)\n"                    \
+    "                                  2 (serialize with re-initialization)\n"
 #else
 #define USAGE_SERIALIZATION ""
 #endif
@@ -1589,7 +1591,7 @@ int main( int argc, char *argv[] )
         else if( strcmp( p, "serialize") == 0 )
         {
             opt.serialize = atoi( q );
-            if( opt.serialize < 0 || opt.serialize > 1)
+            if( opt.serialize < 0 || opt.serialize > 2)
                 goto usage;
         }
         else if( strcmp( p, "eap_tls" ) == 0 )
@@ -2917,7 +2919,7 @@ send_request:
      * 7c. Simulate serialize/deserialize and go back to data exchange
      */
 #if defined(MBEDTLS_SSL_CONTEXT_SERIALIZATION)
-    if( opt.serialize != 0)
+    if( opt.serialize != 0 )
     {
         size_t len;
         unsigned char *buf = NULL;
@@ -2953,6 +2955,28 @@ send_request:
             goto exit;
         }
 
+        if( opt.serialize == 2 )
+        {
+            mbedtls_ssl_free( &ssl );
+
+            mbedtls_ssl_init( &ssl );
+
+            if( ( ret = mbedtls_ssl_setup( &ssl, &conf ) ) != 0 )
+            {
+                mbedtls_printf( " failed\n  ! mbedtls_ssl_setup returned -0x%x\n\n",
+                                -ret );
+                goto exit;
+            }
+
+            if( opt.nbio == 2 )
+                mbedtls_ssl_set_bio( &ssl, &server_fd, delayed_send, delayed_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_printf( " Deserializing connection..." );
 
         if( ( ret = mbedtls_ssl_context_load( &ssl, buf, len ) ) != 0 )
diff --git a/programs/ssl/ssl_server2.c b/programs/ssl/ssl_server2.c
index f3858044a..8502a6a1a 100644
--- a/programs/ssl/ssl_server2.c
+++ b/programs/ssl/ssl_server2.c
@@ -439,7 +439,9 @@ int main( void )
 
 #if defined(MBEDTLS_SSL_CONTEXT_SERIALIZATION)
 #define USAGE_SERIALIZATION \
-    "    serialize=%%d        default: 0 (do not serialize/deserialize)\n"
+    "    serialize=%%d        default: 0 (do not serialize/deserialize)\n" \
+    "                         options: 1 (serialize)\n"                    \
+    "                                  2 (serialize with re-initialization)\n"
 #else
 #define USAGE_SERIALIZATION ""
 #endif
@@ -2301,7 +2303,7 @@ int main( int argc, char *argv[] )
         else if( strcmp( p, "serialize") == 0 )
         {
             opt.serialize = atoi( q );
-            if( opt.serialize < 0 || opt.serialize > 1)
+            if( opt.serialize < 0 || opt.serialize > 2)
                 goto usage;
         }
         else if( strcmp( p, "eap_tls" ) == 0 )
@@ -3928,7 +3930,7 @@ data_exchange:
      * 7b. Simulate serialize/deserialize and go back to data exchange
      */
 #if defined(MBEDTLS_SSL_CONTEXT_SERIALIZATION)
-    if( opt.serialize != 0)
+    if( opt.serialize != 0 )
     {
         size_t len;
         unsigned char *buf = NULL;
@@ -3964,6 +3966,27 @@ data_exchange:
             goto exit;
         }
 
+        if( opt.serialize == 2 )
+        {
+            mbedtls_ssl_free( &ssl );
+
+            mbedtls_ssl_init( &ssl );
+
+            if( ( ret = mbedtls_ssl_setup( &ssl, &conf ) ) != 0 )
+            {
+                mbedtls_printf( " failed\n  ! mbedtls_ssl_setup returned -0x%x\n\n",
+                                -ret );
+                goto exit;
+            }
+
+            if( opt.nbio == 2 )
+                mbedtls_ssl_set_bio( &ssl, &client_fd, delayed_send, delayed_recv, NULL );
+            else
+                mbedtls_ssl_set_bio( &ssl, &client_fd, mbedtls_net_send, mbedtls_net_recv,
+                                    opt.nbio == 0 ? mbedtls_net_recv_timeout : NULL );
+
+        }
+
         mbedtls_printf( " Deserializing connection..." );
 
         if( ( ret = mbedtls_ssl_context_load( &ssl, buf, len ) ) != 0 )