[+] add protection for memory pool (#387)

alibaba · Dec 18, 2023 · e44ad40 · e44ad40
1 parent 20b163d
commit e44ad40
Show file tree

Hide file tree

Showing 6 changed files with 162 additions and 2 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -13,6 +13,7 @@ option (XQC_ENABLE_RENO "enable reno" OFF)
 option (XQC_ENABLE_UNLIMITED "enable unlimited cc" OFF)
 option (XQC_ENABLE_MP_INTEROP "enable MPQUIC interop" OFF)
 option (XQC_NO_PID_PACKET_PROCESS "do not expose path_id in xqc_engine_packet_process" OFF)
+option (XQC_PROTECT_POOL_MEM "enable write protection for pool memory (for debug)" OFF)
 
 if(NOT CMAKE_BUILD_TYPE)
     set(CMAKE_BUILD_TYPE Release)

diff --git a/include/xquic/xquic.h b/include/xquic/xquic.h
@@ -1257,6 +1257,9 @@ typedef struct xqc_conn_settings_s {
      */
     xqc_bool_t                  is_interop_mode;
 
+#ifdef XQC_PROTECT_POOL_MEM
+    uint8_t                     protect_pool_mem;
+#endif
 
 } xqc_conn_settings_t;
 

diff --git a/scripts/case_test.sh b/scripts/case_test.sh
@@ -853,6 +853,18 @@ else
     echo "$errlog"
 fi
 
+clear_log
+echo -e "send 10M data (mempool protected) ...\c"
+result=`${CLIENT_BIN}  -s 10240000 -l e -E -x 600 |grep ">>>>>>>> pass"`
+errlog=`grep_err_log`
+echo "$result"
+if [ -z "$errlog" ] && [ "$result" == ">>>>>>>> pass:1" ]; then
+    case_print_result "send_10M_data_mempool_protected" "pass"
+else
+    case_print_result "send_10M_data_mempool_protected" "fail"
+    echo "$errlog"
+fi
+
 clear_log
 echo -e "send 4K every time ...\c"
 result=`${CLIENT_BIN} -s 10240000 -l e -E -x 49|grep ">>>>>>>> pass"`
@@ -1166,12 +1178,13 @@ fi
 clear_log
 killall test_server
 echo -e "client Initial dcid corruption ...\c"
+sleep 1
 ${SERVER_BIN} -l d -e > /dev/null &
 sleep 1
 client_print_res=`${CLIENT_BIN} -s 1024000 -l d -t 1 -x 22 -E | grep ">>>>>>>> pass"`
 errlog=`grep_err_log`
 server_log_res=`grep "decrypt payload error" slog`
-server_conn_cnt=`grep "xqc_conn_create" slog | grep -v "tra_parameters_set" | wc -l`
+server_conn_cnt=`grep "xqc_conn_create" slog | grep -v "tra_parameters_set" | grep -v "mempool" | wc -l`
 echo "$client_print_res"
 if [ "$client_print_res" != "" ] && [ "$server_log_res" != "" ] && [ $server_conn_cnt -eq 2 ]; then
     case_print_result "client_initial_dcid_corruption" "pass"

diff --git a/src/common/xqc_memory_pool.h b/src/common/xqc_memory_pool.h
@@ -7,9 +7,19 @@
 
 #include <string.h>
 #include <stdint.h>
+#include <xquic/xquic.h>
 
 #include "src/common/xqc_malloc.h"
 
+#ifdef XQC_PROTECT_POOL_MEM
+#ifndef XQC_SYS_WINDOWS
+#include <stdlib.h>
+#include <unistd.h>
+#include <sys/mman.h>
+#include <assert.h>
+#endif
+#endif
+
 /* Interfaces:
  * xqc_memory_pool_t *xqc_create_pool(size_t size)
  * void xqc_destroy_pool(xqc_memory_pool_t* pool)
@@ -36,19 +46,78 @@ typedef struct xqc_memory_pool_s {
     xqc_memory_block_t *current;
     xqc_memory_large_t *large;  /* large chunk list */
     size_t              max;
+#ifdef XQC_PROTECT_POOL_MEM
+    xqc_bool_t          protect_block;
+    size_t              page_size;
+#endif
 } xqc_memory_pool_t;
 
 #define XQC_MAX_MALLOC_FROM_POOL (4096)
 
+#ifdef XQC_PROTECT_POOL_MEM
+static inline void *
+xqc_mempool_malloc_protected(size_t size, size_t page_sz)
+{
+#ifndef XQC_SYS_WINDOWS
+    int ret;
+    void *ptr = NULL;
+    ret = posix_memalign(&ptr, page_sz, page_sz + size);
+    if (ret != 0) {
+        return NULL;
+    }
+    ret = mprotect(ptr, page_sz, PROT_READ);
+    if (ret != 0) {
+        xqc_free(ptr);
+        return NULL;
+    }
+    return (void*)((char*)ptr + page_sz);
+#else
+    return xqc_malloc(size);
+#endif
+}
+
+static inline void
+xqc_mempool_free_protected(void* ptr, size_t page_sz) {
+#ifndef XQC_SYS_WINDOWS
+    void *start;
+    int ret;
+    start = (void*)((char*)ptr - page_sz);
+    ret = mprotect(start, page_sz, PROT_READ | PROT_WRITE | PROT_EXEC);
+    assert(ret == 0);
+    xqc_free(start);
+#else
+    xqc_free(ptr);
+#endif
+}
+#endif
 
+#ifdef XQC_PROTECT_POOL_MEM
+static inline xqc_memory_pool_t *
+xqc_create_pool(size_t size, xqc_bool_t protect_block)
+#else
 static inline xqc_memory_pool_t *
 xqc_create_pool(size_t size)
+#endif
 {
     if (size <= sizeof(xqc_memory_pool_t)) {
         return NULL;
     }
 
+#ifdef XQC_PROTECT_POOL_MEM
+    char *m;
+#ifndef XQC_SYS_WINDOWS
+    size_t page_sz = sysconf(_SC_PAGESIZE);
+#else
+    size_t page_sz = 4096;
+#endif
+    if (protect_block) {
+        m = xqc_mempool_malloc_protected(size, page_sz);
+    } else {
+        m = xqc_malloc(size);
+    }
+#else
     char *m = xqc_malloc(size);
+#endif
     if (m == NULL) {
         return NULL;
     }
@@ -58,6 +127,10 @@ xqc_create_pool(size_t size)
     pool->block.end = m + size;
     pool->block.failed = 0;
     pool->block.next = NULL;
+#ifdef XQC_PROTECT_POOL_MEM
+    pool->protect_block = protect_block;
+    pool->page_size = page_sz;
+#endif
 
     pool->current = &pool->block;
     pool->large = NULL;
@@ -76,23 +149,61 @@ xqc_destroy_pool(xqc_memory_pool_t *pool)
     while (block) {
         xqc_memory_block_t *p = block;
         block = block->next;
+#ifdef XQC_PROTECT_POOL_MEM
+        if (pool->protect_block) {
+            xqc_mempool_free_protected(p, pool->page_size);
+
+        } else {
+            xqc_free(p); 
+        }
+#else
         xqc_free(p);
+#endif
     }
 
     xqc_memory_large_t *large = pool->large;
     while (large) {
         xqc_memory_large_t * p = large;
         large = large->next;
+#ifdef XQC_PROTECT_POOL_MEM
+        if (pool->protect_block) {
+            xqc_mempool_free_protected(p, pool->page_size);
+
+        } else {
+            xqc_free(p); 
+        }
+#else
         xqc_free(p);
+#endif
     }
 
-    xqc_free(pool);
+#ifdef XQC_PROTECT_POOL_MEM
+    if (pool->protect_block) {
+        xqc_mempool_free_protected(pool, pool->page_size);
+
+    } else {
+        xqc_free(pool); 
+    }
+#else
+        xqc_free(pool);
+#endif
 }
 
 static inline void *
 xqc_palloc_large(xqc_memory_pool_t *pool, size_t size)
 {
+#ifdef XQC_PROTECT_POOL_MEM
+    xqc_memory_large_t *p;
+    if (pool->protect_block) {
+        p = xqc_mempool_malloc_protected(size + sizeof(xqc_memory_large_t), pool->page_size);
+
+    } else {
+        p = xqc_malloc(size + sizeof(xqc_memory_large_t));
+    }
+#else
     xqc_memory_large_t *p = xqc_malloc(size + sizeof(xqc_memory_large_t));
+#endif
+
     if (p == NULL) {
         return NULL;
     }
@@ -112,7 +223,18 @@ xqc_palloc_block(xqc_memory_pool_t *pool, size_t size)
 {
     size_t psize = pool->block.end - (char *)pool;
 
+#ifdef XQC_PROTECT_POOL_MEM
+    char *m;
+    if (pool->protect_block) {
+        m = xqc_mempool_malloc_protected(psize, pool->page_size);
+
+    } else {
+        m = xqc_malloc(psize);
+    }
+#else
     char *m = xqc_malloc(psize);
+#endif
+
     if (m == NULL) {
         return NULL;
     }

diff --git a/src/transport/xqc_conn.c b/src/transport/xqc_conn.c
@@ -80,6 +80,9 @@ xqc_conn_settings_t default_conn_settings = {
                                     .rtt_us_thr_low = 500000
                                   },
     .is_interop_mode            = 0,
+#ifdef XQC_PROTECT_POOL_MEM
+    .protect_pool_mem           = 0,
+#endif
 };
 
 
@@ -118,6 +121,9 @@ xqc_server_set_conn_settings(const xqc_conn_settings_t *settings)
     default_conn_settings.recv_rate_bytes_per_sec = settings->recv_rate_bytes_per_sec;
     default_conn_settings.enable_stream_rate_limit = settings->enable_stream_rate_limit;
     default_conn_settings.init_recv_window = settings->init_recv_window;
+#ifdef XQC_PROTECT_POOL_MEM
+    default_conn_settings.protect_pool_mem = settings->protect_pool_mem;
+#endif
 
     if (default_conn_settings.init_recv_window) {
         default_conn_settings.init_recv_window = xqc_max(default_conn_settings.init_recv_window, XQC_QUIC_MAX_MSS);
@@ -477,11 +483,20 @@ xqc_conn_create(xqc_engine_t *engine, xqc_cid_t *dcid, xqc_cid_t *scid,
     const xqc_conn_settings_t *settings, void *user_data, xqc_conn_type_t type)
 {
     xqc_connection_t *xc = NULL;
+#ifdef XQC_PROTECT_POOL_MEM
+    xqc_memory_pool_t *pool = xqc_create_pool(engine->config->conn_pool_size, settings->protect_pool_mem);
+#else
     xqc_memory_pool_t *pool = xqc_create_pool(engine->config->conn_pool_size);
+#endif
     if (pool == NULL) {
         return NULL;
     }
 
+#ifdef XQC_PROTECT_POOL_MEM
+    xqc_log(engine->log, XQC_LOG_DEBUG, "|mempool|protect:%d|page_sz:%z|",
+            pool->protect_block, pool->page_size);
+#endif
+
     xc = xqc_pcalloc(pool, sizeof(xqc_connection_t));
     if (xc == NULL) {
         goto fail;

diff --git a/tests/test_client.c b/tests/test_client.c
@@ -4477,6 +4477,12 @@ int main(int argc, char *argv[]) {
         .recv_rate_bytes_per_sec = rate_limit,
     };
 
+#ifdef XQC_PROTECT_POOL_MEM
+    if (g_test_case == 600) {
+        conn_settings.protect_pool_mem = 1;
+    }
+#endif
+
     xqc_stream_settings_t stream_settings = { .recv_rate_bytes_per_sec = 0 };
 
     if (g_test_case == 109) {