From d5b6a20557743911ff9f90af5d8ad24c699570d3 Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Thu, 12 Mar 2020 13:03:48 -0700
Subject: [PATCH] std: Don't abort process when printing panics in tests

This commit fixes an issue when using `set_print` and friends, notably
used by libtest, to avoid aborting the process if printing panics. This
previously panicked due to borrowing a mutable `RefCell` twice, and this
is worked around by borrowing these cells for less time, instead
taking out and removing contents temporarily.

Closes #69558
---
 src/libstd/io/stdio.rs                   | 12 ++++++++----
 src/test/ui/panic-while-printing.rs      | 24 ++++++++++++++++++++++++
 src/test/ui/test-panic-while-printing.rs | 24 ++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 4 deletions(-)
 create mode 100644 src/test/ui/panic-while-printing.rs
 create mode 100644 src/test/ui/test-panic-while-printing.rs

diff --git a/src/libstd/io/stdio.rs b/src/libstd/io/stdio.rs
index d410faca30d..0fb0757792e 100644
--- a/src/libstd/io/stdio.rs
+++ b/src/libstd/io/stdio.rs
@@ -792,10 +792,14 @@ fn print_to<T>(
 {
     let result = local_s
         .try_with(|s| {
-            if let Ok(mut borrowed) = s.try_borrow_mut() {
-                if let Some(w) = borrowed.as_mut() {
-                    return w.write_fmt(args);
-                }
+            // Note that we completely remove a local sink to write to in case
+            // our printing recursively panics/prints, so the recursive
+            // panic/print goes to the global sink instead of our local sink.
+            let prev = s.borrow_mut().take();
+            if let Some(mut w) = prev {
+                let result = w.write_fmt(args);
+                *s.borrow_mut() = Some(w);
+                return result;
             }
             global_s().write_fmt(args)
         })
diff --git a/src/test/ui/panic-while-printing.rs b/src/test/ui/panic-while-printing.rs
new file mode 100644
index 00000000000..7e9fa16b084
--- /dev/null
+++ b/src/test/ui/panic-while-printing.rs
@@ -0,0 +1,24 @@
+// run-pass
+// ignore-emscripten no subprocess support
+
+#![feature(set_stdio)]
+
+use std::fmt;
+use std::fmt::{Display, Formatter};
+use std::io::set_panic;
+
+pub struct A;
+
+impl Display for A {
+    fn fmt(&self, _f: &mut Formatter<'_>) -> fmt::Result {
+        panic!();
+    }
+}
+
+fn main() {
+    set_panic(Some(Box::new(Vec::new())));
+    assert!(std::panic::catch_unwind(|| {
+        eprintln!("{}", A);
+    })
+    .is_err());
+}
diff --git a/src/test/ui/test-panic-while-printing.rs b/src/test/ui/test-panic-while-printing.rs
new file mode 100644
index 00000000000..23f45407c1a
--- /dev/null
+++ b/src/test/ui/test-panic-while-printing.rs
@@ -0,0 +1,24 @@
+// compile-flags:--test
+// run-pass
+// ignore-emscripten no subprocess support
+
+use std::fmt;
+use std::fmt::{Display, Formatter};
+
+pub struct A(Vec<u32>);
+
+impl Display for A {
+    fn fmt(&self, _f: &mut Formatter<'_>) -> fmt::Result {
+        self.0[0];
+        Ok(())
+    }
+}
+
+#[test]
+fn main() {
+    let result = std::panic::catch_unwind(|| {
+        let a = A(vec![]);
+        eprintln!("{}", a);
+    });
+    assert!(result.is_err());
+}