From f69cf930648999cedd0716c6804f46b0c1fbf504 Mon Sep 17 00:00:00 2001
From: Andrew Kelley <andrew@ziglang.org>
Date: Fri, 7 May 2021 21:23:51 -0700
Subject: [PATCH] std: start code increases stack size as appropriate on linux

closes #8708
---
 lib/std/start.zig | 59 ++++++++++++++++++++++++++++++++---------------
 src/main.zig      | 12 ----------
 2 files changed, 40 insertions(+), 31 deletions(-)

diff --git a/lib/std/start.zig b/lib/std/start.zig
index d9ec173bbc..38c4036bd7 100644
--- a/lib/std/start.zig
+++ b/lib/std/start.zig
@@ -297,26 +297,47 @@ fn posixCallMainAndExit() noreturn {
             std.os.linux.tls.initStaticTLS();
         }
 
-        // TODO This is disabled because what should we do when linking libc and this code
-        // does not execute? And also it's causing a test failure in stack traces in release modes.
+        // Linux ignores the stack size from the ELF file, and instead always gives 8 MiB.
+        // Here we look for the stack size in our program headers and tell the kernel,
+        // no, seriously, give me that stack space, I wasn't joking.
+        {
+            var i: usize = 0;
+            var at_phnum: usize = undefined;
+            var at_phdr: usize = undefined;
+            while (auxv[i].a_type != std.elf.AT_NULL) : (i += 1) {
+                switch (auxv[i].a_type) {
+                    std.elf.AT_PHNUM => at_phnum = auxv[i].a_un.a_val,
+                    std.elf.AT_PHDR => at_phdr = auxv[i].a_un.a_val,
+                    else => continue,
+                }
+            }
+            const phdrs = (@intToPtr([*]std.elf.Phdr, at_phdr))[0..at_phnum];
+            for (phdrs) |*phdr| {
+                switch (phdr.p_type) {
+                    std.elf.PT_GNU_STACK => {
+                        const wanted_stack_size = phdr.p_memsz;
+                        assert(wanted_stack_size % std.mem.page_size == 0);
 
-        //// Linux ignores the stack size from the ELF file, and instead always does 8 MiB. A further
-        //// problem is that it uses PROT_GROWSDOWN which prevents stores to addresses too far down
-        //// the stack and requires "probing". So here we allocate our own stack.
-        //const wanted_stack_size = gnu_stack_phdr.p_memsz;
-        //assert(wanted_stack_size % std.mem.page_size == 0);
-        //// Allocate an extra page as the guard page.
-        //const total_size = wanted_stack_size + std.mem.page_size;
-        //const new_stack = std.os.mmap(
-        //    null,
-        //    total_size,
-        //    std.os.PROT_READ | std.os.PROT_WRITE,
-        //    std.os.MAP_PRIVATE | std.os.MAP_ANONYMOUS,
-        //    -1,
-        //    0,
-        //) catch @panic("out of memory");
-        //std.os.mprotect(new_stack[0..std.mem.page_size], std.os.PROT_NONE) catch {};
-        //std.os.exit(@call(.{.stack = new_stack}, callMainWithArgs, .{argc, argv, envp}));
+                        std.os.setrlimit(.STACK, .{
+                            .cur = wanted_stack_size,
+                            .max = wanted_stack_size,
+                        }) catch {
+                            // If this is a debug build, it will be useful to find out
+                            // why this failed. If it is a release build, we allow the
+                            // stack overflow to cause a segmentation fault. Memory safety
+                            // is not compromised, however, depending on runtime state,
+                            // the application may crash due to provided stack space not
+                            // matching the known upper bound.
+                            if (builtin.mode == .Debug) {
+                                @panic("unable to increase stack size");
+                            }
+                        };
+                        break;
+                    },
+                    else => {},
+                }
+            }
+        }
     }
 
     std.os.exit(@call(.{ .modifier = .always_inline }, callMainWithArgs, .{ argc, argv, envp }));
diff --git a/src/main.zig b/src/main.zig
index ae7c8995a3..01be971602 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -180,18 +180,6 @@ pub fn mainArgs(gpa: *Allocator, arena: *Allocator, args: []const []const u8) !v
 
     defer log_scopes.deinit(gpa);
 
-    if (@import("builtin").target.os.tag == .linux) {
-        // Linux does not respect the stack size specified in the ELF, so we
-        // have to do this at runtime. TODO move this code to start.zig using
-        // the GNU_STACK program header.
-        std.os.setrlimit(.STACK, .{
-            .cur = 16 * 1024 * 1024,
-            .max = 16 * 1024 * 1024,
-        }) catch |err| {
-            warn("unable to increase stack size to 16 MiB", .{});
-        };
-    }
-
     const cmd = args[1];
     const cmd_args = args[2..];
     if (mem.eql(u8, cmd, "build-exe")) {