From 039aea5142538bfe31c5eeb37f63ead8c86fc5fc Mon Sep 17 00:00:00 2001
From: John Safranek <john@wolfssl.com>
Date: Wed, 20 Mar 2024 14:08:50 -0700
Subject: [PATCH] Messaging Filtering

1. Add an error code and string for the message filtering fail.
2. Add a function to check incoming message IDs for appropriateness
   during the client or server handshake.
(ZD 17710)
---
 src/internal.c     | 85 ++++++++++++++++++++++++++++++++++++++++++++++
 wolfssh/error.h    |  3 +-
 wolfssh/internal.h | 14 ++++++++
 3 files changed, 101 insertions(+), 1 deletion(-)

diff --git a/src/internal.c b/src/internal.c
index 7e9727cf6..83ebc709d 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -433,6 +433,9 @@ const char* GetErrorString(int err)
         case WS_SFTP_NOT_FILE_E:
             return "not a regular file";
 
+        case WS_MSGID_NOT_ALLOWED_E:
+            return "message not allowed before user authentication";
+
         default:
             return "Unknown error code";
     }
@@ -545,6 +548,84 @@ static void HandshakeInfoFree(HandshakeInfo* hs, void* heap)
 }
 
 
+#ifndef NO_WOLFSSH_SERVER
+INLINE static int IsMessageAllowedServer(WOLFSSH *ssh, byte msg)
+{
+    /* Has client userauth started? */
+    if (ssh->acceptState < ACCEPT_KEYED) {
+        if (msg > MSGID_KEXDH_LIMIT) {
+            return 0;
+        }
+    }
+    /* Is server userauth complete? */
+    if (ssh->acceptState < ACCEPT_SERVER_USERAUTH_SENT) {
+        /* Explicitly check for messages not allowed before user
+         * authentication has comleted. */
+        if (msg >= MSGID_USERAUTH_LIMIT) {
+            WLOG(WS_LOG_DEBUG, "Message ID %u not allowed by server "
+                    "before user authentication is complete", msg);
+            return 0;
+        }
+        /* Explicitly check for the user authentication messages that
+         * only the server sends, it shouldn't receive them. */
+        if (msg > MSGID_USERAUTH_RESTRICT) {
+            WLOG(WS_LOG_DEBUG, "Message ID %u not allowed by server "
+                    "during user authentication", msg);
+            return 0;
+        }
+    }
+    return 1;
+}
+#endif /* NO_WOLFSSH_SERVER */
+
+
+#ifndef NO_WOLFSSH_CLIENT
+INLINE static int IsMessageAllowedClient(WOLFSSH *ssh, byte msg)
+{
+    /* Has client userauth started? */
+    if (ssh->connectState < CONNECT_CLIENT_KEXDH_INIT_SENT) {
+        if (msg >= MSGID_KEXDH_LIMIT) {
+            return 0;
+        }
+    }
+    /* Is client userauth complete? */
+    if (ssh->connectState < CONNECT_SERVER_USERAUTH_ACCEPT_DONE) {
+        /* Explicitly check for messages not allowed before user
+         * authentication has comleted. */
+        if (msg >= MSGID_USERAUTH_LIMIT) {
+            WLOG(WS_LOG_DEBUG, "Message ID %u not allowed by client "
+                    "before user authentication is complete", msg);
+            return 0;
+        }
+        /* Explicitly check for the user authentication message that
+         * only the client sends, it shouldn't receive it. */
+        if (msg == MSGID_USERAUTH_RESTRICT) {
+            WLOG(WS_LOG_DEBUG, "Message ID %u not allowed by client "
+                    "during user authentication", msg);
+            return 0;
+        }
+    }
+    return 1;
+}
+#endif /* NO_WOLFSSH_CLIENT */
+
+
+INLINE static int IsMessageAllowed(WOLFSSH *ssh, byte msg)
+{
+#ifndef NO_WOLFSSH_SERVER
+    if (ssh->ctx->side == WOLFSSH_ENDPOINT_SERVER) {
+        return IsMessageAllowedServer(ssh, msg);
+    }
+#endif /* NO_WOLFSSH_SERVER */
+#ifndef NO_WOLFSSH_CLIENT
+    if (ssh->ctx->side == WOLFSSH_ENDPOINT_CLIENT) {
+        return IsMessageAllowedClient(ssh, msg);
+    }
+#endif /* NO_WOLFSSH_CLIENT */
+    return 0;
+}
+
+
 #ifdef DEBUG_WOLFSSH
 
 static const char cannedBanner[] =
@@ -7526,6 +7607,10 @@ static int DoPacket(WOLFSSH* ssh, byte* bufferConsumed)
         return WS_OVERFLOW_E;
     }
 
+    if (!IsMessageAllowed(ssh, msg)) {
+        return WS_MSGID_NOT_ALLOWED_E;
+    }
+
     switch (msg) {
 
         case MSGID_DISCONNECT:
diff --git a/wolfssh/error.h b/wolfssh/error.h
index d1fdfdd27..eab353666 100644
--- a/wolfssh/error.h
+++ b/wolfssh/error.h
@@ -132,8 +132,9 @@ enum WS_ErrorCodes {
     WS_KEY_CHECK_VAL_E      = -1091, /* OpenSSH key check value fail */
     WS_KEY_FORMAT_E         = -1092, /* OpenSSH key format fail */
     WS_SFTP_NOT_FILE_E      = -1093, /* Not a regular file */
+    WS_MSGID_NOT_ALLOWED_E  = -1094, /* Message not allowed before userauth */
     
-    WS_LAST_E               = -1093  /* Update this to indicate last error */
+    WS_LAST_E               = -1094  /* Update this to indicate last error */
 };
 
 
diff --git a/wolfssh/internal.h b/wolfssh/internal.h
index 3cd5b3776..c32c40780 100644
--- a/wolfssh/internal.h
+++ b/wolfssh/internal.h
@@ -1087,6 +1087,20 @@ enum WS_MessageIds {
 };
 
 
+#define MSGID_KEXDH_LIMIT 30
+
+/* The endpoints should not allow message IDs greater than or
+ * equal to msgid 80 before user authentication is complete.
+ * Per RFC 4252 section 6. */
+#define MSGID_USERAUTH_LIMIT 80
+
+/* The client should only send the user auth request message
+ * (50), it should not accept it. The server should only receive
+ * the user auth request message, it should not accept the other
+ * user auth messages, it sends them. (>50) */
+#define MSGID_USERAUTH_RESTRICT 50
+
+
 #define CHANNEL_EXTENDED_DATA_STDERR WOLFSSH_EXT_DATA_STDERR