rekey: trigger highwater on per-key packet count

- Track txMsgCount/rxMsgCount per key epoch and reset on NEW_KEYS;
  seq/peerSeq still wrap freely per RFC 4253 Sec 6.4.
- Extend HighwaterCheck to fire highwaterCb when packet count crosses
  msgHighwaterMark (default 2^31, RFC 4344 Sec 3.1).
- Add wolfSSH_CTX_SetMsgHighwater / SetMsgHighwater / GetMsgHighwater.
- Consolidate receive-path HighwaterCheck into DoPacket so byte- and
  packet-count thresholds share a single canonical fire site.

Issue: F-246

Author: ejohnstown

src/internal.c

@@ -188,6 +188,9 @@
     Set the number of Miller-Rabin rounds used when the client checks the
     server's prime group when using GEX key exchange. The default is 8. More
     rounds are better, but also takes a lot longer.
+  WOLFSSH_DEFAULT_MSG_HIGHWATER_MARK
+    Set the default value for the number of messages to send or receive before
+    calling the highwater callback function. By default this forces a rekey.
 */
 
 static const char sshProtoIdStr[] = "SSH-2.0-wolfSSHv"
@@ -544,7 +547,10 @@ static int HashUpdate(wc_HashAlg* hash, enum wc_HashType type,
 static INLINE int HighwaterCheck(WOLFSSH* ssh, byte side)
 {
     int ret = WS_SUCCESS;
+    int fire = 0;
 
+    /* RFC 4253 Sec 9: bound bytes per key (txCount/rxCount reset on rekey)
+     * to limit cipher keystream/IV exhaustion under a single key. */
     if (!ssh->highwaterFlag && ssh->highwaterMark &&
         (ssh->txCount >= ssh->highwaterMark ||
          ssh->rxCount >= ssh->highwaterMark)) {
@@ -553,10 +559,28 @@ static INLINE int HighwaterCheck(WOLFSSH* ssh, byte side)
              (side == WOLFSSH_HWSIDE_TRANSMIT) ? "Transmit" : "Receive");
 
         ssh->highwaterFlag = 1;
+        fire = 1;
+    }
+
+    /* RFC 4344 Sec 3.1: rekey before the 32-bit SSH sequence number wraps
+     * to prevent MAC/nonce reuse. Counter is per-key (resets on rekey),
+     * not the absolute ssh->seq (which does not reset); default 2^31
+     * keeps each epoch comfortably under the 2^32 wrap. */
+    if (!ssh->msgHighwaterFlag && ssh->msgHighwaterMark &&
+        (ssh->txMsgCount >= ssh->msgHighwaterMark ||
+         ssh->rxMsgCount >= ssh->msgHighwaterMark)) {
+
+        WLOG(WS_LOG_DEBUG, "%s over msg high water mark",
+             (side == WOLFSSH_HWSIDE_TRANSMIT) ? "Transmit" : "Receive");
 
-        if (ssh->ctx->highwaterCb)
-            ret = ssh->ctx->highwaterCb(side, ssh->highwaterCtx);
+        ssh->msgHighwaterFlag = 1;
+        fire = 1;
     }
+
+
+    if (fire && ssh->ctx->highwaterCb)
+        ret = ssh->ctx->highwaterCb(side, ssh->highwaterCtx);
+
     return ret;
 }
 
@@ -1025,6 +1049,7 @@ WOLFSSH_CTX* CtxInit(WOLFSSH_CTX* ctx, byte side, void* heap)
     ctx->ioSendCb = wsEmbedSend;
 #endif /* WOLFSSH_USER_IO */
     ctx->highwaterMark = DEFAULT_HIGHWATER_MARK;
+    ctx->msgHighwaterMark = WOLFSSH_DEFAULT_MSG_HIGHWATER_MARK;
     ctx->highwaterCb = wsHighwater;
 #if defined(WOLFSSH_SCP) && !defined(WOLFSSH_SCP_USER_CALLBACKS)
     ctx->scpRecvCb = wsScpRecvCallback;
@@ -1227,6 +1252,7 @@ WOLFSSH* SshInit(WOLFSSH* ssh, WOLFSSH_CTX* ctx)
     ssh->ioReadCtx   = &ssh->rfd;  /* prevent invalid access if not correctly */
     ssh->ioWriteCtx  = &ssh->wfd;  /* set */
     ssh->highwaterMark = ctx->highwaterMark;
+    ssh->msgHighwaterMark = ctx->msgHighwaterMark;
     ssh->highwaterCtx  = (void*)ssh;
     ssh->reqSuccessCtx = (void*)ssh;
     ssh->fs            = NULL;
@@ -6284,7 +6310,9 @@ static int DoNewKeys(WOLFSSH* ssh, byte* buf, word32 len, word32* idx)
 
     if (ret == WS_SUCCESS) {
         ssh->rxCount = 0;
+        ssh->rxMsgCount = 0;
         ssh->highwaterFlag = 0;
+        ssh->msgHighwaterFlag = 0;
 
         /* Clear peer is keying flag */
         ssh->isKeying &= ~WOLFSSH_PEER_IS_KEYING;
@@ -10162,7 +10190,23 @@ static int DoPacket(WOLFSSH* ssh, byte* bufferConsumed)
         idx += padSz;
         ssh->inputBuffer.idx = idx;
         ssh->peerSeq++;
+        ssh->rxMsgCount++;
         *bufferConsumed = 1;
+
+        /* Canonical receive-path highwater check. rxCount was advanced by
+         * Decrypt/DecryptAead in DoReceive and rxMsgCount just above, so
+         * either threshold can fire here on the crossing packet. Run the
+         * check unconditionally so the flag-set side effect and callback
+         * fire on data packets (DoChannelData et al. return informational
+         * status like WS_CHAN_RXD, not WS_SUCCESS). Only fold the result
+         * into ret when ret is currently WS_SUCCESS, so a rekey-trigger
+         * return (e.g. WS_WANT_WRITE from SendKexInit) does not mask a
+         * real packet error or informational status already in ret. */
+        {
+            int hwRet = HighwaterCheck(ssh, WOLFSSH_HWSIDE_RECEIVE);
+            if (ret == WS_SUCCESS)
+                ret = hwRet;
+        }
     }
 
     return ret;
@@ -10602,14 +10646,6 @@ int DoReceive(WOLFSSH* ssh)
                     ssh->error = ret;
                     return WS_FATAL_ERROR;
                 }
-
-                ret = HighwaterCheck(ssh, WOLFSSH_HWSIDE_RECEIVE);
-                if (ret != WS_SUCCESS) {
-                    WLOG(WS_LOG_DEBUG, "PR: First HighwaterCheck fail");
-                    ssh->error = ret;
-                    ret = WS_FATAL_ERROR;
-                    break;
-                }
             }
             FALL_THROUGH;
 
@@ -10694,14 +10730,6 @@ int DoReceive(WOLFSSH* ssh)
                 }
             }
             ssh->processReplyState = PROCESS_PACKET;
-
-            ret = HighwaterCheck(ssh, WOLFSSH_HWSIDE_RECEIVE);
-            if (ret != WS_SUCCESS) {
-                WLOG(WS_LOG_DEBUG, "PR: HighwaterCheck fail");
-                ssh->error = ret;
-                ret = WS_FATAL_ERROR;
-                break;
-            }
             FALL_THROUGH;
 
         case PROCESS_PACKET:
@@ -10991,6 +11019,7 @@ static int BundlePacket(WOLFSSH* ssh)
 
     if (ret == WS_SUCCESS) {
         ssh->seq++;
+        ssh->txMsgCount++;
         ssh->outputBuffer.length = idx;
     }
     else {
@@ -13301,6 +13330,7 @@ int SendNewKeys(WOLFSSH* ssh)
 
     if (ret == WS_SUCCESS) {
         ssh->txCount = 0;
+        ssh->txMsgCount = 0;
     }
 
     if (ret == WS_SUCCESS) {
@@ -17976,6 +18006,11 @@ int wolfSSH_TestChannelPutData(WOLFSSH_CHANNEL* channel, byte* data,
     return ChannelPutData(channel, data, dataSz);
 }
 
+int wolfSSH_TestHighwaterCheck(WOLFSSH* ssh, byte side)
+{
+    return HighwaterCheck(ssh, side);
+}
+
 #ifndef WOLFSSH_NO_DH_GEX_SHA256
 
 int wolfSSH_TestDoKexDhGexRequest(WOLFSSH* ssh, byte* buf, word32 len,

src/ssh.c

@@ -242,12 +242,12 @@ void* wolfSSH_GetFilesystemHandle(WOLFSSH* ssh)
 }
 
 
-int wolfSSH_SetHighwater(WOLFSSH* ssh, word32 highwater)
+int wolfSSH_SetHighwater(WOLFSSH* ssh, word32 level)
 {
     WLOG(WS_LOG_DEBUG, "Entering wolfSSH_SetHighwater()");
 
     if (ssh) {
-        ssh->highwaterMark = highwater;
+        ssh->highwaterMark = level;
 
         return WS_SUCCESS;
     }
@@ -267,13 +267,42 @@ word32 wolfSSH_GetHighwater(WOLFSSH* ssh)
 }
 
 
-void wolfSSH_SetHighwaterCb(WOLFSSH_CTX* ctx, word32 highwater,
-                            WS_CallbackHighwater cb)
+void wolfSSH_CTX_SetMsgHighwater(WOLFSSH_CTX* ctx, word32 level)
+{
+    WLOG(WS_LOG_DEBUG, "Entering wolfSSH_CTX_SetMsgHighwater()");
+
+    if (ctx)
+        ctx->msgHighwaterMark = level;
+}
+
+
+void wolfSSH_SetMsgHighwater(WOLFSSH* ssh, word32 level)
+{
+    WLOG(WS_LOG_DEBUG, "Entering wolfSSH_SetMsgHighwater()");
+
+    if (ssh)
+        ssh->msgHighwaterMark = level;
+}
+
+
+word32 wolfSSH_GetMsgHighwater(WOLFSSH* ssh)
+{
+    WLOG(WS_LOG_DEBUG, "Entering wolfSSH_GetMsgHighwater()");
+
+    if (ssh)
+        return ssh->msgHighwaterMark;
+
+    return 0;
+}
+
+
+void wolfSSH_SetHighwaterCb(WOLFSSH_CTX* ctx, word32 level,
+        WS_CallbackHighwater cb)
 {
     WLOG(WS_LOG_DEBUG, "Entering wolfSSH_SetHighwaterCb()");
 
     if (ctx) {
-        ctx->highwaterMark = highwater;
+        ctx->highwaterMark = level;
         ctx->highwaterCb = cb;
     }
 }

tests/unit.c

@@ -1181,6 +1181,139 @@ static int test_ChannelPutData(void)
     return result;
 }
 
+/* Counter callback for test_MsgHighwater. Records each invocation without
+ * triggering wolfSSH_TriggerKeyExchange (which needs a live session). */
+typedef struct HwTestCtx {
+    int  count;
+    byte lastSide;
+} HwTestCtx;
+
+static int HwTestCb(byte side, void* ctx)
+{
+    HwTestCtx* hc = (HwTestCtx*)ctx;
+    if (hc != NULL) {
+        hc->count++;
+        hc->lastSide = side;
+    }
+    return WS_SUCCESS;
+}
+
+/* Exercise the wolfSSH_*MsgHighwater APIs and the per-key packet-count
+ * threshold path inside HighwaterCheck. Covers:
+ *   - NULL safety on getters/setters
+ *   - CTX default value matches WOLFSSH_DEFAULT_MSG_HIGHWATER_MARK
+ *   - CTX/SSH setter round-trip and CTX -> SSH inheritance on wolfSSH_new
+ *   - SSH setter does not bleed back into the CTX
+ *   - Threshold crossing fires the highwater callback exactly once per epoch
+ *     (msgHighwaterFlag gates re-firing under the same keys)
+ *   - Receive side fires independently of the transmit side
+ *   - msgHighwaterMark == 0 disables the per-key packet check */
+static int test_MsgHighwater(void)
+{
+    WOLFSSH_CTX* ctx = NULL;
+    WOLFSSH*     ssh = NULL;
+    HwTestCtx    hc;
+    int          result = 0;
+
+    if (wolfSSH_GetMsgHighwater(NULL) != 0)
+        return -800;
+    wolfSSH_CTX_SetMsgHighwater(NULL, 1234);
+    wolfSSH_SetMsgHighwater(NULL, 1234);
+
+    ctx = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_SERVER, NULL);
+    if (ctx == NULL)
+        return -801;
+
+    if (ctx->msgHighwaterMark != WOLFSSH_DEFAULT_MSG_HIGHWATER_MARK) {
+        result = -802;
+        goto done;
+    }
+
+    wolfSSH_CTX_SetMsgHighwater(ctx, 4096);
+    if (ctx->msgHighwaterMark != 4096) {
+        result = -803;
+        goto done;
+    }
+
+    ssh = wolfSSH_new(ctx);
+    if (ssh == NULL) {
+        result = -804;
+        goto done;
+    }
+    if (wolfSSH_GetMsgHighwater(ssh) != 4096) {
+        result = -805;
+        goto done;
+    }
+
+    wolfSSH_SetMsgHighwater(ssh, 16);
+    if (wolfSSH_GetMsgHighwater(ssh) != 16) {
+        result = -806;
+        goto done;
+    }
+    if (ctx->msgHighwaterMark != 4096) {
+        result = -807;
+        goto done;
+    }
+
+    /* Install a counter callback. ssh->highwaterMark stays at the inherited
+     * default (~1 GiB) and txCount/rxCount are not touched, so the byte-count
+     * branch cannot fire and only the packet-count branch is under test. */
+    WMEMSET(&hc, 0, sizeof(hc));
+    wolfSSH_SetHighwaterCb(ctx, ctx->highwaterMark, HwTestCb);
+    wolfSSH_SetHighwaterCtx(ssh, &hc);
+    wolfSSH_SetMsgHighwater(ssh, 8);
+
+    ssh->txMsgCount = 7;
+    if (wolfSSH_TestHighwaterCheck(ssh, WOLFSSH_HWSIDE_TRANSMIT) != WS_SUCCESS
+            || hc.count != 0) {
+        result = -808;
+        goto done;
+    }
+
+    ssh->txMsgCount = 8;
+    if (wolfSSH_TestHighwaterCheck(ssh, WOLFSSH_HWSIDE_TRANSMIT) != WS_SUCCESS
+            || hc.count != 1
+            || hc.lastSide != WOLFSSH_HWSIDE_TRANSMIT) {
+        result = -809;
+        goto done;
+    }
+
+    /* Flag-gated: further packets in the same epoch must not re-fire. */
+    ssh->txMsgCount = 100;
+    if (wolfSSH_TestHighwaterCheck(ssh, WOLFSSH_HWSIDE_TRANSMIT) != WS_SUCCESS
+            || hc.count != 1) {
+        result = -810;
+        goto done;
+    }
+
+    /* Simulate a fresh key epoch (msgHighwaterFlag and rx/txMsgCount are
+     * reset by DoNewKeys/SendNewKeys) and verify the receive side fires. */
+    ssh->msgHighwaterFlag = 0;
+    ssh->rxMsgCount = 8;
+    if (wolfSSH_TestHighwaterCheck(ssh, WOLFSSH_HWSIDE_RECEIVE) != WS_SUCCESS
+            || hc.count != 2
+            || hc.lastSide != WOLFSSH_HWSIDE_RECEIVE) {
+        result = -811;
+        goto done;
+    }
+
+    /* mark == 0 disables the per-key packet check entirely. */
+    wolfSSH_SetMsgHighwater(ssh, 0);
+    ssh->msgHighwaterFlag = 0;
+    ssh->txMsgCount = 0xFFFFFFFFu;
+    ssh->rxMsgCount = 0xFFFFFFFFu;
+    if (wolfSSH_TestHighwaterCheck(ssh, WOLFSSH_HWSIDE_TRANSMIT) != WS_SUCCESS
+            || hc.count != 2) {
+        result = -812;
+        goto done;
+    }
+
+done:
+    wolfSSH_free(ssh);
+    wolfSSH_CTX_free(ctx);
+    return result;
+}
+
 /* Verify DoChannelRequest sends CHANNEL_SUCCESS for known types and
  * CHANNEL_FAILURE for unrecognized ones (RFC 4254 Section 5.4).
  *
@@ -1609,6 +1742,10 @@ int wolfSSH_UnitTest(int argc, char** argv)
     unitResult = test_ChannelPutData();
     printf("ChannelPutData: %s\n", (unitResult == 0 ? "SUCCESS" : "FAILED"));
     testResult = testResult || unitResult;
+
+    unitResult = test_MsgHighwater();
+    printf("MsgHighwater: %s\n", (unitResult == 0 ? "SUCCESS" : "FAILED"));
+    testResult = testResult || unitResult;
 #endif
 
 #ifdef WOLFSSH_KEYGEN