ssh: add a rekey-state accessor

src/internal.c

@@ -15199,6 +15199,7 @@ int SendKexInit(WOLFSSH* ssh)
             macAlgoNamesSz = 0, noneNamesSz = 0;
 
     int ret = WS_SUCCESS;
+    int delivered = 0;
 
     WLOG(WS_LOG_DEBUG, "Entering SendKexInit()");
 
@@ -15225,8 +15226,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) {
@@ -15354,11 +15353,23 @@ int SendKexInit(WOLFSSH* ssh)
     }
 
     if (ret == WS_SUCCESS) {
-        ret = wolfSSH_SendPacket(ssh);
+        word32 flushes = ssh->txFlushCount;
+
+        ret = SendPacketFlush(ssh);
+        delivered = SendPacketDelivered(ssh, flushes, ret);
     }
 
-    if (ret != WS_WANT_WRITE && ret != WS_SUCCESS)
+    if (delivered) {
+        /* Set self is keying flag once the KEX init is sent or queued, before
+         * HighwaterCheck() can run a callback that reads it. */
+        ssh->isKeying |= WOLFSSH_SELF_IS_KEYING;
+    }
+    else {
         PurgePacket(ssh);
+    }
+
+    if (ret == WS_SUCCESS)
+        ret = HighwaterCheck(ssh, WOLFSSH_HWSIDE_TRANSMIT);
 
     WLOG(WS_LOG_DEBUG, "Leaving SendKexInit(), ret = %d", ret);
     return ret;

src/ssh.c

@@ -4709,6 +4709,12 @@ int wolfSSH_OutputPending(const WOLFSSH* ssh)
 }
 
 
+int wolfSSH_RekeyPending(const WOLFSSH* ssh)
+{
+    return (ssh != NULL && ssh->isKeying != 0);
+}
+
+
 #ifdef WOLFSSH_FWD
 
 int wolfSSH_CTX_SetFwdCb(WOLFSSH_CTX* ctx,

tests/regress.c

@@ -15792,6 +15792,39 @@ static void TestClientParseDestination(void)
 }
 
 
+/* Covers each keying bit alone, both together, and a NULL session. */
+static void TestRekeyPendingAccessor(void)
+{
+    WOLFSSH_CTX* ctx;
+    WOLFSSH* ssh;
+
+    AssertIntEQ(wolfSSH_RekeyPending(NULL), 0);
+    AssertIntEQ(wolfSSH_OutputPending(NULL), 0);
+
+    ctx = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_CLIENT, NULL);
+    AssertNotNull(ctx);
+    ssh = wolfSSH_new(ctx);
+    AssertNotNull(ssh);
+
+    AssertIntEQ(wolfSSH_RekeyPending(ssh), 0);
+
+    ssh->isKeying = WOLFSSH_PEER_IS_KEYING;
+    AssertTrue(wolfSSH_RekeyPending(ssh) != 0);
+
+    ssh->isKeying = WOLFSSH_SELF_IS_KEYING;
+    AssertTrue(wolfSSH_RekeyPending(ssh) != 0);
+
+    ssh->isKeying = WOLFSSH_SELF_IS_KEYING | WOLFSSH_PEER_IS_KEYING;
+    AssertTrue(wolfSSH_RekeyPending(ssh) != 0);
+
+    ssh->isKeying = 0;
+    AssertIntEQ(wolfSSH_RekeyPending(ssh), 0);
+
+    wolfSSH_free(ssh);
+    wolfSSH_CTX_free(ctx);
+}
+
+
 #if defined(WOLFSSH_TEST_INTERNAL) || defined(WOLFSSL_BASE64_ENCODE)
 /* Write contents to path exactly as given, with no terminator added, so a
  * test can seed a file whose last line ends without a newline. */
@@ -16203,6 +16236,7 @@ int main(int argc, char** argv)
 #endif
 
     TestClientParseDestination();
+    TestRekeyPendingAccessor();
 #ifdef WOLFSSH_TEST_INTERNAL
     TestAppendKeyToFile();
     TestAppendNoTrailingNewline();

tests/testsuite.c

@@ -241,13 +241,6 @@ int wolfSSH_TestsuiteTest(int argc, char** argv)
 
     wolfSSH_Init();
 
-    /* Linked against the installed library, so this also proves
-     * wolfSSH_OutputPending() is exported and not hidden. */
-    if (wolfSSH_OutputPending(NULL) != 0) {
-        fprintf(stderr, "wolfSSH_OutputPending(NULL) was not zero\n");
-        return EXIT_FAILURE;
-    }
-
     #if defined(FIPS_VERSION_GE) && FIPS_VERSION_GE(5,2)
     {
         int i;

tests/unit.c

@@ -4160,10 +4160,13 @@ static int test_ChannelPutData(void)
     return result;
 }
 
-/* Counter callback for test_MsgHighwater. Records each invocation without
- * triggering wolfSSH_TriggerKeyExchange (which needs a live session). */
+/* Counter callback for the highwater tests. Records each invocation, and the
+ * keying state of ssh when one is set, without triggering
+ * wolfSSH_TriggerKeyExchange (which needs a live session). */
 typedef struct HwTestCtx {
+    WOLFSSH* ssh;
     int  count;
+    int  keying;
     byte lastSide;
 } HwTestCtx;
 
@@ -4173,6 +4176,7 @@ static int HwTestCb(byte side, void* ctx)
     if (hc != NULL) {
         hc->count++;
         hc->lastSide = side;
+        hc->keying = wolfSSH_RekeyPending(hc->ssh);
     }
     return WS_SUCCESS;
 }
@@ -4618,6 +4622,13 @@ static WS_MAYBE_UNUSED int OobIoSend(WOLFSSH* ssh, void* buf, word32 sz,
     return (int)sz + 1;
 }
 
+/* Reports an error whenever the byte or message highwater mark fires. */
+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;
@@ -4907,6 +4918,24 @@ static WS_MAYBE_UNUSED int PacketIoRecv(WOLFSSH* ssh, void* buf, word32 sz, void
     return (int)n;
 }
 
+/* Write budget for the IOSend mocks: that many writes are refused with a
+ * would-block before the mock acts. */
+static int s_sendRefusals = 0;
+
+/* Refuses the first s_sendRefusals writes with a would-block, taking no bytes,
+ * then resets the socket. */
+static WS_MAYBE_UNUSED int RefuseThenResetIoSend(WOLFSSH* ssh, void* buf,
+        word32 sz, void* ctx)
+{
+    (void)ssh; (void)buf; (void)sz; (void)ctx;
+
+    if (s_sendRefusals > 0) {
+        s_sendRefusals--;
+        return WS_CBIO_ERR_WANT_WRITE;
+    }
+    return WS_CBIO_ERR_CONN_RST;
+}
+
 /* Builds a plaintext CHANNEL_EXTENDED_DATA (stderr) SSH packet addressed to
  * channelId, carrying 10 bytes of payload set to fill, into pkt (needs 32
  * bytes) and returns its size. A bare session negotiates no cipher
@@ -5758,13 +5787,6 @@ static int test_ChannelExtDataBufferGrowth(void)
 
 #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
@@ -7927,7 +7949,6 @@ static int test_StreamReadEofOtherChannel(void)
 
 static byte s_sentBuf[512];
 static word32 s_sentSz = 0;
-static int s_sendRefusals = 0;
 
 /* Refuses the first s_sendRefusals writes with a would-block, then takes
  * everything and keeps a copy of what reached the transport. */
@@ -7949,23 +7970,6 @@ static int RefuseThenCaptureIoSend(WOLFSSH* ssh, void* buf, word32 sz,
 }
 
 
-/* Refuses the sends DoChannelClose() makes, then resets the socket under the
- * worker's flush. */
-static int RefuseThenResetIoSend(WOLFSSH* ssh, void* buf, word32 sz, void* ctx)
-{
-    WOLFSSH_UNUSED(ssh);
-    WOLFSSH_UNUSED(buf);
-    WOLFSSH_UNUSED(sz);
-    WOLFSSH_UNUSED(ctx);
-
-    if (s_sendRefusals > 0) {
-        s_sendRefusals--;
-        return WS_CBIO_ERR_WANT_WRITE;
-    }
-    return WS_CBIO_ERR_CONN_RST;
-}
-
-
 /* DoPacket() consumes the peer's CHANNEL_CLOSE whatever DoChannelClose()
  * returns, so the reply gets one chance to be built. A blocked socket must not
  * cost it: the EOF and the close both have to be bundled, and the channel
@@ -8809,6 +8813,94 @@ static int test_TriggerKeyExchangeKeepsError(void)
     wolfSSH_CTX_free(ctx);
     return result;
 }
+
+
+/* Covers a KEX init the transport rejects, one it refuses with a
+ * would-block, and one it takes whole with a highwater callback that fails
+ * or reads the keying flag. */
+static int test_KexInitSendAwayGatesKeying(void)
+{
+    WOLFSSH_CTX* ctx = NULL;
+    WOLFSSH*     ssh = NULL;
+    HwTestCtx    hc;
+    int          result = 0;
+    int          ret;
+
+    ctx = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_CLIENT, NULL);
+    if (ctx == NULL)
+        return -1925;
+    /* No refusals, so the first write resets the socket. */
+    s_sendRefusals = 0;
+    wolfSSH_SetIOSend(ctx, RefuseThenResetIoSend);
+
+    ssh = wolfSSH_new(ctx);
+    if (ssh == NULL) { result = -1926; goto done; }
+
+    ret = wolfSSH_TriggerKeyExchange(ssh);
+    if (ret != WS_SOCKET_ERROR_E) { result = -1927; goto done; }
+    if (wolfSSH_RekeyPending(ssh)) { result = -1928; goto done; }
+    if (wolfSSH_OutputPending(ssh)) { result = -1929; goto done; }
+
+    wolfSSH_free(ssh);
+    ssh = NULL;
+
+    /* One refusal leaves the packet framed and owed instead. */
+    s_sendRefusals = 1;
+    ssh = wolfSSH_new(ctx);
+    if (ssh == NULL) { result = -1930; goto done; }
+
+    ret = wolfSSH_TriggerKeyExchange(ssh);
+    if (ret != WS_WANT_WRITE) { result = -1931; goto done; }
+    if (!wolfSSH_RekeyPending(ssh)) { result = -1932; goto done; }
+    if (!wolfSSH_OutputPending(ssh)) { result = -1933; 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 = -1934; goto done; }
+    wolfSSH_SetIOSend(ctx, DiscardIoSend);
+    wolfSSH_SetHighwaterCb(ctx, 1, FailHighwater);
+
+    ssh = wolfSSH_new(ctx);
+    if (ssh == NULL) { result = -1935; goto done; }
+
+    ret = wolfSSH_TriggerKeyExchange(ssh);
+    if (ret == WS_SUCCESS) { result = -1936; goto done; }
+    if (!wolfSSH_RekeyPending(ssh)) { result = -1937; goto done; }
+
+    wolfSSH_free(ssh);
+    ssh = NULL;
+    wolfSSH_CTX_free(ctx);
+    ctx = NULL;
+
+    /* The highwater callback runs from the KEX init's own flush, with the
+     * packet already on the wire. */
+    ctx = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_CLIENT, NULL);
+    if (ctx == NULL) { result = -1938; goto done; }
+    wolfSSH_SetIOSend(ctx, DiscardIoSend);
+    wolfSSH_SetHighwaterCb(ctx, 1, HwTestCb);
+
+    ssh = wolfSSH_new(ctx);
+    if (ssh == NULL) { result = -1939; goto done; }
+    WMEMSET(&hc, 0, sizeof(hc));
+    hc.ssh = ssh;
+    wolfSSH_SetHighwaterCtx(ssh, &hc);
+
+    ret = wolfSSH_TriggerKeyExchange(ssh);
+    if (ret != WS_SUCCESS) { result = -1940; goto done; }
+    if (hc.count != 1 || !hc.keying) { result = -1941; goto done; }
+
+done:
+    s_sendRefusals = 0;
+    wolfSSH_free(ssh);
+    wolfSSH_CTX_free(ctx);
+    return result;
+}
 #endif /* NO_WOLFSSH_CLIENT */
 
 
@@ -23099,6 +23191,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

@@ -780,7 +780,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

wolfssh/ssh.h

@@ -104,6 +104,8 @@ WOLFSSH_API void wolfSSH_free(WOLFSSH* ssh);
  * the peer's disconnect, which is how most sessions end.
  * To ask whether a write is still owed, call wolfSSH_OutputPending() rather
  * than reading a status: it answers after any return, including a success.
+ * To ask whether a key exchange is in flight, call wolfSSH_RekeyPending()
+ * rather than reading a status: it answers after any return.
  *
  * For WS_CHAN_RXD, WS_EXTDATA, WS_EOF, WS_SUCCESS and a WS_REKEYING that
  * displaced WS_SUCCESS or WS_CHAN_RXD, channelId (when not NULL) names the
@@ -118,6 +120,11 @@ WOLFSSH_API int wolfSSH_GetLastRxId(WOLFSSH* ssh, word32* channelId);
 /* Returns nonzero if a write is still owed. Session state */
 WOLFSSH_API int wolfSSH_OutputPending(const WOLFSSH* ssh);
 
+/* Returns nonzero during a key exchange, the first one included, and 0
+ * otherwise or when ssh is NULL. Only NEWKEYS from both sides clears it, so
+ * it stays set after a failed exchange: end a loop on wolfSSH_worker(). */
+WOLFSSH_API int wolfSSH_RekeyPending(const WOLFSSH* ssh);
+
 WOLFSSH_API int wolfSSH_set_fd(WOLFSSH* ssh, WS_SOCKET_T fd);
 WOLFSSH_API WS_SOCKET_T wolfSSH_get_fd(const WOLFSSH* ssh);