Improve the TLS bidirectional shutdown

Author: dgarske

examples/tls/tls_client.c

@@ -639,9 +639,17 @@ int TPM2_TLS_ClientArgs(void* userCtx, int argc, char *argv[])
         printf("Failure %d (0x%x): %s\n", rc, rc, wolfTPM2_GetRCString(rc));
     }
 
-    /* Bidirectional shutdown */
-    while (wolfSSL_shutdown(ssl) == WOLFSSL_SHUTDOWN_NOT_DONE) {
-        printf("Shutdown not complete\n");
+    if (wolfSSL_shutdown(ssl) == WOLFSSL_SHUTDOWN_NOT_DONE) {
+        /* Bidirectional shutdown */
+        if (SocketWaitData(&sockIoCtx, 2 /* seconds */) == 1) {
+            int ret = wolfSSL_shutdown(ssl);
+            if (ret == WOLFSSL_SUCCESS) {
+                printf("Bidirectional shutdown complete\n");
+            }
+            else if (ret != WOLFSSL_SHUTDOWN_NOT_DONE) {
+                fprintf(stderr, "Bidirectional shutdown failed\n");
+            }
+        }
     }
 
     wolfSSL_free(ssl);

examples/tls/tls_common.h

@@ -322,6 +322,32 @@ static inline int SetupSocketAndConnect(SockIoCbCtx* sockIoCtx, const char* host
     return 0;
 }
 
+static inline int SocketWaitData(SockIoCbCtx* sockIoCtx, int timeout_sec)
+{
+    int res;
+    struct timeval timeout;
+    fd_set fds, errfds;
+    FD_ZERO(&fds);
+    FD_ZERO(&errfds);
+    FD_SET(sockIoCtx->fd, &fds);
+    FD_SET(sockIoCtx->fd, &errfds);
+    timeout.tv_sec = timeout_sec;
+    timeout.tv_usec = 0;
+    res = select(sockIoCtx->fd + 1, &fds, NULL, &errfds, &timeout);
+    if (res == 0) {
+        return 0; /* timeout */
+    }
+    else if (res > 0) {
+        if (FD_ISSET(sockIoCtx->fd, &fds)) {
+            return 1; /* ready to read */
+        }
+        else if (FD_ISSET(sockIoCtx->fd, &errfds)) {
+            return -1; /* error */
+        }
+    }
+    return 0; /* select failed */
+}
+
 static inline void CloseAndCleanupSocket(SockIoCbCtx* sockIoCtx)
 {
     if (sockIoCtx->fd != -1) {
@@ -343,6 +369,7 @@ static inline void CloseAndCleanupSocket(SockIoCbCtx* sockIoCtx)
 
     int SetupSocketAndListen(SockIoCbCtx* sockIoCtx, word32 port);
     int SocketWaitClient(SockIoCbCtx* sockIoCtx);
+    int SocketWaitData(SockIoCbCtx* sockIoCtx, int timeout_sec);
 #endif /* !WOLFSSL_USER_IO */
 
 /******************************************************************************/

examples/tls/tls_server.c

@@ -672,9 +672,17 @@ int TPM2_TLS_ServerArgs(void* userCtx, int argc, char *argv[])
             }
         }
 
+        if (wolfSSL_shutdown(ssl) == WOLFSSL_SHUTDOWN_NOT_DONE) {
         /* Bidirectional shutdown */
-        while (wolfSSL_shutdown(ssl) == WOLFSSL_SHUTDOWN_NOT_DONE) {
-            printf("Shutdown not complete\n");
+            if (SocketWaitData(&sockIoCtx, 2 /* seconds */) == 1) {
+                int ret = wolfSSL_shutdown(ssl);
+                if (ret == WOLFSSL_SUCCESS) {
+                    printf("Bidirectional shutdown complete\n");
+                }
+                else if (ret != WOLFSSL_SHUTDOWN_NOT_DONE) {
+                    fprintf(stderr, "Bidirectional shutdown failed\n");
+                }
+            }
         }
 
         wolfSSL_free(ssl);
@@ -691,11 +699,6 @@ int TPM2_TLS_ServerArgs(void* userCtx, int argc, char *argv[])
     }
 
     if (ssl != NULL) {
-        /* Bidirectional shutdown */
-        while (wolfSSL_shutdown(ssl) == WOLFSSL_SHUTDOWN_NOT_DONE) {
-            printf("Shutdown not complete\n");
-        }
-
         wolfSSL_free(ssl);
     }
     wolfSSL_CTX_free(ctx);