Support `#[unix_sigpipe = "inherit|sig_dfl|sig_ign"]` on `fn main()`

This makes it possible to instruct libstd to never touch the signal
handler for `SIGPIPE`, which makes programs pipeable by default (e.g.
with `./your-program | head -n 1`) without `ErrorKind::BrokenPipe`
errors.

11 files changed, 102 insertions, 29 deletions

diff --git a/compiler/rustc_codegen_cranelift/src/main_shim.rs b/compiler/rustc_codegen_cranelift/src/main_shim.rs
index c67b6e98b32..3c024a84d90 100644
--- a/compiler/rustc_codegen_cranelift/src/main_shim.rs
+++ b/compiler/rustc_codegen_cranelift/src/main_shim.rs
@@ -1,7 +1,7 @@
 use rustc_hir::LangItem;
 use rustc_middle::ty::subst::GenericArg;
 use rustc_middle::ty::AssocKind;
-use rustc_session::config::EntryFnType;
+use rustc_session::config::{sigpipe, EntryFnType};
 use rustc_span::symbol::Ident;
 
 use crate::prelude::*;
@@ -15,12 +15,12 @@ pub(crate) fn maybe_create_entry_wrapper(
     is_jit: bool,
     is_primary_cgu: bool,
 ) {
-    let (main_def_id, is_main_fn) = match tcx.entry_fn(()) {
+    let (main_def_id, (is_main_fn, sigpipe)) = match tcx.entry_fn(()) {
         Some((def_id, entry_ty)) => (
             def_id,
             match entry_ty {
-                EntryFnType::Main => true,
-                EntryFnType::Start => false,
+                EntryFnType::Main { sigpipe } => (true, sigpipe),
+                EntryFnType::Start => (false, sigpipe::DEFAULT),
             },
         ),
         None => return,
@@ -35,7 +35,7 @@ pub(crate) fn maybe_create_entry_wrapper(
         return;
     }
 
-    create_entry_fn(tcx, module, unwind_context, main_def_id, is_jit, is_main_fn);
+    create_entry_fn(tcx, module, unwind_context, main_def_id, is_jit, is_main_fn, sigpipe);
 
     fn create_entry_fn(
         tcx: TyCtxt<'_>,
@@ -44,6 +44,7 @@ pub(crate) fn maybe_create_entry_wrapper(
         rust_main_def_id: DefId,
         ignore_lang_start_wrapper: bool,
         is_main_fn: bool,
+        sigpipe: u8,
     ) {
         let main_ret_ty = tcx.fn_sig(rust_main_def_id).output();
         // Given that `main()` has no arguments,
@@ -83,6 +84,7 @@ pub(crate) fn maybe_create_entry_wrapper(
             bcx.switch_to_block(block);
             let arg_argc = bcx.append_block_param(block, m.target_config().pointer_type());
             let arg_argv = bcx.append_block_param(block, m.target_config().pointer_type());
+            let arg_sigpipe = bcx.ins().iconst(types::I8, sigpipe as i64);
 
             let main_func_ref = m.declare_func_in_func(main_func_id, &mut bcx.func);
 
@@ -143,7 +145,8 @@ pub(crate) fn maybe_create_entry_wrapper(
                 let main_val = bcx.ins().func_addr(m.target_config().pointer_type(), main_func_ref);
 
                 let func_ref = m.declare_func_in_func(start_func_id, &mut bcx.func);
-                let call_inst = bcx.ins().call(func_ref, &[main_val, arg_argc, arg_argv]);
+                let call_inst =
+                    bcx.ins().call(func_ref, &[main_val, arg_argc, arg_argv, arg_sigpipe]);
                 bcx.inst_results(call_inst)[0]
             } else {
                 // using user-defined start fn
diff --git a/compiler/rustc_codegen_ssa/src/base.rs b/compiler/rustc_codegen_ssa/src/base.rs
index 4c6be3f9108..6e482062383 100644
--- a/compiler/rustc_codegen_ssa/src/base.rs
+++ b/compiler/rustc_codegen_ssa/src/base.rs
@@ -389,15 +389,14 @@ pub fn maybe_create_entry_wrapper<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
 
     let main_llfn = cx.get_fn_addr(instance);
 
-    let use_start_lang_item = EntryFnType::Start != entry_type;
-    let entry_fn = create_entry_fn::<Bx>(cx, main_llfn, main_def_id, use_start_lang_item);
+    let entry_fn = create_entry_fn::<Bx>(cx, main_llfn, main_def_id, entry_type);
     return Some(entry_fn);
 
     fn create_entry_fn<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
         cx: &'a Bx::CodegenCx,
         rust_main: Bx::Value,
         rust_main_def_id: DefId,
-        use_start_lang_item: bool,
+        entry_type: EntryFnType,
     ) -> Bx::Function {
         // The entry function is either `int main(void)` or `int main(int argc, char **argv)`,
         // depending on whether the target needs `argc` and `argv` to be passed in.
@@ -442,7 +441,7 @@ pub fn maybe_create_entry_wrapper<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
         let i8pp_ty = cx.type_ptr_to(cx.type_i8p());
         let (arg_argc, arg_argv) = get_argc_argv(cx, &mut bx);
 
-        let (start_fn, start_ty, args) = if use_start_lang_item {
+        let (start_fn, start_ty, args) = if let EntryFnType::Main { sigpipe } = entry_type {
             let start_def_id = cx.tcx().require_lang_item(LangItem::Start, None);
             let start_fn = cx.get_fn_addr(
                 ty::Instance::resolve(
@@ -454,8 +453,13 @@ pub fn maybe_create_entry_wrapper<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
                 .unwrap()
                 .unwrap(),
             );
-            let start_ty = cx.type_func(&[cx.val_ty(rust_main), isize_ty, i8pp_ty], isize_ty);
-            (start_fn, start_ty, vec![rust_main, arg_argc, arg_argv])
+
+            let i8_ty = cx.type_i8();
+            let arg_sigpipe = bx.const_u8(sigpipe);
+
+            let start_ty =
+                cx.type_func(&[cx.val_ty(rust_main), isize_ty, i8pp_ty, i8_ty], isize_ty);
+            (start_fn, start_ty, vec![rust_main, arg_argc, arg_argv, arg_sigpipe])
         } else {
             debug!("using user-defined start fn");
             let start_ty = cx.type_func(&[isize_ty, i8pp_ty], isize_ty);
diff --git a/compiler/rustc_feature/src/active.rs b/compiler/rustc_feature/src/active.rs
index 22178dd2123..89103a87623 100644
--- a/compiler/rustc_feature/src/active.rs
+++ b/compiler/rustc_feature/src/active.rs
@@ -519,6 +519,8 @@ declare_features! (
     /// Allows creation of instances of a struct by moving fields that have
     /// not changed from prior instances of the same struct (RFC #2528)
     (active, type_changing_struct_update, "1.58.0", Some(86555), None),
+    /// Enables rustc to generate code that instructs libstd to NOT ignore SIGPIPE.
+    (active, unix_sigpipe, "CURRENT_RUSTC_VERSION", Some(97889), None),
     /// Allows unsized fn parameters.
     (active, unsized_fn_params, "1.49.0", Some(48055), None),
     /// Allows unsized rvalues at arguments and parameters.
diff --git a/compiler/rustc_feature/src/builtin_attrs.rs b/compiler/rustc_feature/src/builtin_attrs.rs
index d520efed9b8..00621341ee3 100644
--- a/compiler/rustc_feature/src/builtin_attrs.rs
+++ b/compiler/rustc_feature/src/builtin_attrs.rs
@@ -359,6 +359,7 @@ pub const BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
     ),
 
     // Entry point:
+    gated!(unix_sigpipe, Normal, template!(Word, NameValueStr: "inherit|sig_ign|sig_dfl"), ErrorFollowing, experimental!(unix_sigpipe)),
     ungated!(start, Normal, template!(Word), WarnFollowing),
     ungated!(no_start, CrateLevel, template!(Word), WarnFollowing),
     ungated!(no_main, CrateLevel, template!(Word), WarnFollowing),
diff --git a/compiler/rustc_monomorphize/src/collector.rs b/compiler/rustc_monomorphize/src/collector.rs
index 82ef16a7f72..fceca58530c 100644
--- a/compiler/rustc_monomorphize/src/collector.rs
+++ b/compiler/rustc_monomorphize/src/collector.rs
@@ -1315,7 +1315,7 @@ impl<'v> RootCollector<'_, 'v> {
     /// the return type of `main`. This is not needed when
     /// the user writes their own `start` manually.
     fn push_extra_entry_roots(&mut self) {
-        let Some((main_def_id, EntryFnType::Main)) = self.entry_fn else {
+        let Some((main_def_id, EntryFnType::Main { .. })) = self.entry_fn else {
             return;
         };
 
diff --git a/compiler/rustc_passes/src/check_attr.rs b/compiler/rustc_passes/src/check_attr.rs
index 5789531d2ff..94a1e63e3f9 100644
--- a/compiler/rustc_passes/src/check_attr.rs
+++ b/compiler/rustc_passes/src/check_attr.rs
@@ -2117,6 +2117,7 @@ fn check_invalid_crate_level_attr(tcx: TyCtxt<'_>, attrs: &[Attribute]) {
         sym::automatically_derived,
         sym::start,
         sym::rustc_main,
+        sym::unix_sigpipe,
         sym::derive,
         sym::test,
         sym::test_case,
diff --git a/compiler/rustc_passes/src/entry.rs b/compiler/rustc_passes/src/entry.rs
index 7381019a620..cd10170d3ba 100644
--- a/compiler/rustc_passes/src/entry.rs
+++ b/compiler/rustc_passes/src/entry.rs
@@ -1,11 +1,11 @@
-use rustc_ast::{entry::EntryPointType, Attribute};
+use rustc_ast::entry::EntryPointType;
 use rustc_errors::struct_span_err;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, LocalDefId, CRATE_DEF_ID, LOCAL_CRATE};
 use rustc_hir::{ItemId, Node, CRATE_HIR_ID};
 use rustc_middle::ty::query::Providers;
 use rustc_middle::ty::{DefIdTree, TyCtxt};
-use rustc_session::config::{CrateType, EntryFnType};
+use rustc_session::config::{sigpipe, CrateType, EntryFnType};
 use rustc_session::parse::feature_err;
 use rustc_span::symbol::sym;
 use rustc_span::{Span, Symbol, DUMMY_SP};
@@ -71,14 +71,12 @@ fn entry_point_type(ctxt: &EntryContext<'_>, id: ItemId, at_root: bool) -> Entry
     }
 }
 
-fn err_if_attr_found(ctxt: &EntryContext<'_>, attrs: &[Attribute], sym: Symbol) {
+fn err_if_attr_found(ctxt: &EntryContext<'_>, id: ItemId, sym: Symbol, details: &str) {
+    let attrs = ctxt.tcx.hir().attrs(id.hir_id());
     if let Some(attr) = ctxt.tcx.sess.find_by_name(attrs, sym) {
         ctxt.tcx
             .sess
-            .struct_span_err(
-                attr.span,
-                &format!("`{}` attribute can only be used on functions", sym),
-            )
+            .struct_span_err(attr.span, &format!("`{}` attribute {}", sym, details))
             .emit();
     }
 }
@@ -87,14 +85,16 @@ fn find_item(id: ItemId, ctxt: &mut EntryContext<'_>) {
     let at_root = ctxt.tcx.opt_local_parent(id.def_id) == Some(CRATE_DEF_ID);
 
     match entry_point_type(ctxt, id, at_root) {
-        EntryPointType::None => (),
+        EntryPointType::None => {
+            err_if_attr_found(ctxt, id, sym::unix_sigpipe, "can only be used on `fn main()`");
+        }
         _ if !matches!(ctxt.tcx.def_kind(id.def_id), DefKind::Fn) => {
-            let attrs = ctxt.tcx.hir().attrs(id.hir_id());
-            err_if_attr_found(ctxt, attrs, sym::start);
-            err_if_attr_found(ctxt, attrs, sym::rustc_main);
+            err_if_attr_found(ctxt, id, sym::start, "can only be used on functions");
+            err_if_attr_found(ctxt, id, sym::rustc_main, "can only be used on functions");
         }
         EntryPointType::MainNamed => (),
         EntryPointType::OtherMain => {
+            err_if_attr_found(ctxt, id, sym::unix_sigpipe, "can only be used on root `fn main()`");
             ctxt.non_main_fns.push(ctxt.tcx.def_span(id.def_id));
         }
         EntryPointType::RustcMainAttr => {
@@ -116,6 +116,7 @@ fn find_item(id: ItemId, ctxt: &mut EntryContext<'_>) {
             }
         }
         EntryPointType::Start => {
+            err_if_attr_found(ctxt, id, sym::unix_sigpipe, "can only be used on `fn main()`");
             if ctxt.start_fn.is_none() {
                 ctxt.start_fn = Some((id.def_id, ctxt.tcx.def_span(id.def_id.to_def_id())));
             } else {
@@ -136,8 +137,9 @@ fn find_item(id: ItemId, ctxt: &mut EntryContext<'_>) {
 fn configure_main(tcx: TyCtxt<'_>, visitor: &EntryContext<'_>) -> Option<(DefId, EntryFnType)> {
     if let Some((def_id, _)) = visitor.start_fn {
         Some((def_id.to_def_id(), EntryFnType::Start))
-    } else if let Some((def_id, _)) = visitor.attr_main_fn {
-        Some((def_id.to_def_id(), EntryFnType::Main))
+    } else if let Some((local_def_id, _)) = visitor.attr_main_fn {
+        let def_id = local_def_id.to_def_id();
+        Some((def_id, EntryFnType::Main { sigpipe: sigpipe(tcx, def_id) }))
     } else {
         if let Some(main_def) = tcx.resolutions(()).main_def && let Some(def_id) = main_def.opt_fn_def_id() {
             // non-local main imports are handled below
@@ -161,13 +163,39 @@ fn configure_main(tcx: TyCtxt<'_>, visitor: &EntryContext<'_>) -> Option<(DefId,
                 )
                 .emit();
             }
-            return Some((def_id, EntryFnType::Main));
+            return Some((def_id, EntryFnType::Main { sigpipe: sigpipe(tcx, def_id) }));
         }
         no_main_err(tcx, visitor);
         None
     }
 }
 
+fn sigpipe(tcx: TyCtxt<'_>, def_id: DefId) -> u8 {
+    if let Some(attr) = tcx.get_attr(def_id, sym::unix_sigpipe) {
+        match (attr.value_str(), attr.meta_item_list()) {
+            (Some(sym::inherit), None) => sigpipe::INHERIT,
+            (Some(sym::sig_ign), None) => sigpipe::SIG_IGN,
+            (Some(sym::sig_dfl), None) => sigpipe::SIG_DFL,
+            (_, Some(_)) => {
+                // Keep going so that `fn emit_malformed_attribute()` can print
+                // an excellent error message
+                sigpipe::DEFAULT
+            }
+            _ => {
+                tcx.sess
+                    .struct_span_err(
+                        attr.span,
+                        "valid values for `#[unix_sigpipe = \"...\"]` are `inherit`, `sig_ign`, or `sig_dfl`",
+                    )
+                    .emit();
+                sigpipe::DEFAULT
+            }
+        }
+    } else {
+        sigpipe::DEFAULT
+    }
+}
+
 fn no_main_err(tcx: TyCtxt<'_>, visitor: &EntryContext<'_>) {
     let sp = tcx.def_span(CRATE_DEF_ID);
     if *tcx.sess.parse_sess.reached_eof.borrow() {
diff --git a/compiler/rustc_session/src/config.rs b/compiler/rustc_session/src/config.rs
index 162fc9aa0a6..e7c8538300e 100644
--- a/compiler/rustc_session/src/config.rs
+++ b/compiler/rustc_session/src/config.rs
@@ -36,6 +36,8 @@ use std::iter::{self, FromIterator};
 use std::path::{Path, PathBuf};
 use std::str::{self, FromStr};
 
+pub mod sigpipe;
+
 /// The different settings that the `-C strip` flag can have.
 #[derive(Clone, Copy, PartialEq, Hash, Debug)]
 pub enum Strip {
@@ -798,7 +800,15 @@ impl UnstableOptions {
 // The type of entry function, so users can have their own entry functions
 #[derive(Copy, Clone, PartialEq, Hash, Debug, HashStable_Generic)]
 pub enum EntryFnType {
-    Main,
+    Main {
+        /// Specifies what to do with `SIGPIPE` before calling `fn main()`.
+        ///
+        /// What values that are valid and what they mean must be in sync
+        /// across rustc and libstd, but we don't want it public in libstd,
+        /// so we take a bit of an unusual approach with simple constants
+        /// and an `include!()`.
+        sigpipe: u8,
+    },
     Start,
 }
 
diff --git a/compiler/rustc_session/src/config/sigpipe.rs b/compiler/rustc_session/src/config/sigpipe.rs
new file mode 100644
index 00000000000..31a3d9de05b
--- /dev/null
+++ b/compiler/rustc_session/src/config/sigpipe.rs
@@ -0,0 +1,20 @@
+/// Do not touch `SIGPIPE`. Use whatever the parent process uses.
+#[allow(dead_code)]
+pub const INHERIT: u8 = 1;
+
+/// Change `SIGPIPE` to `SIG_IGN` so that failed writes results in `EPIPE`
+/// that are eventually converted to `ErrorKind::BrokenPipe`.
+#[allow(dead_code)]
+pub const SIG_IGN: u8 = 2;
+
+/// Change `SIGPIPE` to `SIG_DFL` so that the process is killed when trying
+/// to write to a closed pipe. This is usually the desired behavior for CLI
+/// apps that produce textual output that you want to pipe to other programs
+/// such as `head -n 1`.
+#[allow(dead_code)]
+pub const SIG_DFL: u8 = 3;
+
+/// `SIG_IGN` has been the Rust default since 2014. See
+/// <https://github.com/rust-lang/rust/issues/62569>.
+#[allow(dead_code)]
+pub const DEFAULT: u8 = SIG_IGN;
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
index 1b47d11f6f8..92f1b2d4ea6 100644
--- a/compiler/rustc_span/src/symbol.rs
+++ b/compiler/rustc_span/src/symbol.rs
@@ -819,6 +819,7 @@ symbols! {
         infer_outlives_requirements,
         infer_static_outlives_requirements,
         inherent_associated_types,
+        inherit,
         inlateout,
         inline,
         inline_const,
@@ -1301,6 +1302,8 @@ symbols! {
         should_panic,
         shr,
         shr_assign,
+        sig_dfl,
+        sig_ign,
         simd,
         simd_add,
         simd_and,
@@ -1519,6 +1522,7 @@ symbols! {
         unit,
         universal_impl_trait,
         unix,
+        unix_sigpipe,
         unlikely,
         unmarked_api,
         unpin,
diff --git a/compiler/rustc_typeck/src/lib.rs b/compiler/rustc_typeck/src/lib.rs
index 5699f642baf..82b86d97d95 100644
--- a/compiler/rustc_typeck/src/lib.rs
+++ b/compiler/rustc_typeck/src/lib.rs
@@ -443,7 +443,7 @@ fn check_start_fn_ty(tcx: TyCtxt<'_>, start_def_id: DefId) {
 
 fn check_for_entry_fn(tcx: TyCtxt<'_>) {
     match tcx.entry_fn(()) {
-        Some((def_id, EntryFnType::Main)) => check_main_fn_ty(tcx, def_id),
+        Some((def_id, EntryFnType::Main { .. })) => check_main_fn_ty(tcx, def_id),
         Some((def_id, EntryFnType::Start)) => check_start_fn_ty(tcx, def_id),
         _ => {}
     }