WIP: client message filter, handing KEX

Author: ejohnstown

src/internal.c

@@ -647,38 +647,64 @@ INLINE static int IsMessageAllowedServer(WOLFSSH *ssh, byte msg)
 #ifndef NO_WOLFSSH_CLIENT
 INLINE static int IsMessageAllowedClient(WOLFSSH *ssh, byte msg)
 {
-    /* Transport Layer Generic messages are always allowed. */
-    if (MSGIDLIMIT_TRANS_GEN(msg)) {
-        return 1;
-    }
-
-    /* The client should only send the user auth request message,
-     * it should not accept it. */
-    if (msg == MSGID_USERAUTH_REQUEST) {
+    /* Only the client should send these messages, never receive. */
+    if (msg == MSGID_SERVICE_REQUEST || msg == MSGID_USERAUTH_REQUEST) {
         WLOG(WS_LOG_DEBUG, "Message ID %u not allowed by %s %s",
                 msg, "client", "ever");
         return 0;
     }
 
+    /* Transport Layer Generic messages are always allowed. */
+    if (MSGIDLIMIT_TRANS_GEN(msg)) {
+        return 1;
+    }
+
     /* Is KEX complete? */
-    if (ssh->connectState < CONNECT_KEYED && ssh->handshake) {
-        /* If expecting a specific message, and didn't receive it, error. */
-        if (ssh->handshake->expectMsgId != MSGID_NONE) {
-            if (msg != ssh->handshake->expectMsgId) {
-                WLOG(WS_LOG_DEBUG, "Message ID %u not the expected message %u",
-                        msg, ssh->handshake->expectMsgId);
+    if (MSGIDLIMIT_TRANS(msg)) {
+        if (ssh->isKeying) {
+            /* MSGID_KEXINIT not allowed when keying. */
+            if (msg == MSGID_KEXINIT) {
+                WLOG(WS_LOG_DEBUG, "Message ID %u not allowed by %s %s",
+                        msg, "client", "when keying");
                 return 0;
             }
-            ssh->handshake->expectMsgId = MSGID_NONE;
-            return 1;
+
+            /* Error if expecting a specific message and didn't receive. */
+            if (ssh->handshake && ssh->handshake->expectMsgId != MSGID_NONE) {
+                if (msg != ssh->handshake->expectMsgId) {
+                    WLOG(WS_LOG_DEBUG,
+                            "Message ID %u not the expected message %u",
+                            msg, ssh->handshake->expectMsgId);
+                    return 0;
+                }
+                else {
+                    /* Got the expected message, clear expectation. */
+                    ssh->handshake->expectMsgId = MSGID_NONE;
+                    return 1;
+                }
+            }
+        }
+        else {
+            /* MSGID_KEXINIT only allowed when not keying. */
+            if (msg == MSGID_KEXINIT) {
+                return 1;
+            }
+
+            /* All other transport KEX and ALGO messages are not allowed
+             * when not keying. */
+            WLOG(WS_LOG_DEBUG, "Message ID %u not allowed by %s %s",
+                    msg, "client", "when not keying");
+            return 0;
         }
     }
+
     /* Has client userauth started? */
     if (ssh->connectState < CONNECT_CLIENT_KEXDH_INIT_SENT) {
         if (msg >= MSGID_KEXDH_GEX_REQUEST) {
             return 0;
         }
     }
+
     /* Is client userauth complete? */
     if (ssh->connectState < CONNECT_SERVER_USERAUTH_ACCEPT_DONE) {
         /* The endpoints should not allow message IDs greater than or
@@ -689,13 +715,6 @@ INLINE static int IsMessageAllowedClient(WOLFSSH *ssh, byte msg)
                     msg, "client", "before user authentication is complete");
             return 0;
         }
-        /* Explicitly check for the user authentication request message.
-         * The client only sends the message, it shouldn't receive it. */
-        if (msg == MSGID_USERAUTH_REQUEST) {
-            WLOG(WS_LOG_DEBUG, "Message ID %u not allowed by %s %s",
-                    msg, "client", "during user authentication");
-            return 0;
-        }
     }
     else {
         if (MSGIDLIMIT_AUTH(msg)) {
@@ -704,6 +723,7 @@ INLINE static int IsMessageAllowedClient(WOLFSSH *ssh, byte msg)
             return 0;
         }
     }
+
     return 1;
 }
 #endif /* NO_WOLFSSH_CLIENT */
@@ -1095,7 +1115,7 @@ WOLFSSH* SshInit(WOLFSSH* ssh, WOLFSSH_CTX* ctx)
     ssh->fs            = NULL;
     ssh->acceptState = ACCEPT_BEGIN;
     ssh->clientState = CLIENT_BEGIN;
-    ssh->isKeying    = 1;
+    ssh->isKeying    = 0;
     ssh->authId      = ID_USERAUTH_PUBLICKEY;
     ssh->supportedAuth[0] = ID_USERAUTH_PUBLICKEY;
     ssh->supportedAuth[1] = ID_USERAUTH_PASSWORD;
@@ -4326,10 +4346,14 @@ static int DoKexInit(WOLFSSH* ssh, byte* buf, word32 len, word32* idx)
         byte scratchLen[LENGTH_SZ];
         word32 strSz = 0;
 
-        if (!ssh->isKeying) {
+        ssh->kexinitRxd = 1;
+        if (!ssh->kexinitTxd) {
             WLOG(WS_LOG_DEBUG, "Keying initiated");
             ret = SendKexInit(ssh);
         }
+        else {
+            ssh->handshake->expectMsgId = MSGID_KEXINIT;
+        }
 
         /* account for possible want write case from SendKexInit */
         if (ret == WS_SUCCESS || ret == WS_WANT_WRITE)
@@ -5941,6 +5965,8 @@ static int DoNewKeys(WOLFSSH* ssh, byte* buf, word32 len, word32* idx)
         ssh->rxCount = 0;
         ssh->highwaterFlag = 0;
         ssh->isKeying = 0;
+        ssh->kexinitTxd = 0;
+        ssh->kexinitRxd = 0;
         HandshakeInfoFree(ssh->handshake, ssh->ctx->heap);
         ssh->handshake = NULL;
         WLOG(WS_LOG_DEBUG, "Keying completed");
@@ -9383,7 +9409,7 @@ static int DoPacket(WOLFSSH* ssh, byte* bufferConsumed)
         case MSGID_KEXINIT:
             WLOG(WS_LOG_DEBUG, "Decoding MSGID_KEXINIT");
             ret = DoKexInit(ssh, buf + idx, payloadSz, &payloadIdx);
-            if (ssh->isKeying == 1 &&
+            if (ssh->kexinitTxd == 1 &&
                     ssh->connectState == CONNECT_SERVER_CHANNEL_REQUEST_DONE) {
                 if (ssh->handshake->kexId == ID_DH_GEX_SHA256) {
 #if !defined(WOLFSSH_NO_DH) && !defined(WOLFSSH_NO_DH_GEX_SHA256)
@@ -10479,7 +10505,6 @@ int SendKexInit(WOLFSSH* ssh)
     }
 
     if (ret == WS_SUCCESS) {
-        ssh->isKeying = 1;
         if (ssh->handshake == NULL) {
             ssh->handshake = HandshakeInfoNew(ssh->ctx->heap);
             if (ssh->handshake == NULL) {
@@ -10602,8 +10627,10 @@ int SendKexInit(WOLFSSH* ssh)
         ret = BundlePacket(ssh);
     }
 
-    if (ret == WS_SUCCESS)
+    if (ret == WS_SUCCESS) {
+        ssh->kexinitTxd = 1;
         ret = wolfSSH_SendPacket(ssh);
+    }
 
     if (ret != WS_WANT_WRITE && ret != WS_SUCCESS)
         PurgePacket(ssh);

wolfssh/internal.h

@@ -739,6 +739,8 @@ struct WOLFSSH {
     byte serverState;
     byte processReplyState;
     byte isKeying;
+    byte kexinitTxd;
+    byte kexinitRxd;
     byte authId;           /* if using public key or password */
     byte supportedAuth[4]; /* supported auth IDs public key , password */