From dc66602e12487c12a4528e3974b9c3b36cfc9a45 Mon Sep 17 00:00:00 2001
From: JacobBarthelmeh <jacob@wolfssl.com>
Date: Fri, 17 May 2024 16:11:30 -0600
Subject: [PATCH] add example auth pending

---
 examples/echoserver/echoserver.c | 31 +++++++++++++++++++++++++------
 1 file changed, 25 insertions(+), 6 deletions(-)

diff --git a/examples/echoserver/echoserver.c b/examples/echoserver/echoserver.c
index 0951c13c3..fe82d57e1 100644
--- a/examples/echoserver/echoserver.c
+++ b/examples/echoserver/echoserver.c
@@ -1355,7 +1355,8 @@ static int NonBlockSSH_accept(WOLFSSH* ssh)
 
     while ((ret != WS_SUCCESS
                 && ret != WS_SCP_COMPLETE && ret != WS_SFTP_COMPLETE)
-            && (error == WS_WANT_READ || error == WS_WANT_WRITE)) {
+            && (error == WS_WANT_READ || error == WS_WANT_WRITE ||
+                error == WS_AUTH_PENDING)) {
 
         if (error == WS_WANT_READ)
             printf("... server would read block\n");
@@ -1365,7 +1366,8 @@ static int NonBlockSSH_accept(WOLFSSH* ssh)
         select_ret = tcp_select(sockfd, 1);
         if (select_ret == WS_SELECT_RECV_READY  ||
             select_ret == WS_SELECT_ERROR_READY ||
-            error      == WS_WANT_WRITE)
+            error      == WS_WANT_WRITE ||
+            error      == WS_AUTH_PENDING)
         {
             ret = wolfSSH_accept(ssh);
             error = wolfSSH_get_error(ssh);
@@ -1387,11 +1389,16 @@ static THREAD_RETURN WOLFSSH_THREAD server_worker(void* vArgs)
 
     passwdRetry = MAX_PASSWD_RETRY;
 
-    if (!threadCtx->nonBlock)
+    if (!threadCtx->nonBlock) {
         ret = wolfSSH_accept(threadCtx->ssh);
-    else
+        if (wolfSSH_get_error(threadCtx->ssh) == WS_AUTH_PENDING) {
+            printf("Auth pending error, use -N for non blocking\n");
+            printf("Trying to close down the connection\n");
+        }
+    }
+    else {
         ret = NonBlockSSH_accept(threadCtx->ssh);
-
+    }
 #ifdef WOLFSSH_SCP
     /* finish off SCP operation */
     if (ret == WS_SCP_INIT) {
@@ -2010,6 +2017,7 @@ static int wsUserAuthResult(byte res,
 }
 
 
+static int userAuthWouldBlock = 0;
 static int wsUserAuth(byte authType,
                       WS_UserAuthData* authData,
                       void* ctx)
@@ -2023,6 +2031,12 @@ static int wsUserAuth(byte authType,
         return WOLFSSH_USERAUTH_FAILURE;
     }
 
+    if (userAuthWouldBlock > 0) {
+        printf("User Auth would block ....\n");
+        userAuthWouldBlock--;
+        return WOLFSSH_USERAUTH_WOULD_BLOCK;
+    }
+
     if (authType != WOLFSSH_USERAUTH_PASSWORD &&
 #ifdef WOLFSSH_ALLOW_USERAUTH_NONE
         authType != WOLFSSH_USERAUTH_NONE &&
@@ -2239,6 +2253,7 @@ static void ShowUsage(void)
     printf(" -a <file>     load in a root CA certificate file\n");
 #endif
     printf(" -k            set the list of key algos to use\n");
+    printf(" -b <num>      test user auth would block\n");
 }
 
 
@@ -2300,7 +2315,7 @@ THREAD_RETURN WOLFSSH_THREAD echoserver_test(void* args)
     serverArgs->return_code = EXIT_SUCCESS;
 
     if (argc > 0) {
-        const char* optlist = "?1a:d:efEp:R:Ni:j:I:J:K:P:k:";
+        const char* optlist = "?1a:d:efEp:R:Ni:j:I:J:K:P:k:b:";
         myoptind = 0;
         while ((ch = mygetopt(argc, argv, optlist)) != -1) {
             switch (ch) {
@@ -2384,6 +2399,10 @@ THREAD_RETURN WOLFSSH_THREAD echoserver_test(void* args)
                     passwdList = StrListAdd(passwdList, myoptarg);
                     break;
 
+                case 'b':
+                    userAuthWouldBlock = atoi(myoptarg);
+                    break;
+
                 default:
                     ShowUsage();
                     serverArgs->return_code = MY_EX_USAGE;