internal: set the keying flag once the KEX init reaches the transport

- SendKexInit() takes ssh->txFlushCount before wolfSSH_SendPacket()
  and sets WOLFSSH_SELF_IS_KEYING when SendPacketDelivered() reports
  the packet away, in place of setting it before the packet is built.
- PurgePacket() runs on the same decision, so a packet the transport
  took is not purged behind an error the highwater callback raised.
- internal.h describes the flag as set once the KEX init is sent or
  queued.
- tests/unit.c covers a KEX init whose send fails outright, one that
  short-writes, and one the transport takes whole behind a highwater
  callback that fails.
- FailHighwater() moves beside the other shared send callbacks and
  gains WS_MAYBE_UNUSED, so the client test uses it too.

src/internal.c

@@ -15194,6 +15194,7 @@ int SendKexInit(WOLFSSH* ssh)
             macAlgoNamesSz = 0, noneNamesSz = 0;
 
     int ret = WS_SUCCESS;
+    int delivered = 0;
 
     WLOG(WS_LOG_DEBUG, "Entering SendKexInit()");
 
@@ -15220,8 +15221,6 @@ int SendKexInit(WOLFSSH* ssh)
     }
 
     if (ret == WS_SUCCESS) {
-        /* Set self is keying flag since we started sending the KEX init msg */
-        ssh->isKeying |= WOLFSSH_SELF_IS_KEYING;
         if (ssh->handshake == NULL) {
             ssh->handshake = HandshakeInfoNew(ssh->ctx->heap);
             if (ssh->handshake == NULL) {
@@ -15349,11 +15348,19 @@ int SendKexInit(WOLFSSH* ssh)
     }
 
     if (ret == WS_SUCCESS) {
+        word32 flushes = ssh->txFlushCount;
+
         ret = wolfSSH_SendPacket(ssh);
+        delivered = SendPacketDelivered(ssh, flushes, ret);
     }
 
-    if (ret != WS_WANT_WRITE && ret != WS_SUCCESS)
+    if (delivered) {
+        /* Set self is keying flag now the KEX init msg is away */
+        ssh->isKeying |= WOLFSSH_SELF_IS_KEYING;
+    }
+    else {
         PurgePacket(ssh);
+    }
 
     WLOG(WS_LOG_DEBUG, "Leaving SendKexInit(), ret = %d", ret);
     return ret;

tests/unit.c

@@ -4565,6 +4565,13 @@ static WS_MAYBE_UNUSED int OobIoSend(WOLFSSH* ssh, void* buf, word32 sz,
     return (int)sz + 1;
 }
 
+/* Fires once the message highwater mark is crossed and reports an error. */
+static WS_MAYBE_UNUSED int FailHighwater(byte side, void* ctx)
+{
+    (void)side; (void)ctx;
+    return WS_FATAL_ERROR;
+}
+
 static int test_DoChannelExtendedData_overflow(void)
 {
     WOLFSSH_CTX*     ctx = NULL;
@@ -5705,13 +5712,6 @@ done:
 
 #ifndef NO_WOLFSSH_SERVER
 
-/* Fires once the message highwater mark is crossed and reports an error. */
-static int FailHighwater(byte side, void* ctx)
-{
-    (void)side; (void)ctx;
-    return WS_FATAL_ERROR;
-}
-
 /* wolfSSH_SendPacket() runs the highwater check after the packet is on the wire
  * and returns the highwater callback's status, so a failing callback makes a
  * delivered WINDOW_ADJUST look like a failed send. Credit re-parked then is
@@ -8758,6 +8758,81 @@ done:
     wolfSSH_CTX_free(ctx);
     return result;
 }
+
+
+/* Covers a KEX init whose send fails outright and one that short-writes. */
+static int test_KexInitSendAwayGatesKeying(void)
+{
+    WOLFSSH_CTX* ctx = NULL;
+    WOLFSSH*     ssh = NULL;
+    int          result = 0;
+    int          ret;
+
+    ctx = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_CLIENT, NULL);
+    if (ctx == NULL)
+        return -1897;
+    /* No refusals, so the first write resets the socket. */
+    s_sendRefusals = 0;
+    wolfSSH_SetIOSend(ctx, RefuseThenResetIoSend);
+    wolfSSH_SetIORecv(ctx, PacketIoRecv);
+
+    ssh = wolfSSH_new(ctx);
+    if (ssh == NULL) { result = -1898; goto done; }
+
+    ret = wolfSSH_TriggerKeyExchange(ssh);
+    if (ret == WS_SUCCESS || ret == WS_WANT_WRITE) {
+        result = -1899;
+        goto done;
+    }
+    if (wolfSSH_RekeyPending(ssh)) { result = -1900; goto done; }
+
+    wolfSSH_free(ssh);
+
+    /* One refusal short-writes instead. */
+    s_sendRefusals = 1;
+    ssh = wolfSSH_new(ctx);
+    if (ssh == NULL) { result = -1903; goto done; }
+
+    ret = wolfSSH_TriggerKeyExchange(ssh);
+    if (ret != WS_SUCCESS && ret != WS_WANT_WRITE) {
+        result = -1904;
+        goto done;
+    }
+    if (!wolfSSH_RekeyPending(ssh)) { result = -1905; goto done; }
+
+    wolfSSH_free(ssh);
+    ssh = NULL;
+    wolfSSH_CTX_free(ctx);
+    ctx = NULL;
+
+    /* The transport takes the whole packet and the highwater callback then
+     * fails, so the error arrives with the KEX init already sent. */
+    ctx = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_CLIENT, NULL);
+    if (ctx == NULL) { result = -1906; goto done; }
+    wolfSSH_SetIOSend(ctx, DiscardIoSend);
+    wolfSSH_SetIORecv(ctx, PacketIoRecv);
+    wolfSSH_SetHighwaterCb(ctx, 1, FailHighwater);
+
+    ssh = wolfSSH_new(ctx);
+    if (ssh == NULL) { result = -1907; goto done; }
+    if (wolfSSH_SetHighwater(ssh, 1) != WS_SUCCESS) {
+        result = -1908;
+        goto done;
+    }
+
+    ret = wolfSSH_TriggerKeyExchange(ssh);
+    if (ret == WS_SUCCESS) { result = -1909; goto done; }
+    if (!wolfSSH_RekeyPending(ssh)) { result = -1910; goto done; }
+
+done:
+    s_sendRefusals = 0;
+    s_recvPkt = NULL;
+    s_recvPktSz = 0;
+    s_recvPktOff = 0;
+    wolfSSH_free(ssh);
+    wolfSSH_CTX_free(ctx);
+    return result;
+}
 #endif /* NO_WOLFSSH_CLIENT */
 
 
@@ -22798,6 +22873,11 @@ int wolfSSH_UnitTest(int argc, char** argv)
     printf("TriggerKeyExchangeKeepsError: %s\n",
            (unitResult == 0 ? "SUCCESS" : "FAILED"));
     testResult = testResult || unitResult;
+
+    unitResult = test_KexInitSendAwayGatesKeying();
+    printf("KexInitSendAwayGatesKeying: %s\n",
+           (unitResult == 0 ? "SUCCESS" : "FAILED"));
+    testResult = testResult || unitResult;
 #endif
 
 

wolfssh/internal.h

@@ -736,7 +736,7 @@ enum NameIdType {
 #define WOLFSSH_PROTOID_LIMIT 255
 
 /* Keep track of keying state for both sides of the connection.
- * WOLFSSH_SELF_IS_KEYING gets set on sending KEX init and
+ * WOLFSSH_SELF_IS_KEYING gets set once the KEX init is sent or queued and
  * WOLFSSH_PEER_IS_KEYING gets set on receiving KEX init */
 #define WOLFSSH_PEER_IS_KEYING 0x01
 #define WOLFSSH_SELF_IS_KEYING 0x02