From 5c2dc2b7856a66338e2f113540fd2e8fcb847f9e Mon Sep 17 00:00:00 2001
From: Benjamin Chodroff <benjamin.chodroff@gmail.com>
Date: Wed, 8 Aug 2018 11:10:45 +0800
Subject: [PATCH] fix memory corruption and other 32bit overflows

---
 network/nw_buf.c  | 10 ++++++++--
 utils/ut_rpc.c    | 10 ++++++++--
 utils/ut_rpc.h    |  1 +
 utils/ut_ws_svr.c |  2 +-
 utils/ut_ws_svr.h |  2 ++
 5 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/network/nw_buf.c b/network/nw_buf.c
index f338af5f..fa42b0a5 100644
--- a/network/nw_buf.c
+++ b/network/nw_buf.c
@@ -8,7 +8,9 @@
 # include "nw_buf.h"
 
 # define NW_BUF_POOL_INIT_SIZE 64
+# define NW_BUF_POOL_MAX_SIZE  65535
 # define NW_CACHE_INIT_SIZE    64
+# define NW_CACHE_MAX_SIZE     65535
 
 size_t nw_buf_size(nw_buf *buf)
 {
@@ -85,7 +87,7 @@ void nw_buf_free(nw_buf_pool *pool, nw_buf *buf)
 {
     if (pool->free < pool->free_total) {
         pool->free_arr[pool->free++] = buf;
-    } else {
+    } else if (pool->free_total < NW_BUF_POOL_MAX_SIZE) {
         uint32_t new_free_total = pool->free_total * 2;
         void *new_arr = realloc(pool->free_arr, new_free_total * sizeof(nw_buf *));
         if (new_arr) {
@@ -95,6 +97,8 @@ void nw_buf_free(nw_buf_pool *pool, nw_buf *buf)
         } else {
             free(buf);
         }
+    } else {
+        free(buf);
     }
 }
 
@@ -230,7 +234,7 @@ void nw_cache_free(nw_cache *cache, void *obj)
 {
     if (cache->free < cache->free_total) {
         cache->free_arr[cache->free++] = obj;
-    } else {
+    } else if (cache->free_total < NW_CACHE_MAX_SIZE) {
         uint32_t new_free_total = cache->free_total * 2;
         void *new_arr = realloc(cache->free_arr, new_free_total * sizeof(void *));
         if (new_arr) {
@@ -240,6 +244,8 @@ void nw_cache_free(nw_cache *cache, void *obj)
         } else {
             free(obj);
         }
+    } else {
+        free(obj);
     }
 }
 
diff --git a/utils/ut_rpc.c b/utils/ut_rpc.c
index 4ade964b..14d0429a 100644
--- a/utils/ut_rpc.c
+++ b/utils/ut_rpc.c
@@ -44,13 +44,19 @@ int rpc_pack(rpc_pkg *pkg, void **data, uint32_t *size)
 {
     static void *send_buf;
     static size_t send_buf_size;
-    uint32_t pkg_size = RPC_PKG_HEAD_SIZE + pkg->ext_size + pkg->body_size;
+    uint32_t pkg_size;
+    if (pkg->body_size > RPC_PKG_MAX_BODY_SIZE) {
+        return -1;
+    }
+    pkg_size = RPC_PKG_HEAD_SIZE + pkg->ext_size + pkg->body_size;
     if (send_buf_size < pkg_size) {
         if (send_buf)
             free(send_buf);
         send_buf_size = pkg_size * 2;
         send_buf = malloc(send_buf_size);
-        assert(send_buf != NULL);
+        if (send_buf == NULL) {
+            return -1;
+        }
     }
 
     memcpy(send_buf, pkg, RPC_PKG_HEAD_SIZE);
diff --git a/utils/ut_rpc.h b/utils/ut_rpc.h
index a30cf534..2c7b2179 100644
--- a/utils/ut_rpc.h
+++ b/utils/ut_rpc.h
@@ -32,6 +32,7 @@ typedef struct rpc_pkg {
 # pragma pack()
 
 # define RPC_PKG_HEAD_SIZE (sizeof(rpc_pkg) - sizeof(void *) * 2)
+# define RPC_PKG_MAX_BODY_SIZE ((UINT32_MAX / 2) - UINT16_MAX - sizeof(rpc_pkg))
 
 int rpc_decode(nw_ses *ses, void *data, size_t max);
 int rpc_pack(rpc_pkg *pkg, void **data, uint32_t *size);
diff --git a/utils/ut_ws_svr.c b/utils/ut_ws_svr.c
index 5dfb87b4..5a4b8b10 100644
--- a/utils/ut_ws_svr.c
+++ b/utils/ut_ws_svr.c
@@ -129,7 +129,7 @@ static int on_http_message_complete(http_parser* parser)
     if (upgrade == NULL || strcasecmp(upgrade, "websocket") != 0)
         goto error;
     const char *connection = http_request_get_header(info->request, "Connection");
-    if (connection == NULL)
+    if (connection == NULL || strlen(connection) > UT_WS_SVR_MAX_HEADER_SIZE)
         goto error;
     else {
         bool found_upgrade = false;
diff --git a/utils/ut_ws_svr.h b/utils/ut_ws_svr.h
index d3fcf0f9..11c53d47 100644
--- a/utils/ut_ws_svr.h
+++ b/utils/ut_ws_svr.h
@@ -12,6 +12,8 @@
 # include "nw_buf.h"
 # include "nw_timer.h"
 
+# define UT_WS_SVR_MAX_HEADER_SIZE 1024
+
 typedef struct ws_svr_cfg {
     uint32_t bind_count;
     nw_svr_bind *bind_arr;