From b04aba89d3a1931470983212925443e7aefca1e1 Mon Sep 17 00:00:00 2001
From: Rui Ueyama <ruiu@bluewhale.systems>
Date: Thu, 29 Jun 2023 15:56:35 +0900
Subject: [PATCH] Add `-z stack-size`

https://github.com/rui314/mold/issues/917
---
 elf/cmdline.cc           |  3 +++
 elf/mold.h               |  1 +
 elf/output-chunks.cc     | 13 ++++++++-----
 test/elf/z-stack-size.sh |  9 +++++++++
 4 files changed, 21 insertions(+), 5 deletions(-)
 create mode 100755 test/elf/z-stack-size.sh

diff --git a/elf/cmdline.cc b/elf/cmdline.cc
index a887e07d0c..332db414e1 100644
--- a/elf/cmdline.cc
+++ b/elf/cmdline.cc
@@ -191,6 +191,7 @@ inline const char helpmsg[] = R"(
                               Separate all loadable segments to different pages
     -z separate-code          Separate code and data into different pages
     -z noseparate-code        Allow overlap in pages
+  -z stack-size=VALUE         Set size of stack segment
   -z relro                    Make some sections read-only after relocation (default)
     -z norelro
   -z text                     Report error if DT_TEXTREL is set
@@ -866,6 +867,8 @@ std::vector<std::string> parse_nonpositional_args(Context<E> &ctx) {
       z_separate_code = SEPARATE_CODE;
     } else if (read_z_flag("noseparate-code")) {
       z_separate_code = NOSEPARATE_CODE;
+    } else if (read_z_arg("stack-size")) {
+      ctx.arg.z_stack_size = parse_number(ctx, "-z stack-size", arg);
     } else if (read_z_flag("dynamic-undefined-weak")) {
       ctx.arg.z_dynamic_undefined_weak = true;
     } else if (read_z_flag("nodynamic-undefined-weak")) {
diff --git a/elf/mold.h b/elf/mold.h
index ba9d4a94c2..a4eb650b78 100644
--- a/elf/mold.h
+++ b/elf/mold.h
@@ -1702,6 +1702,7 @@ struct Context {
     i64 filler = -1;
     i64 spare_dynamic_tags = 5;
     i64 thread_count = 0;
+    i64 z_stack_size = 0;
     std::string_view emulation;
     std::optional<Glob> unique;
     std::optional<u64> physical_image_base;
diff --git a/elf/output-chunks.cc b/elf/output-chunks.cc
index 88b37f11e7..caf4ff830a 100644
--- a/elf/output-chunks.cc
+++ b/elf/output-chunks.cc
@@ -325,11 +325,14 @@ static std::vector<ElfPhdr<E>> create_phdr(Context<E> &ctx) {
 
   // Add PT_GNU_STACK, which is a marker segment that doesn't really
   // contain any segments. It controls executable bit of stack area.
-  vec.push_back(ElfPhdr<E>{
-    .p_type = PT_GNU_STACK,
-    .p_flags = ctx.arg.z_execstack ? (PF_R | PF_W | PF_X) : (PF_R | PF_W),
-    .p_align = 1,
-  });
+  {
+    ElfPhdr<E> phdr = {};
+    phdr.p_type = PT_GNU_STACK;
+    phdr.p_flags = ctx.arg.z_execstack ? (PF_R | PF_W | PF_X) : (PF_R | PF_W);
+    phdr.p_memsz = ctx.arg.z_stack_size;
+    phdr.p_align = 1;
+    vec.push_back(phdr);
+  }
 
   // Create a PT_GNU_RELRO.
   if (ctx.arg.z_relro) {
diff --git a/test/elf/z-stack-size.sh b/test/elf/z-stack-size.sh
new file mode 100755
index 0000000000..d3fe717504
--- /dev/null
+++ b/test/elf/z-stack-size.sh
@@ -0,0 +1,9 @@
+#!/bin/bash
+. $(dirname $0)/common.inc
+
+cat <<EOF | $CC -o $t/a.o -c -xc -
+int main() {}
+EOF
+
+$CC -B. -o $t/exe $t/a.o -Wl,-z,stack-size=0x900000
+readelf -W --segments $t/exe | grep -q 'GNU_STACK .* 0x900000 RW'