rekey: trigger highwater on per-key packet count

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,26 @@ static INLINE int HighwaterCheck(WOLFSSH* ssh, byte side)
              (side == WOLFSSH_HWSIDE_TRANSMIT) ? "Transmit" : "Receive");
 
         ssh->highwaterFlag = 1;
+        fire = 1;
+    }
+
+    /* RFC 4344 Sec 3.1: bound packets per key (txMsgCount/rxMsgCount reset on
+     * rekey) to limit cipher/IV exhaustion under a single key; this is not a
+     * guard for the absolute SSH sequence number, which never resets. */
+    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 +1047,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 +1250,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 +6308,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;
@@ -10151,6 +10177,8 @@ static int DoPacket(WOLFSSH* ssh, byte* bufferConsumed)
 
     /* if the auth is still pending, don't discard the packet data */
     if (ret != WS_AUTH_PENDING) {
+        int hwRet;
+
         if (payloadSz > 0) {
             idx += payloadIdx;
             if (idx + padSz > len) {
@@ -10162,7 +10190,14 @@ static int DoPacket(WOLFSSH* ssh, byte* bufferConsumed)
         idx += padSz;
         ssh->inputBuffer.idx = idx;
         ssh->peerSeq++;
+        ssh->rxMsgCount++;
         *bufferConsumed = 1;
+
+        /* Run after rxMsgCount++ so the msg-count highwater fires on the
+         * threshold-crossing packet and reports the correct side. */
+        hwRet = HighwaterCheck(ssh, WOLFSSH_HWSIDE_RECEIVE);
+        if (hwRet != WS_SUCCESS)
+            ret = hwRet;
     }
 
     return ret;
@@ -10991,6 +11026,7 @@ static int BundlePacket(WOLFSSH* ssh)
 
     if (ret == WS_SUCCESS) {
         ssh->seq++;
+        ssh->txMsgCount++;
         ssh->outputBuffer.length = idx;
     }
     else {
@@ -13301,6 +13337,7 @@ int SendNewKeys(WOLFSSH* ssh)
 
     if (ret == WS_SUCCESS) {
         ssh->txCount = 0;
+        ssh->txMsgCount = 0;
     }
 
     if (ret == WS_SUCCESS) {

src/ssh.c

@@ -267,6 +267,45 @@ word32 wolfSSH_GetHighwater(WOLFSSH* ssh)
 }
 
 
+int wolfSSH_CTX_SetMsgHighwater(WOLFSSH_CTX* ctx, word32 highwater)
+{
+    WLOG(WS_LOG_DEBUG, "Entering wolfSSH_CTX_SetMsgHighwater()");
+
+    if (ctx) {
+        ctx->msgHighwaterMark = highwater;
+
+        return WS_SUCCESS;
+    }
+
+    return WS_BAD_ARGUMENT;
+}
+
+
+int wolfSSH_SetMsgHighwater(WOLFSSH* ssh, word32 highwater)
+{
+    WLOG(WS_LOG_DEBUG, "Entering wolfSSH_SetMsgHighwater()");
+
+    if (ssh) {
+        ssh->msgHighwaterMark = highwater;
+
+        return WS_SUCCESS;
+    }
+
+    return WS_BAD_ARGUMENT;
+}
+
+
+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 highwater,
                             WS_CallbackHighwater cb)
 {

wolfssh/internal.h

@@ -443,6 +443,12 @@ enum NameIdType {
 #ifndef DEFAULT_HIGHWATER_MARK
     #define DEFAULT_HIGHWATER_MARK ((1024 * 1024 * 1024) - (32 * 1024))
 #endif
+#ifndef WOLFSSH_DEFAULT_MSG_HIGHWATER_MARK
+    /* RFC 4344 Sec 3.1: bound packets per key epoch to limit cipher/IV
+     * exhaustion under a single key (not the absolute SSH sequence number,
+     * which is not reset by rekey); default to 2^31 packets per key. */
+    #define WOLFSSH_DEFAULT_MSG_HIGHWATER_MARK 0x80000000U
+#endif
 #ifndef DEFAULT_WINDOW_SZ
     #define DEFAULT_WINDOW_SZ (128 * 1024)
 #endif
@@ -597,6 +603,7 @@ struct WOLFSSH_CTX {
     byte publicKeyAlgo[WOLFSSH_MAX_PUB_KEY_ALGO];
     word32 publicKeyAlgoCount;
     word32 highwaterMark;
+    word32 msgHighwaterMark;
     const char* banner;
     const char* sshProtoIdStr;
     const char* algoListKex;
@@ -749,8 +756,12 @@ struct WOLFSSH {
     int wflags;            /* optional write flags */
     word32 txCount;
     word32 rxCount;
+    word32 txMsgCount;     /* Packets sent under current keys */
+    word32 rxMsgCount;     /* Packets received under current keys */
     word32 highwaterMark;
+    word32 msgHighwaterMark; /* Per-key packet limit (RFC 4344 Sec 3.1) */
     byte highwaterFlag;    /* Set when highwater CB called */
+    byte msgHighwaterFlag; /* Set when msg-count highwater CB called */
     void* highwaterCtx;    /* Highwater CB context */
     void* globalReqCtx;    /* Global Request CB context */
     void* reqSuccessCtx;   /* Global Request Success CB context */

wolfssh/ssh.h

@@ -78,14 +78,19 @@ WOLFSSH_API int wolfSSH_SetFilesystemHandle(WOLFSSH*, void*);
 WOLFSSH_API void* wolfSSH_GetFilesystemHandle(WOLFSSH*);
 
 /* data high water mark functions */
-WOLFSSH_API int wolfSSH_SetHighwater(WOLFSSH*, word32);
-WOLFSSH_API word32 wolfSSH_GetHighwater(WOLFSSH*);
-
-typedef int (*WS_CallbackHighwater)(byte, void*);
-WOLFSSH_API void wolfSSH_SetHighwaterCb(WOLFSSH_CTX*, word32,
-                                        WS_CallbackHighwater);
-WOLFSSH_API void wolfSSH_SetHighwaterCtx(WOLFSSH*, void*);
-WOLFSSH_API void* wolfSSH_GetHighwaterCtx(WOLFSSH*);
+WOLFSSH_API int wolfSSH_SetHighwater(WOLFSSH* ssh, word32 highwater);
+WOLFSSH_API word32 wolfSSH_GetHighwater(WOLFSSH* ssh);
+
+/* packet count high water mark functions (RFC 4344 Sec 3.1) */
+WOLFSSH_API int wolfSSH_CTX_SetMsgHighwater(WOLFSSH_CTX* ctx, word32 highwater);
+WOLFSSH_API int wolfSSH_SetMsgHighwater(WOLFSSH* ssh, word32 highwater);
+WOLFSSH_API word32 wolfSSH_GetMsgHighwater(WOLFSSH* ssh);
+
+typedef int (*WS_CallbackHighwater)(byte side, void* ctx);
+WOLFSSH_API void wolfSSH_SetHighwaterCb(WOLFSSH_CTX* ctx,
+        word32 highwater, WS_CallbackHighwater cb);
+WOLFSSH_API void wolfSSH_SetHighwaterCtx(WOLFSSH* ssh, void* ctx);
+WOLFSSH_API void* wolfSSH_GetHighwaterCtx(WOLFSSH* ssh);
 
 WOLFSSH_API int wolfSSH_ReadKey_buffer_ex(const byte* in, word32 inSz, int format,
         byte** out, word32* outSz, const byte** outType, word32* outTypeSz,