Merge pull request #293 from shintaro-iwasaki/pr/stack_canary

mem: implement a stack overflow detection mechanism
pmodels · Jan 27, 2021 · 245853d · 245853d
2 parents dc7ce7f + 23eb88d
commit 245853d
Show file tree

Hide file tree

Showing 5 changed files with 98 additions and 12 deletions.
diff --git a/configure.ac b/configure.ac
@@ -200,6 +200,14 @@ AC_ARG_ENABLE([tool],
     AS_HELP_STRING([--enable-tool],
                    [enable the tool interface, which is disabled by default.]))
 
+# --enable-stack-overflow-check
+AC_ARG_ENABLE([stack-overflow-check],
+[  --enable-stack-overflow-check@<:@=OPT@:>@ enable a stack overflow check
+        canary|canary-8     - use an 8-byte stack canary.
+        canary-XX           - use an XX-byte stack canary.
+        none|no
+],,[enable_stack_overflow_check=no])
+
 # --disable-simple-mutex
 AC_ARG_ENABLE([simple-mutex],
     AS_HELP_STRING([--disable-simple-mutex], [use an experimental mutex implementation]),,
@@ -679,6 +687,37 @@ AS_IF([test "x$enable_tool" != "xyes"],
       [AC_DEFINE(ABT_CONFIG_DISABLE_TOOL_INTERFACE, 1,
                  [Define to use the tool interface])])
 
+# --enable-stack-overflow-check
+stack_overflow_check_type="ABTI_STACK_CHECK_TYPE_NONE"
+stack_overflow_canary_size=0
+case "$enable_stack_overflow_check" in
+    canary)
+        stack_overflow_check_type="ABTI_STACK_CHECK_TYPE_CANARY"
+        stack_overflow_canary_size="8"
+    ;;
+    canary-*)
+        stack_overflow_canary_size_tmp="`echo $enable_stack_overflow_check | sed -e 's/canary-//g'`"
+        # stack_overflow_canary_size_tmp must be an integer
+        if test x"`echo $stack_overflow_canary_size_tmp | sed -e 's/[[0-9]]//g'`" = x"" ; then
+            stack_overflow_check_type="ABTI_STACK_CHECK_TYPE_CANARY"
+            stack_overflow_canary_size="$stack_overflow_canary_size_tmp"
+        else
+            AC_MSG_WARN([Unknown value $enable_stack_overflow_check for --enable-stack-overflow-check])
+        fi
+    ;;
+    none|no)
+        stack_overflow_check_type="ABTI_STACK_CHECK_TYPE_NONE"
+    ;;
+    *)
+        AC_MSG_WARN([Unknown value $enable_stack_overflow_check for --enable-stack-overflow-check])
+    ;;
+esac
+
+AC_DEFINE_UNQUOTED([ABT_CONFIG_STACK_CHECK_TYPE], [$stack_overflow_check_type],
+          [Define an algorithm of a stack overflow check])
+AC_DEFINE_UNQUOTED([ABT_CONFIG_STACK_CHECK_CANARY_SIZE], [$stack_overflow_canary_size],
+          [Define the size of a stack canary when it is enabled])
+
 # check if __atomic builtins are supported
 AC_TRY_COMPILE(
 [#include <stdint.h>],

diff --git a/src/include/abt.h.in b/src/include/abt.h.in
@@ -616,6 +616,8 @@ enum ABT_info_query_kind {
     ABT_INFO_QUERY_KIND_DYNAMIC_PROMOTION,
     /** Whether the stack unwinding feature is enabled or not */
     ABT_INFO_QUERY_KIND_ENABLED_STACK_UNWIND,
+    /** Whether the stack overflow check is enabled or not */
+    ABT_INFO_QUERY_KIND_ENABLED_STACK_OVERFLOW_CHECK,
 };
 
 /**

diff --git a/src/include/abti.h b/src/include/abti.h
@@ -61,6 +61,9 @@
 #define ABT_THREAD_TYPE_FULLY_FLEDGED 0
 #define ABT_THREAD_TYPE_DYNAMIC_PROMOTION 1
 
+#define ABTI_STACK_CHECK_TYPE_NONE 0
+#define ABTI_STACK_CHECK_TYPE_CANARY 1
+
 enum ABTI_xstream_type {
     ABTI_XSTREAM_TYPE_PRIMARY,
     ABTI_XSTREAM_TYPE_SECONDARY

diff --git a/src/include/abti_mem.h b/src/include/abti_mem.h
@@ -28,7 +28,40 @@ void ABTI_mem_finalize(ABTI_global *p_global);
 void ABTI_mem_finalize_local(ABTI_xstream *p_local_xstream);
 int ABTI_mem_check_lp_alloc(int lp_alloc);
 
+#define ABTI_STACK_CANARY_VALUE ((uint64_t)0xbaadc0debaadc0de)
+
 /* Inline functions */
+static inline void ABTI_mem_register_stack(void *p_stack, size_t stacksize)
+{
+#if ABT_CONFIG_STACK_CHECK_TYPE == ABTI_STACK_CHECK_TYPE_CANARY
+    /* Write down stack canary. */
+    if (p_stack) {
+        uint64_t i;
+        for (i = 0;
+             i < ABTU_roundup_uint64(ABT_CONFIG_STACK_CHECK_CANARY_SIZE, 8);
+             i += sizeof(uint64_t)) {
+            ((uint64_t *)p_stack)[i] = ABTI_STACK_CANARY_VALUE;
+        }
+    }
+#endif
+    ABTI_VALGRIND_REGISTER_STACK(p_stack, stacksize);
+}
+
+static inline void ABTI_mem_unregister_stack(void *p_stack)
+{
+#if ABT_CONFIG_STACK_CHECK_TYPE == ABTI_STACK_CHECK_TYPE_CANARY
+    if (p_stack) {
+        uint64_t i;
+        for (i = 0;
+             i < ABTU_roundup_uint64(ABT_CONFIG_STACK_CHECK_CANARY_SIZE, 8);
+             i += sizeof(uint64_t)) {
+            ABTI_ASSERT(((uint64_t *)p_stack)[i] == ABTI_STACK_CANARY_VALUE);
+        }
+    }
+#endif
+    ABTI_VALGRIND_UNREGISTER_STACK(p_stack);
+}
+
 ABTU_ret_err static inline int
 ABTI_mem_alloc_nythread_malloc(ABTI_thread **pp_thread)
 {
@@ -160,7 +193,7 @@ ABTI_mem_alloc_ythread_default(ABTI_local *p_local, ABTI_ythread **pp_ythread)
     /* Initialize members of ABTI_thread_attr. */
     p_ythread->p_stack = p_stack;
     p_ythread->stacksize = stacksize;
-    ABTI_VALGRIND_REGISTER_STACK(p_ythread->p_stack, p_ythread->stacksize);
+    ABTI_mem_register_stack(p_ythread->p_stack, p_ythread->stacksize);
     *pp_ythread = p_ythread;
     return ABT_SUCCESS;
 }
@@ -189,7 +222,7 @@ ABTU_ret_err static inline int ABTI_mem_alloc_ythread_mempool_desc_stack(
     /* Copy members of p_attr. */
     p_ythread->p_stack = p_stack;
     p_ythread->stacksize = stacksize;
-    ABTI_VALGRIND_REGISTER_STACK(p_ythread->p_stack, p_ythread->stacksize);
+    ABTI_mem_register_stack(p_ythread->p_stack, p_ythread->stacksize);
     *pp_ythread = p_ythread;
     return ABT_SUCCESS;
 }
@@ -211,7 +244,7 @@ ABTI_mem_alloc_ythread_malloc_desc_stack(ABTI_thread_attr *p_attr,
     p_ythread->thread.type = ABTI_THREAD_TYPE_MEM_MALLOC_DESC_STACK;
     p_ythread->stacksize = stacksize;
     p_ythread->p_stack = p_stack;
-    ABTI_VALGRIND_REGISTER_STACK(p_ythread->p_stack, p_ythread->stacksize);
+    ABTI_mem_register_stack(p_ythread->p_stack, p_ythread->stacksize);
     *pp_ythread = p_ythread;
     return ABT_SUCCESS;
 }
@@ -236,7 +269,7 @@ ABTU_ret_err static inline int ABTI_mem_alloc_ythread_mempool_desc(
     p_ythread->stacksize = p_attr->stacksize;
     p_ythread->p_stack = p_attr->p_stack;
     /* Note that the valgrind registration is ignored iff p_stack is NULL. */
-    ABTI_VALGRIND_REGISTER_STACK(p_ythread->p_stack, p_ythread->stacksize);
+    ABTI_mem_register_stack(p_ythread->p_stack, p_ythread->stacksize);
     *pp_ythread = p_ythread;
     return ABT_SUCCESS;
 }
@@ -248,7 +281,7 @@ static inline void ABTI_mem_free_thread(ABTI_local *p_local,
 #ifdef ABT_CONFIG_USE_MEM_POOL
     if (p_thread->type & ABTI_THREAD_TYPE_MEM_MEMPOOL_DESC_STACK) {
         ABTI_ythread *p_ythread = ABTI_thread_get_ythread(p_thread);
-        ABTI_VALGRIND_UNREGISTER_STACK(p_ythread->p_stack);
+        ABTI_mem_unregister_stack(p_ythread->p_stack);
 
         ABTI_xstream *p_local_xstream = ABTI_local_get_xstream_or_null(p_local);
         /* Came from a memory pool. */
@@ -266,24 +299,20 @@ static inline void ABTI_mem_free_thread(ABTI_local *p_local,
 #endif
         if (p_thread->type & ABTI_THREAD_TYPE_MEM_MEMPOOL_DESC) {
         /* Non-yieldable thread or yieldable thread without stack. */
-#ifdef HAVE_VALGRIND_SUPPORT
         ABTI_ythread *p_ythread = ABTI_thread_get_ythread_or_null(p_thread);
         if (p_ythread)
-            ABTI_VALGRIND_UNREGISTER_STACK(p_ythread->p_stack);
-#endif
+            ABTI_mem_unregister_stack(p_ythread->p_stack);
         ABTI_mem_free_nythread(p_local, p_thread);
     } else if (p_thread->type & ABTI_THREAD_TYPE_MEM_MALLOC_DESC_STACK) {
         ABTI_ythread *p_ythread = ABTI_thread_get_ythread(p_thread);
-        ABTI_VALGRIND_UNREGISTER_STACK(p_ythread->p_stack);
+        ABTI_mem_unregister_stack(p_ythread->p_stack);
         ABTU_free(p_ythread->p_stack);
     } else {
         ABTI_ASSERT(p_thread->type & ABTI_THREAD_TYPE_MEM_MALLOC_DESC);
         ABTI_STATIC_ASSERT(offsetof(ABTI_ythread, thread) == 0);
-#ifdef HAVE_VALGRIND_SUPPORT
         ABTI_ythread *p_ythread = ABTI_thread_get_ythread_or_null(p_thread);
         if (p_ythread)
-            ABTI_VALGRIND_UNREGISTER_STACK(p_ythread->p_stack);
-#endif
+            ABTI_mem_unregister_stack(p_ythread->p_stack);
         ABTU_free(p_thread);
     }
 }

diff --git a/src/info.c b/src/info.c
@@ -170,6 +170,12 @@ static void info_trigger_print_all_thread_stacks(
  *   to \c ABT_TRUE if Argobots is configured to enable the stack unwinding
  *   feature.  Otherwise, \c val is set to \c ABT_FALSE.
  *
+ * - \c ABT_INFO_QUERY_KIND_ENABLED_STACK_OVERFLOW_CHECK
+ *
+ *   \c val must be a pointer to a variable of type \c int.  \c val is set to 1
+ *   if Argobots is configured to use a stack canary to check stack overflow.
+ *   Otherwise, \c val is set to 0.
+ *
  * @changev20
  * \DOC_DESC_V1X_RETURN_INFO_IF_POSSIBLE
  * @endchangev20
@@ -333,6 +339,13 @@ int ABT_info_query_config(ABT_info_query_kind query_kind, void *val)
             *((ABT_bool *)val) = ABT_TRUE;
 #else
             *((ABT_bool *)val) = ABT_FALSE;
+#endif
+            break;
+        case ABT_INFO_QUERY_KIND_ENABLED_STACK_OVERFLOW_CHECK:
+#if ABT_CONFIG_STACK_CHECK_TYPE == ABTI_STACK_CHECK_TYPE_CANARY
+            *((int *)val) = 1;
+#else
+            *((int *)val) = 0;
 #endif
             break;
         default: