From a185b56b7caca17c7aa9d6f702fe1b2209c82e4e Mon Sep 17 00:00:00 2001
From: Tatsuyuki Ishi <ishitatsuyuki@gmail.com>
Date: Sun, 18 Mar 2018 23:02:06 +0900
Subject: [PATCH] Address review comments

---
 src/librustc_driver/lib.rs    |  8 +++++---
 src/libstd/rt.rs              | 15 +++++++++++++++
 src/libstd/sys/unix/thread.rs | 20 +++++++++++++++++++-
 src/libstd/thread/mod.rs      |  8 --------
 4 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/src/librustc_driver/lib.rs b/src/librustc_driver/lib.rs
index 76179723941..e39a2c2f5dc 100644
--- a/src/librustc_driver/lib.rs
+++ b/src/librustc_driver/lib.rs
@@ -1476,13 +1476,15 @@ pub fn in_rustc_thread<F, R>(f: F) -> Result<R, Box<Any + Send>>
         } else if rlim.rlim_max < STACK_SIZE as libc::rlim_t {
             true
         } else {
+            std::rt::deinit_stack_guard();
             rlim.rlim_cur = STACK_SIZE as libc::rlim_t;
             if libc::setrlimit(libc::RLIMIT_STACK, &mut rlim) != 0 {
                 let err = io::Error::last_os_error();
-                error!("in_rustc_thread: error calling setrlimit: {}", err);
-                true
+                // We have already deinited the stack. Further corruption is
+                // not allowed.
+                panic!("in_rustc_thread: error calling setrlimit: {}", err);
             } else {
-                std::thread::update_stack_guard();
+                std::rt::update_stack_guard();
                 false
             }
         }
diff --git a/src/libstd/rt.rs b/src/libstd/rt.rs
index e1392762a59..8f945470b7e 100644
--- a/src/libstd/rt.rs
+++ b/src/libstd/rt.rs
@@ -73,3 +73,18 @@ fn lang_start<T: ::process::Termination + 'static>
 {
     lang_start_internal(&move || main().report(), argc, argv)
 }
+
+/// Function used for reverting changes to the main stack before setrlimit().
+/// This is POSIX (non-Linux) specific and unlikely to be directly stabilized.
+#[unstable(feature = "rustc_stack_internals", issue = "0")]
+pub unsafe fn deinit_stack_guard() {
+    ::sys::thread::guard::deinit();
+}
+
+/// Function used for resetting the main stack guard address after setrlimit().
+/// This is POSIX specific and unlikely to be directly stabilized.
+#[unstable(feature = "rustc_stack_internals", issue = "0")]
+pub unsafe fn update_stack_guard() {
+    let main_guard = ::sys::thread::guard::init();
+    ::sys_common::thread_info::reset_guard(main_guard);
+}
diff --git a/src/libstd/sys/unix/thread.rs b/src/libstd/sys/unix/thread.rs
index 72cdb9440b8..d94e11a5207 100644
--- a/src/libstd/sys/unix/thread.rs
+++ b/src/libstd/sys/unix/thread.rs
@@ -222,7 +222,7 @@ pub mod guard {
 #[cfg_attr(test, allow(dead_code))]
 pub mod guard {
     use libc;
-    use libc::mmap;
+    use libc::{mmap, munmap};
     use libc::{PROT_NONE, MAP_PRIVATE, MAP_ANON, MAP_FAILED, MAP_FIXED};
     use ops::Range;
     use sys::os;
@@ -336,6 +336,24 @@ pub mod guard {
         }
     }
 
+    pub unsafe fn deinit() {
+        if !cfg!(target_os = "linux") {
+            if let Some(mut stackaddr) = get_stack_start() {
+                // Ensure address is aligned. Same as above.
+                let remainder = (stackaddr as usize) % PAGE_SIZE;
+                if remainder != 0 {
+                    stackaddr = ((stackaddr as usize) + PAGE_SIZE - remainder)
+                        as *mut libc::c_void;
+                }
+
+                // Undo the guard page mapping.
+                if munmap(stackaddr, PAGE_SIZE) != 0 {
+                    panic!("unable to deallocate the guard page");
+                }
+            }
+        }
+    }
+
     #[cfg(any(target_os = "macos",
               target_os = "bitrig",
               target_os = "openbsd",
diff --git a/src/libstd/thread/mod.rs b/src/libstd/thread/mod.rs
index b686ddc205e..71aee673cfe 100644
--- a/src/libstd/thread/mod.rs
+++ b/src/libstd/thread/mod.rs
@@ -208,14 +208,6 @@ pub use self::local::{LocalKey, AccessError};
 #[unstable(feature = "libstd_thread_internals", issue = "0")]
 #[doc(hidden)] pub use self::local::os::Key as __OsLocalKeyInner;
 
-/// Function used for resetting the main stack guard address after setrlimit().
-/// This is POSIX specific and unlikely to be directly stabilized.
-#[unstable(feature = "rustc_stack_internals", issue = "0")]
-pub unsafe fn update_stack_guard() {
-    let main_guard = imp::guard::init();
-    thread_info::reset_guard(main_guard);
-}
-
 ////////////////////////////////////////////////////////////////////////////////
 // Builder
 ////////////////////////////////////////////////////////////////////////////////