From 265c65a7bc767600c1e0876e265652c4b17c1734 Mon Sep 17 00:00:00 2001
From: Coldwings <coldwings@me.com>
Date: Tue, 23 Apr 2024 18:10:41 +0800
Subject: [PATCH] Add mprotect for the top-end of photon stack (#464)

* Add mprotect for the top-end of photon stack

* Fix some tests that may stack overflow in aarch64
---
 rpc/test/test-ooo.cpp | 12 ++++++------
 thread/test/test.cpp  | 10 +++++-----
 thread/thread.cpp     | 11 ++++++++---
 thread/timer.h        |  2 +-
 4 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/rpc/test/test-ooo.cpp b/rpc/test/test-ooo.cpp
index b6c2d38f..59428709 100644
--- a/rpc/test/test-ooo.cpp
+++ b/rpc/test/test-ooo.cpp
@@ -92,7 +92,7 @@ void* test_ooo_execution(void* args_)
 TEST(OutOfOrder, Execution) {
     OooEngine engine;
     for (int i = 0; i < 5; ++i)
-        thread_create(test_ooo_execution, &engine, 64 * 1024);
+        thread_create(test_ooo_execution, &engine, DEFAULT_STACK_SIZE);
 
     thread_yield();
     wait_for_completion();
@@ -167,7 +167,7 @@ TEST(OutOfOrder, keep_same_tag) {
               delete_ooo_execution_engine(engine);
           });
     for (int i=0;i<thread_num;i++) {
-        thread_create11(64*1024, test_ooo_same_tag, engine, &heavy_issue, &heavy_complete);
+        thread_create11(DEFAULT_STACK_SIZE, test_ooo_same_tag, engine, &heavy_issue, &heavy_complete);
     }
     thread_join(thread_enable_join(thread_create11(process_thread)));
 }
@@ -181,7 +181,7 @@ TEST(OutOfOrder, heavy_test) {
               delete_ooo_execution_engine(engine);
           });
     for (int i=0;i<thread_num;i++) {
-        thread_create11(64*1024, test_ooo, engine, &heavy_issue, &heavy_complete);
+        thread_create11(DEFAULT_STACK_SIZE, test_ooo, engine, &heavy_issue, &heavy_complete);
     }
     thread_join(thread_enable_join(thread_create11(process_thread)));
 }
@@ -226,7 +226,7 @@ TEST(OutOfOrder, error_issue) {
               delete_ooo_execution_engine(engine);
           });
     for (int i=0;i<thread_num;i++) {
-        thread_create11(64*1024, test_ooo, engine, &error_issue, &heavy_complete);
+        thread_create11(DEFAULT_STACK_SIZE, test_ooo, engine, &error_issue, &heavy_complete);
     }
     thread_join(thread_enable_join(thread_create11(process_thread)));
     log_output = log_output_stdout;
@@ -247,7 +247,7 @@ TEST(OutOfOrder, error_complete) {
               delete_ooo_execution_engine(engine);
           });
     for (int i=0;i<thread_num;i++) {
-        thread_create11(64*1024, test_ooo, engine, &heavy_issue, &error_complete);
+        thread_create11(DEFAULT_STACK_SIZE, test_ooo, engine, &heavy_issue, &error_complete);
     }
     thread_join(thread_enable_join(thread_create11(process_thread)));
     thread_sleep(1);
@@ -267,7 +267,7 @@ TEST(OutOfOrder, error_process) {
               delete_ooo_execution_engine(engine);
           });
     for (int i=0;i<thread_num;i++) {
-        thread_create11(64*1024, test_ooo, engine, &heavy_issue, &heavy_complete);
+        thread_create11(DEFAULT_STACK_SIZE, test_ooo, engine, &heavy_issue, &heavy_complete);
     }
     thread_join(thread_enable_join(thread_create11(error_process)));
     thread_sleep(1);
diff --git a/thread/test/test.cpp b/thread/test/test.cpp
index fe40f110..e1264e31 100644
--- a/thread/test/test.cpp
+++ b/thread/test/test.cpp
@@ -718,7 +718,7 @@ void *func1(void *)
 
 TEST(ThreadPool, test)
 {
-    ThreadPool<64> pool(64*1024);
+    ThreadPool<64> pool(DEFAULT_STACK_SIZE);
     vector<TPControl*> ths;
     ths.resize(FLAGS_ths_total);
     for (int i = 0; i<FLAGS_ths_total; i++)
@@ -733,7 +733,7 @@ TEST(ThreadPool, test)
 
 TEST(ThreadPool, migrate) {
     WorkPool wp(4, 0, 0, -1);
-    ThreadPool<64> pool(64 * 1024);
+    ThreadPool<64> pool(DEFAULT_STACK_SIZE);
     vector<TPControl*> ths;
     ths.resize(FLAGS_ths_total);
     for (int i = 0; i < FLAGS_ths_total; i++) {
@@ -750,7 +750,7 @@ TEST(ThreadPool, migrate) {
 
 TEST(ThreadPool, multithread) {
     WorkPool wp(4, 0, 0, -1);
-    ThreadPool<64> pool(64 * 1024);
+    ThreadPool<64> pool(DEFAULT_STACK_SIZE);
     vector<TPControl*> ths;
     ths.resize(FLAGS_ths_total);
     for (int i = 0; i < FLAGS_ths_total; i++) {
@@ -798,7 +798,7 @@ TEST(RWLock, checklock) {
         uint64_t arg = (i << 32) | (rand()%10 < 7 ? photon::RLOCK : photon::WLOCK);
         handles.emplace_back(
             photon::thread_enable_join(
-                photon::thread_create(&rwlocktest, (void*)arg, 64*1024)
+                photon::thread_create(&rwlocktest, (void*)arg)
             )
         );
     }
@@ -1086,7 +1086,7 @@ TEST(smp, rwlock) {
             uint64_t arg = (i << 32) | (rand()%10 < 7 ? photon::RLOCK : photon::WLOCK);
             handles.emplace_back(
                 photon::thread_enable_join(
-                    photon::thread_create(&smprwlocktest, (void*)arg, 64*1024)
+                    photon::thread_create(&smprwlocktest, (void*)arg)
                 )
             );
         }
diff --git a/thread/thread.cpp b/thread/thread.cpp
index 33163c89..40c243a9 100644
--- a/thread/thread.cpp
+++ b/thread/thread.cpp
@@ -305,6 +305,8 @@ namespace photon
             assert(state == states::DONE);
             // `buf` and `stack_size` will always store on register
             // when calling deallocating.
+            char* protect_head = (char*)align_up((uint64_t)buf, PAGE_SIZE);
+            mprotect(protect_head, PAGE_SIZE, PROT_READ | PROT_WRITE);
             photon_thread_dealloc(buf, stack_size);
         }
     };
@@ -935,12 +937,15 @@ R"(
             LOG_ERROR_RETURN(ENOSYS, nullptr, "Photon not initialized in this vCPU (OS thread)");
         size_t randomizer = (rand() % 32) * (1024 + 8);
         stack_size = align_up(randomizer + stack_size + sizeof(thread), PAGE_SIZE);
+        stack_size += PAGE_SIZE * 2;  // extra 2 pages for alignment and set guard page
         char* ptr = (char*)photon_thread_alloc(stack_size);
-        uint64_t p = (uint64_t) ptr + stack_size - sizeof(thread) - randomizer;
+        char* protect_head = (char*)align_up((uint64_t)ptr, PAGE_SIZE);
+        mprotect(protect_head, PAGE_SIZE, PROT_NONE);
+        uint64_t p = (uint64_t)ptr + stack_size - sizeof(thread) - randomizer;
         p = align_down(p, 64);
-        auto th = new((char*) p) thread;
+        auto th = new ((char*)p) thread;
         th->buf = ptr;
-        th->stackful_alloc_top = ptr;
+        th->stackful_alloc_top = protect_head + PAGE_SIZE;
         th->start = start;
         th->stack_size = stack_size;
         th->arg = arg;
diff --git a/thread/timer.h b/thread/timer.h
index c8e3dda5..d8e789b2 100644
--- a/thread/timer.h
+++ b/thread/timer.h
@@ -35,7 +35,7 @@ namespace photon
         // it has a `stack_size`, and the `on_timer` is invoked within the thread's context.
         // The timer object is deleted automatically after it is finished.
         Timer(uint64_t default_timeout, Entry on_timer, bool repeating = true,
-              uint64_t stack_size = 1024 * 64)
+              uint64_t stack_size = DEFAULT_STACK_SIZE)
         {
             _on_timer = on_timer;
             _default_timeout = default_timeout;