From e08660e612de78744ba89aad7284d01ed400ae3f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20P=C3=A9gouri=C3=A9-Gonnard?= <mpg@elzevir.fr>
Date: Sat, 16 Aug 2014 11:28:40 +0200
Subject: [PATCH] Fix ssl_read() and close_notify error handling in programs

---
 programs/ssl/ssl_client2.c | 45 +++++++++++++++++++++++---------------
 programs/ssl/ssl_server2.c | 31 ++++++++++++++++++--------
 2 files changed, 49 insertions(+), 27 deletions(-)

diff --git a/programs/ssl/ssl_client2.c b/programs/ssl/ssl_client2.c
index 3394f7235f..dda4263014 100644
--- a/programs/ssl/ssl_client2.c
+++ b/programs/ssl/ssl_client2.c
@@ -1108,23 +1108,29 @@ send_request:
         memset( buf, 0, sizeof( buf ) );
         ret = ssl_read( &ssl, buf, len );
 
-        if( ret == POLARSSL_ERR_NET_WANT_READ || ret == POLARSSL_ERR_NET_WANT_WRITE )
+        if( ret == POLARSSL_ERR_NET_WANT_READ ||
+            ret == POLARSSL_ERR_NET_WANT_WRITE )
             continue;
 
-        if( ret == POLARSSL_ERR_SSL_PEER_CLOSE_NOTIFY )
-            break;
-
-        if( ret < 0 )
+        if( ret <= 0 )
         {
-            printf( "failed\n  ! ssl_read returned -0x%x\n\n", -ret );
-            break;
-        }
+            switch( ret )
+            {
+                case POLARSSL_ERR_SSL_PEER_CLOSE_NOTIFY:
+                    printf( " connection was closed gracefully\n" );
+                    ret = 0;
+                    goto reconnect;
 
-        if( ret == 0 )
-        {
-            printf("\n\nEOF\n\n");
-            ssl_close_notify( &ssl );
-            break;
+                case 0:
+                case POLARSSL_ERR_NET_CONN_RESET:
+                    printf( " connection was reset by peer\n" );
+                    ret = 0;
+                    goto reconnect;
+
+                default:
+                    printf( " ssl_read returned -0x%x\n", -ret );
+                    goto exit;
+            }
         }
 
         len = ret;
@@ -1133,6 +1139,10 @@ send_request:
     }
     while( 1 );
 
+    /*
+     * 9. Reconnect?
+     */
+reconnect:
     if( opt.reconnect != 0 )
     {
         --opt.reconnect;
@@ -1181,10 +1191,10 @@ send_request:
         goto send_request;
     }
 
+    /*
+     * Cleanup and exit
+     */
 exit:
-    if( ret == POLARSSL_ERR_SSL_PEER_CLOSE_NOTIFY )
-        ret = 0;
-
 #ifdef POLARSSL_ERROR_C
     if( ret != 0 )
     {
@@ -1196,6 +1206,7 @@ exit:
 
     if( server_fd )
         net_close( server_fd );
+
 #if defined(POLARSSL_X509_CRT_PARSE_C)
     x509_crt_free( &clicert );
     x509_crt_free( &cacert );
@@ -1206,8 +1217,6 @@ exit:
     ctr_drbg_free( &ctr_drbg );
     entropy_free( &entropy );
 
-    memset( &ssl, 0, sizeof( ssl ) );
-
 #if defined(_WIN32)
     printf( "  + Press Enter to exit this program.\n" );
     fflush( stdout ); getchar();
diff --git a/programs/ssl/ssl_server2.c b/programs/ssl/ssl_server2.c
index 91522135ec..20df7819b8 100644
--- a/programs/ssl/ssl_server2.c
+++ b/programs/ssl/ssl_server2.c
@@ -1514,7 +1514,8 @@ reset:
         memset( buf, 0, sizeof( buf ) );
         ret = ssl_read( &ssl, buf, len );
 
-        if( ret == POLARSSL_ERR_NET_WANT_READ || ret == POLARSSL_ERR_NET_WANT_WRITE )
+        if( ret == POLARSSL_ERR_NET_WANT_READ ||
+            ret == POLARSSL_ERR_NET_WANT_WRITE )
             continue;
 
         if( ret <= 0 )
@@ -1523,18 +1524,18 @@ reset:
             {
                 case POLARSSL_ERR_SSL_PEER_CLOSE_NOTIFY:
                     printf( " connection was closed gracefully\n" );
-                    break;
+                    goto close_notify;
 
+                case 0:
                 case POLARSSL_ERR_NET_CONN_RESET:
                     printf( " connection was reset by peer\n" );
-                    break;
+                    ret = POLARSSL_ERR_NET_CONN_RESET;
+                    goto reset;
 
                 default:
                     printf( " ssl_read returned -0x%x\n", -ret );
-                    break;
+                    goto close_notify;
             }
-
-            break;
         }
 
         if( ssl_get_bytes_avail( &ssl ) == 0 )
@@ -1658,10 +1659,22 @@ reset:
         printf( " ok\n" );
     }
 
+    /*
+     * 8. Close the connection cleanly
+     */
+close_notify:
     printf( "  . Closing the connection..." );
 
     while( ( ret = ssl_close_notify( &ssl ) ) < 0 )
     {
+        printf( " ret = %d (-0x%04X)", ret, -ret );
+        if( ret == POLARSSL_ERR_NET_CONN_RESET )
+        {
+            printf( " connection was reset by peer\n" );
+            ret = 0;
+            goto reset;
+        }
+
         if( ret != POLARSSL_ERR_NET_WANT_READ &&
             ret != POLARSSL_ERR_NET_WANT_WRITE )
         {
@@ -1671,12 +1684,12 @@ reset:
     }
 
     printf( " ok\n" );
-
-    ret = 0;
     goto reset;
 
+    /*
+     * Cleanup and exit
+     */
 exit:
-
 #ifdef POLARSSL_ERROR_C
     if( ret != 0 )
     {