From c0c6da75a6fb474740a6fb99f207f44945ccf909 Mon Sep 17 00:00:00 2001
From: JacobBarthelmeh <jacob@wolfssl.com>
Date: Tue, 27 May 2025 17:11:27 -0600
Subject: [PATCH] adjustment to client side keyboard auth behavior and auth
 test case

---
 src/internal.c     | 19 +++++++++++++++----
 tests/auth.c       | 26 +++++++++++++++-----------
 wolfssh/internal.h |  1 +
 3 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index 2cfce115..d6caa820 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -7878,7 +7878,10 @@ static int DoUserAuthFailure(WOLFSSH* ssh,
                             break;
 #ifdef WOLFSSH_KEYBOARD_INTERACTIVE
                         case ID_USERAUTH_KEYBOARD:
-                            authType |= WOLFSSH_USERAUTH_KEYBOARD;
+                            /* try a different auth method if failing */
+                            if (ssh->kbAuthAttempts < 3) {
+                                authType |= WOLFSSH_USERAUTH_KEYBOARD;
+                            }
                             break;
 #endif
 #if !defined(WOLFSSH_NO_RSA) || !defined(WOLFSSH_NO_ECDSA)
@@ -13382,6 +13385,11 @@ int SendUserAuthKeyboardRequest(WOLFSSH* ssh, WS_UserAuthData* authData)
         if (ret == WOLFSSH_USERAUTH_SUCCESS) {
             ret = WS_SUCCESS;
         }
+        else {
+            WLOG(WS_LOG_DEBUG, "Issue with keyboard auth setup, try another "
+                "auth type");
+            return  SendUserAuthFailure(ssh, 0);
+        }
     }
 
     if (authData->sf.keyboard.promptCount > 0 &&
@@ -13407,10 +13415,12 @@ int SendUserAuthKeyboardRequest(WOLFSSH* ssh, WS_UserAuthData* authData)
         ret = PreparePacket(ssh, payloadSz);
     }
 
-    output = ssh->outputBuffer.buffer;
-    idx = ssh->outputBuffer.length;
+    if (ret == WS_SUCCESS) {
+        output = ssh->outputBuffer.buffer;
+        idx = ssh->outputBuffer.length;
 
-    output[idx++] = MSGID_USERAUTH_INFO_REQUEST;
+        output[idx++] = MSGID_USERAUTH_INFO_REQUEST;
+    }
 
     if (ret == WS_SUCCESS) {
         ret = BuildUserAuthRequestKeyboard(ssh, output, &idx, authData);
@@ -15075,6 +15085,7 @@ int SendUserAuthRequest(WOLFSSH* ssh, byte authType, int addSig)
             /* submethods */
             c32toa(0, output + idx);
             idx += LENGTH_SZ;
+            ssh->kbAuthAttempts++;
         }
 #endif
         else if (authId == ID_USERAUTH_PUBLICKEY)
diff --git a/tests/auth.c b/tests/auth.c
index cf57a3c1..10dd1bae 100644
--- a/tests/auth.c
+++ b/tests/auth.c
@@ -222,11 +222,23 @@ static int load_key(byte isEcc, byte* buf, word32 bufSz)
 
 static int serverUserAuth(byte authType, WS_UserAuthData* authData, void* ctx)
 {
-    (void) ctx;
-    if (authType != WOLFSSH_USERAUTH_KEYBOARD) {
+    WS_UserAuthData_Keyboard* prompts = (WS_UserAuthData_Keyboard*)ctx;
+
+    if (ctx == NULL) {
         return WOLFSSH_USERAUTH_FAILURE;
     }
 
+    if (authType != WOLFSSH_USERAUTH_KEYBOARD &&
+            authType != WOLFSSH_USERAUTH_KEYBOARD_SETUP) {
+        return WOLFSSH_USERAUTH_FAILURE;
+    }
+
+    if (authType == WOLFSSH_USERAUTH_KEYBOARD_SETUP) {
+        WMEMCPY(&authData->sf.keyboard, prompts,
+            sizeof(WS_UserAuthData_Keyboard));
+        return WS_SUCCESS;
+    }
+
     if (authData->sf.keyboard.responseCount != kbResponseCount) {
         return WOLFSSH_USERAUTH_FAILURE;
     }
@@ -251,14 +263,6 @@ static int serverUserAuth(byte authType, WS_UserAuthData* authData, void* ctx)
     return WOLFSSH_USERAUTH_SUCCESS;
 }
 
-static int serverKeyboardCallback(WS_UserAuthData_Keyboard *kbAuth, void *ctx)
-{
-    (void) ctx;
-    WMEMCPY(kbAuth, &promptData, sizeof(WS_UserAuthData_Keyboard));
-
-    return WS_SUCCESS;
-}
-
 static INLINE void SignalTcpReady(tcp_ready* ready, word16 port)
 {
     pthread_mutex_lock(&ready->mutex);
@@ -332,13 +336,13 @@ static THREAD_RETURN WOLFSSH_THREAD server_thread(void* args)
     }
 
     wolfSSH_SetUserAuth(ctx, serverUserAuth);
-    wolfSSH_SetKeyboardAuthPrompts(ctx, serverKeyboardCallback);
     ssh = wolfSSH_new(ctx);
     if (ssh == NULL) {
         ES_ERROR("Couldn't allocate SSH data.\n");
     }
     keyLoadBuf = buf;
     bufSz = EXAMPLE_KEYLOAD_BUFFER_SZ;
+    wolfSSH_SetUserAuthCtx(ssh, &promptData);
 
     bufSz = load_key(peerEcc, keyLoadBuf, bufSz);
     if (bufSz == 0) {
diff --git a/wolfssh/internal.h b/wolfssh/internal.h
index 7f341d4d..2cd7b97b 100644
--- a/wolfssh/internal.h
+++ b/wolfssh/internal.h
@@ -919,6 +919,7 @@ struct WOLFSSH {
     void* keyingCompletionCtx;
 #ifdef WOLFSSH_KEYBOARD_INTERACTIVE
     WS_UserAuthData_Keyboard kbAuth;
+    byte kbAuthAttempts;
 #endif
 };