From 6561da9ce29ea3c59b574b3f5bade803e7df2e86 Mon Sep 17 00:00:00 2001
From: John Safranek <john@wolfssl.com>
Date: Mon, 3 Oct 2016 17:22:15 -0700
Subject: [PATCH] Updates to the data highwater marks for a callback.

---
 src/internal.c     |  7 ++++++-
 src/ssh.c          | 42 +++++++++++++++++++++++++++++++++++++-----
 wolfssh/internal.h |  3 +++
 wolfssh/ssh.h      | 19 +++++++++++++++++--
 4 files changed, 63 insertions(+), 8 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index cd2817c5..8bb47cf2 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -162,6 +162,7 @@ WOLFSSH_CTX* CtxInit(WOLFSSH_CTX* ctx, void* heap)
     ctx->ioRecvCb = wsEmbedRecv;
     ctx->ioSendCb = wsEmbedSend;
 #endif /* WOLFSSH_USER_IO */
+    ctx->countHighwater = DEFAULT_COUNT_HIGHWATER;
 
     return ctx;
 }
@@ -214,7 +215,7 @@ WOLFSSH* SshInit(WOLFSSH* ssh, WOLFSSH_CTX* ctx)
     ssh->wfd         = -1;         /* set to invalid */
     ssh->ioReadCtx   = &ssh->rfd;  /* prevent invalid access if not correctly */
     ssh->ioWriteCtx  = &ssh->wfd;  /* set */
-    ssh->countHighwater = DEFAULT_COUNT_HIGHWATER;
+    ssh->countHighwater = ctx->countHighwater;
     ssh->acceptState = ACCEPT_BEGIN;
     ssh->clientState = CLIENT_BEGIN;
     ssh->nextChannel = DEFAULT_NEXT_CHANNEL;
@@ -2406,6 +2407,8 @@ static INLINE int Encrypt(WOLFSSH* ssh, uint8_t* cipher, const uint8_t* input,
     ssh->txCount += sz;
     if (ssh->countHighwater && ssh->txCount > ssh->countHighwater) {
         WLOG(WS_LOG_DEBUG, "Transmit over high water mark");
+        if (ssh->ctx->highwaterCb)
+            ssh->ctx->highwaterCb(WOLFSSH_HWSIDE_TRANSMIT, ssh->highwaterCtx);
     }
 
     return ret;
@@ -2441,6 +2444,8 @@ static INLINE int Decrypt(WOLFSSH* ssh, uint8_t* plain, const uint8_t* input,
     ssh->rxCount += sz;
     if (ssh->countHighwater && ssh->rxCount > ssh->countHighwater) {
         WLOG(WS_LOG_DEBUG, "Receive over high water mark");
+        if (ssh->ctx->highwaterCb)
+            ssh->ctx->highwaterCb(WOLFSSH_HWSIDE_RECEIVE, ssh->highwaterCtx);
     }
 
     return ret;
diff --git a/src/ssh.c b/src/ssh.c
index fe9789ab..2348859a 100644
--- a/src/ssh.c
+++ b/src/ssh.c
@@ -154,9 +154,9 @@ int wolfSSH_get_fd(const WOLFSSH* ssh)
 }
 
 
-int wolfSSH_set_highwater(WOLFSSH* ssh, uint32_t highwater)
+int wolfSSH_SetHighwater(WOLFSSH* ssh, uint32_t highwater)
 {
-    WLOG(WS_LOG_DEBUG, "Entering wolfSSH_set_highwater()");
+    WLOG(WS_LOG_DEBUG, "Entering wolfSSH_SetHighwater()");
 
     if (ssh) {
         ssh->countHighwater = highwater;
@@ -168,14 +168,46 @@ int wolfSSH_set_highwater(WOLFSSH* ssh, uint32_t highwater)
 }
 
 
-uint32_t wolfSSH_get_highwater(WOLFSSH* ssh)
+uint32_t wolfSSH_GetHighwater(WOLFSSH* ssh)
 {
-    WLOG(WS_LOG_DEBUG, "Entering wolfSSH_get_highwater()");
+    WLOG(WS_LOG_DEBUG, "Entering wolfSSH_GetHighwater()");
 
     if (ssh)
         return ssh->countHighwater;
 
-    return WS_BAD_ARGUMENT;
+    return 0;
+}
+
+
+void wolfSSH_SetHighwaterCb(WOLFSSH_CTX* ctx, uint32_t highwater,
+                            WS_CallbackHighwater cb)
+{
+    WLOG(WS_LOG_DEBUG, "Entering wolfSSH_SetHighwaterCb()");
+
+    if (ctx) {
+        ctx->countHighwater = highwater;
+        ctx->highwaterCb = cb;
+    }
+}
+
+
+void wolfSSH_SetHighwaterCtx(WOLFSSH* ssh, void* ctx)
+{
+    WLOG(WS_LOG_DEBUG, "Entering wolfSSH_SetHighwaterCtx()");
+
+    if (ssh)
+        ssh->highwaterCtx = ctx;
+}
+
+
+void* wolfSSH_GetHighwaterCtx(WOLFSSH* ssh)
+{
+    WLOG(WS_LOG_DEBUG, "Entering wolfSSH_GetHighwaterCtx()");
+
+    if (ssh)
+        return ssh->highwaterCtx;
+
+    return NULL;
 }
 
 
diff --git a/wolfssh/internal.h b/wolfssh/internal.h
index 60d817b2..9cd5dc89 100644
--- a/wolfssh/internal.h
+++ b/wolfssh/internal.h
@@ -139,6 +139,7 @@ struct WOLFSSH_CTX {
     WS_CallbackIORecv   ioRecvCb;    /* I/O Receive Callback */
     WS_CallbackIOSend   ioSendCb;    /* I/O Send Callback */
     WS_CallbackUserAuth userAuthCb;  /* User Authentication Callback */
+    WS_CallbackHighwater highwaterCb; /* Data Highwater Mark Callback */
 
     uint8_t*            cert;        /* Owned by CTX */
     uint32_t            certSz;
@@ -146,6 +147,7 @@ struct WOLFSSH_CTX {
     uint32_t            caCertSz;
     uint8_t*            privateKey;  /* Owned by CTX */
     uint32_t            privateKeySz;
+    uint32_t            countHighwater;
 };
 
 
@@ -195,6 +197,7 @@ struct WOLFSSH {
     uint32_t       txCount;
     uint32_t       rxCount;
     uint32_t       countHighwater;
+    void*          highwaterCtx;
     uint32_t       curSz;
     uint32_t       seq;
     uint32_t       peerSeq;
diff --git a/wolfssh/ssh.h b/wolfssh/ssh.h
index be18430a..7394a01a 100644
--- a/wolfssh/ssh.h
+++ b/wolfssh/ssh.h
@@ -59,8 +59,16 @@ WOLFSSH_API void     wolfSSH_free(WOLFSSH*);
 WOLFSSH_API int  wolfSSH_set_fd(WOLFSSH*, int);
 WOLFSSH_API int  wolfSSH_get_fd(const WOLFSSH*);
 
-WOLFSSH_API int wolfSSH_set_highwater(WOLFSSH*, uint32_t);
-WOLFSSH_API uint32_t wolfSSH_get_highwater(WOLFSSH*);
+/* data high water mark functions */
+WOLFSSH_API int wolfSSH_SetHighwater(WOLFSSH*, uint32_t);
+WOLFSSH_API uint32_t wolfSSH_GetHighwater(WOLFSSH*);
+
+typedef int (*WS_CallbackHighwater)(uint8_t, void*);
+WOLFSSH_API void wolfSSH_SetHighwaterCb(WOLFSSH_CTX*, uint32_t,
+                                        WS_CallbackHighwater);
+WOLFSSH_API void wolfSSH_SetHighwaterCtx(WOLFSSH*, void*);
+WOLFSSH_API void* wolfSSH_GetHighwaterCtx(WOLFSSH*);
+
 
 WOLFSSH_API int wolfSSH_get_error(const WOLFSSH*);
 WOLFSSH_API const char* wolfSSH_get_error_name(const WOLFSSH*);
@@ -134,6 +142,13 @@ WOLFSSH_API int wolfSSH_KDF(uint8_t, uint8_t, uint8_t*, uint32_t,
                 const uint8_t*, uint32_t, const uint8_t*, uint32_t,
                 const uint8_t*, uint32_t);
 
+
+enum WS_HighwaterSide {
+    WOLFSSH_HWSIDE_TRANSMIT,
+    WOLFSSH_HWSIDE_RECEIVE
+};
+
+
 enum WS_EndpointTypes {
     WOLFSSH_ENDPOINT_SERVER,
     WOLFSSH_ENDPOINT_CLIENT