From 8db7cae42bbc8cf27abd95641d5a0c268b1f91b2 Mon Sep 17 00:00:00 2001 From: "Tom D." <15268361+anastygnome@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:43:04 +0200 Subject: [PATCH] uucore: add allocation-failure handling --- build.rs | 23 +- docs/src/l10n.md | 38 +++ src/bin/coreutils.rs | 12 +- src/bin/uudoc.rs | 8 +- src/common/validation.rs | 9 +- src/uu/false/src/false.rs | 5 + src/uu/true/src/true.rs | 5 + src/uucore/locales/en-US.ftl | 3 + src/uucore/locales/fr-FR.ftl | 1 + src/uucore/src/lib/lib.rs | 8 + src/uucore/src/lib/mods.rs | 1 + src/uucore/src/lib/mods/allocation.rs | 328 +++++++++++++++++++ src/uucore/src/lib/mods/clap_localization.rs | 2 + src/uucore_procs/src/lib.rs | 40 ++- 14 files changed, 465 insertions(+), 18 deletions(-) create mode 100644 src/uucore/src/lib/mods/allocation.rs diff --git a/build.rs b/build.rs index 2b61b84d2e4..78f91dd7852 100644 --- a/build.rs +++ b/build.rs @@ -66,7 +66,10 @@ pub fn main() { let mut mf = File::create(Path::new(&out_dir).join("uutils_map.rs")).unwrap(); mf.write_all( - "type UtilityMap = phf::OrderedMap<&'static str, (fn(T) -> i32, fn() -> Command)>;\n\ + "type UtilityMap = phf::OrderedMap<\n\ + &'static str,\n\ + (fn(T) -> i32, fn() -> Command, &'static allocation::AllocErrorConfig),\n\ + >;\n\ \n\ #[allow(clippy::too_many_lines)] #[allow(clippy::unreadable_literal)] @@ -77,27 +80,27 @@ pub fn main() { let mut phf_map = phf_codegen::OrderedMap::<&str>::new(); let mut entries = Vec::new(); + let map_value = |krate: &str| { + format!("({krate}::uumain, {krate}::uu_app, &{krate}::UU_ALLOC_ERROR_CONFIG)") + }; for krate in &crates { - let map_value = format!("({krate}::uumain, {krate}::uu_app)"); match krate.as_ref() { // 'test' is named uu_test to avoid collision with rust core crate 'test'. // It can also be invoked by name '[' for the '[ expr ] syntax'. "uu_test" => { - entries.push(("test", map_value.clone())); - entries.push(("[", map_value.clone())); + entries.push(("test", map_value(krate))); + entries.push(("[", map_value(krate))); } k if k.starts_with(OVERRIDE_PREFIX) => { - entries.push((&k[OVERRIDE_PREFIX.len()..], map_value.clone())); + entries.push((&k[OVERRIDE_PREFIX.len()..], map_value(krate))); } "false" | "true" => { - entries.push(( - krate.as_str(), - format!("(r#{krate}::uumain, r#{krate}::uu_app)"), - )); + let raw_krate = format!("r#{krate}"); + entries.push((krate.as_str(), map_value(&raw_krate))); } _ => { - entries.push((krate.as_str(), map_value.clone())); + entries.push((krate.as_str(), map_value(krate))); } } } diff --git a/docs/src/l10n.md b/docs/src/l10n.md index bb00aa2041c..28dac6962a8 100644 --- a/docs/src/l10n.md +++ b/docs/src/l10n.md @@ -53,6 +53,44 @@ The string parameter determines the lookup path for Fluent files. **English alwa --- +## 🧯 Allocation-failure diagnostics + +Standalone utilities and the `coreutils` multi-call binary install a +process-wide allocator. If an allocation fails, the allocator writes a +diagnostic without allocating more memory and exits with status `1` by default: + +``` +tsort: memory exhausted +``` + +The message is looked up in the following order: + +1. `-memory-exhausted` in the current Fluent bundle +2. `common-memory-exhausted` in the current bundle or its English fallback +3. The built-in `memory exhausted` text when localization is unavailable + +For example, a utility can provide its own localized message with: + +``` +tsort-memory-exhausted = memory exhausted while sorting +``` + +The allocation-failure status can be changed for utilities using the +`uucore::main` attribute: + +``` +#[uucore::main(alloc_error_exit_code = 2)] +pub fn uumain(args: impl uucore::Args) -> UResult<()> { + // ... +} +``` + +`alloc_error_exit_code` accepts values from `1` through `255`. It only controls +the status used when the global allocator handles an allocation failure; +ordinary utility errors keep their existing status. + +--- + ## 🌐 Locale Detection Locale selection is automatic and performed via: diff --git a/src/bin/coreutils.rs b/src/bin/coreutils.rs index 46418112286..8ddee54a77b 100644 --- a/src/bin/coreutils.rs +++ b/src/bin/coreutils.rs @@ -9,7 +9,11 @@ use itertools::Itertools as _; use std::cmp; use std::ffi::OsString; use std::io::{self, Write}; -use uucore::{Args, error::strip_errno}; +use uucore::{Args, allocation, error::strip_errno}; + +#[global_allocator] +static UU_ALLOCATOR: allocation::UuAllocator = + allocation::UuAllocator::new(&allocation::COREUTILS_ALLOC_ERROR_CONFIG); const VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -113,13 +117,17 @@ fn main() { _ => {} } - if let Some(&(uumain, _)) = utils.get(util) { + if let Some((&alloc_error_program, &(uumain, _, alloc_error_config))) = + utils.get_entry(util) + { + allocation::activate(alloc_error_program, alloc_error_config); // TODO: plug the deactivation of the translation // and load the English strings directly at compilation time in the // binary to avoid the load of the flt // Could be something like: // #[cfg(not(feature = "only_english"))] validation::setup_localization_or_exit(util); + allocation::localize_message(); exit(uumain(vec![util_os].into_iter().chain(args))); } // GNU coreutils --help string shows help for coreutils diff --git a/src/bin/uudoc.rs b/src/bin/uudoc.rs index 3f1f709a6c0..42404e110e9 100644 --- a/src/bin/uudoc.rs +++ b/src/bin/uudoc.rs @@ -29,8 +29,8 @@ use textwrap::{fill, indent, termwidth}; use zip::ZipArchive; use coreutils::validation; -use uucore::Args; use uucore::locale::get_message; +use uucore::{Args, allocation}; include!(concat!(env!("OUT_DIR"), "/uutils_map.rs")); @@ -112,7 +112,7 @@ fn usage(utils: &UtilityMap) { /// Generates the coreutils app for the utility map fn gen_coreutils_app(util_map: &UtilityMap) -> Command { let mut command = Command::new("coreutils"); - for (name, (_, sub_app)) in util_map { + for (name, (_, sub_app, _)) in util_map { // Recreate a small subcommand with only the relevant info // (name & short description) let about = sub_app() @@ -321,7 +321,7 @@ fn main() -> io::Result<()> { }; let mut utils = utils.entries().collect::>(); - utils.sort(); + utils.sort_by_key(|(name, _)| *name); println!("Writing util per platform table"); { @@ -361,7 +361,7 @@ fn main() -> io::Result<()> { } println!("Writing to utils"); - for (&name, (_, command)) in utils { + for (&name, (_, command, _)) in utils { let (utils_name, usage_name, command) = match name { "[" => { continue; diff --git a/src/common/validation.rs b/src/common/validation.rs index 4b84ca705a4..6506cf90ec9 100644 --- a/src/common/validation.rs +++ b/src/common/validation.rs @@ -42,7 +42,14 @@ pub fn exit(code: i32) -> ! { /// Gets all available utilities including "coreutils" #[allow(clippy::type_complexity)] pub fn get_all_utilities( - util_map: &phf::OrderedMap<&'static str, (fn(T) -> i32, fn() -> clap::Command)>, + util_map: &phf::OrderedMap< + &'static str, + ( + fn(T) -> i32, + fn() -> clap::Command, + &'static uucore::allocation::AllocErrorConfig, + ), + >, ) -> Vec<&'static str> { std::iter::once("coreutils") .chain(util_map.keys().copied()) diff --git a/src/uu/false/src/false.rs b/src/uu/false/src/false.rs index f143df5ac5e..336123d554a 100644 --- a/src/uu/false/src/false.rs +++ b/src/uu/false/src/false.rs @@ -7,6 +7,11 @@ use std::io::{self, Write as _}; use uucore::error::strip_errno; use uucore::{crate_version, show_error, translate}; +#[doc(hidden)] +/// Allocation-failure policy for this utility. +pub const UU_ALLOC_ERROR_CONFIG: uucore::allocation::AllocErrorConfig = + uucore::allocation::AllocErrorConfig::default_for(env!("CARGO_PKG_NAME")); + // uucore::main does not support no-result pub fn uumain(mut args: impl uucore::Args) -> i32 { // skip binary name diff --git a/src/uu/true/src/true.rs b/src/uu/true/src/true.rs index 83d5620edb2..e94d2a84750 100644 --- a/src/uu/true/src/true.rs +++ b/src/uu/true/src/true.rs @@ -7,6 +7,11 @@ use std::io::{self, Write as _}; use uucore::error::strip_errno; use uucore::{crate_version, show_error, translate}; +#[doc(hidden)] +/// Allocation-failure policy for this utility. +pub const UU_ALLOC_ERROR_CONFIG: uucore::allocation::AllocErrorConfig = + uucore::allocation::AllocErrorConfig::default_for(env!("CARGO_PKG_NAME")); + // uucore::main does not support no-result pub fn uumain(mut args: impl uucore::Args) -> i32 { // skip binary name diff --git a/src/uucore/locales/en-US.ftl b/src/uucore/locales/en-US.ftl index 4267ae52b03..89f85b45d0f 100644 --- a/src/uucore/locales/en-US.ftl +++ b/src/uucore/locales/en-US.ftl @@ -9,6 +9,9 @@ common-help = help common-version = version common-read-error = read error common-write-error = write error +# Fatal allocation failure. Utilities may define `-memory-exhausted` in +# their own Fluent resource to override this diagnostic for that utility. +common-memory-exhausted = memory exhausted # Common clap error messages clap-error-unexpected-argument = { $error_word }: unexpected argument '{ $arg }' found diff --git a/src/uucore/locales/fr-FR.ftl b/src/uucore/locales/fr-FR.ftl index 8adc24cc8f3..4ac2ec5297b 100644 --- a/src/uucore/locales/fr-FR.ftl +++ b/src/uucore/locales/fr-FR.ftl @@ -9,6 +9,7 @@ common-help = aide common-version = version common-write-error = erreur d'écriture common-read-error = erreur de lecture +common-memory-exhausted = mémoire épuisée # Messages d'erreur clap communs clap-error-unexpected-argument = { $error_word } : argument inattendu '{ $arg }' trouvé clap-error-unexpected-argument-simple = argument inattendu diff --git a/src/uucore/src/lib/lib.rs b/src/uucore/src/lib/lib.rs index 553eac23fe9..740c9b127c9 100644 --- a/src/uucore/src/lib/lib.rs +++ b/src/uucore/src/lib/lib.rs @@ -22,6 +22,7 @@ mod mods; // core cross-platform modules pub use uucore_procs::*; // * cross-platform modules +pub use crate::mods::allocation; pub use crate::mods::clap_localization; pub use crate::mods::display; pub use crate::mods::error; @@ -206,10 +207,16 @@ pub fn get_canonical_util_name(util_name: &str) -> &str { #[macro_export] macro_rules! bin_inner { ($util:ident, $post:expr) => { + #[global_allocator] + static UU_ALLOCATOR: $crate::allocation::UuAllocator = + $crate::allocation::UuAllocator::new(&$util::UU_ALLOC_ERROR_CONFIG); + pub fn main() { use std::io::Write; use uucore::locale; + $crate::allocation::activate(stringify!($util), &$util::UU_ALLOC_ERROR_CONFIG); + // Preserve inherited SIGPIPE settings (e.g., from env --default-signal=PIPE) uucore::panic::preserve_inherited_sigpipe(); @@ -226,6 +233,7 @@ macro_rules! bin_inner { } std::process::exit(99) }); + $crate::allocation::localize_message(); // execute utility code let code = $util::uumain(uucore::args_os()); diff --git a/src/uucore/src/lib/mods.rs b/src/uucore/src/lib/mods.rs index e33bf031958..97f8274ee63 100644 --- a/src/uucore/src/lib/mods.rs +++ b/src/uucore/src/lib/mods.rs @@ -4,6 +4,7 @@ // file that was distributed with this source code. // mods ~ cross-platforms modules (core/bundler file) +pub mod allocation; pub mod clap_localization; pub mod display; pub mod error; diff --git a/src/uucore/src/lib/mods/allocation.rs b/src/uucore/src/lib/mods/allocation.rs new file mode 100644 index 00000000000..9a56ca2e235 --- /dev/null +++ b/src/uucore/src/lib/mods/allocation.rs @@ -0,0 +1,328 @@ +// This file is part of the uutils coreutils package. +// +// For the full copyright and license information, please view the LICENSE +// file that was distributed with this source code. + +//! Process-wide allocation failure handling for uutils binaries. + +use super::locale; +use std::alloc::{GlobalAlloc, Layout, System}; +use std::ptr; +use std::slice; +use std::str; +use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicUsize, Ordering}; + +/// GNU-compatible allocation failure text used before localization is ready. +pub const DEFAULT_MESSAGE: &str = "memory exhausted"; + +/// GNU-compatible default allocation failure exit status. +pub const DEFAULT_EXIT_CODE: u8 = 1; + +/// Shared Fluent message used when a utility does not provide its own OOM text. +pub const DEFAULT_MESSAGE_ID: &str = "common-memory-exhausted"; + +/// Per-utility allocation failure configuration. +pub struct AllocErrorConfig { + utility: &'static str, + exit_code: u8, +} + +impl AllocErrorConfig { + /// Create allocation failure configuration for a utility. + pub const fn new(utility: &'static str, exit_code: u8) -> Self { + Self { utility, exit_code } + } + + /// Create allocation failure configuration using the GNU-compatible exit status. + pub const fn default_for(utility: &'static str) -> Self { + Self::new(utility, DEFAULT_EXIT_CODE) + } +} + +/// Allocation failure configuration used before the multicall binary selects a utility. +pub static COREUTILS_ALLOC_ERROR_CONFIG: AllocErrorConfig = + AllocErrorConfig::default_for("coreutils"); + +static ALLOCATION_FAILURE_IN_PROGRESS: AtomicBool = AtomicBool::new(false); +static ACTIVE_CONFIG: AtomicPtr = AtomicPtr::new(ptr::null_mut()); +static PROGRAM_PTR: AtomicPtr = AtomicPtr::new(ptr::null_mut()); +static PROGRAM_LEN: AtomicUsize = AtomicUsize::new(0); +static MESSAGE_PTR: AtomicPtr = AtomicPtr::new(ptr::null_mut()); +static MESSAGE_LEN: AtomicUsize = AtomicUsize::new(0); + +/// Select the allocation failure policy for the utility that is about to run. +/// +/// `program` and `config` must both have static storage because an allocation +/// failure can occur at any later point in the process. +pub fn activate(program: &'static str, config: &'static AllocErrorConfig) { + // Clear a previously localized message before switching utilities. This is + // mostly useful to embedders/tests; the normal multicall path runs one util. + MESSAGE_PTR.store(ptr::null_mut(), Ordering::Release); + MESSAGE_LEN.store(0, Ordering::Relaxed); + + // Publish the program before committing the config. An OOM racing this + // update may use the previous exit code, but it will never observe the new + // config with an unpublished program name. + publish_text(&PROGRAM_PTR, &PROGRAM_LEN, program); + ACTIVE_CONFIG.store(ptr::from_ref(config).cast_mut(), Ordering::Release); +} + +/// Resolve the allocation failure message through the current Fluent bundle. +/// +/// This must be called after localization has been initialized. Utility Fluent +/// resources are loaded with overriding semantics. A utility may customize its +/// diagnostic by defining `-memory-exhausted` (for example, +/// `tsort-memory-exhausted`) in its Fluent resource. If that key is absent, +/// `common-memory-exhausted` is used. If localization is unavailable, the +/// allocator keeps using [`DEFAULT_MESSAGE`]. +pub fn localize_message() { + let config = active_config(&COREUTILS_ALLOC_ERROR_CONFIG); + let utility = default_program(config); + let utility_message_id = format!("{utility}-memory-exhausted"); + let utility_message = locale::get_message(&utility_message_id); + + let message = if utility_message == utility_message_id { + locale::get_message(DEFAULT_MESSAGE_ID) + } else { + utility_message + }; + + if message == DEFAULT_MESSAGE_ID || message == DEFAULT_MESSAGE { + return; + } + + // The OOM path cannot own or allocate a String. Keep the localized value + // alive for the remainder of the process and publish only its borrowed bytes. + let message = Box::leak(message.into_boxed_str()); + publish_text(&MESSAGE_PTR, &MESSAGE_LEN, message); +} + +fn publish_text(ptr: &AtomicPtr, len: &AtomicUsize, text: &'static str) { + // Publish the length before the pointer. A reader that observes the pointer + // through an Acquire load also observes the matching length. + len.store(text.len(), Ordering::Relaxed); + ptr.store(text.as_ptr().cast_mut(), Ordering::Release); +} + +fn load_text(ptr: &AtomicPtr, len: &AtomicUsize) -> Option<&'static str> { + let ptr = ptr.load(Ordering::Acquire); + if ptr.is_null() { + return None; + } + + let len = len.load(Ordering::Relaxed); + + // SAFETY: publish_text() only publishes pointers to &'static str values, + // and the Release/Acquire pair makes the matching length visible here. + unsafe { Some(str::from_utf8_unchecked(slice::from_raw_parts(ptr, len))) } +} + +fn active_config(default: &'static AllocErrorConfig) -> &'static AllocErrorConfig { + let ptr = ACTIVE_CONFIG.load(Ordering::Acquire); + if ptr.is_null() { + default + } else { + // SAFETY: activate() only stores pointers obtained from &'static values. + unsafe { &*ptr } + } +} + +fn default_program(config: &AllocErrorConfig) -> &str { + config.utility.strip_prefix("uu_").unwrap_or(config.utility) +} + +fn allocation_failed(default: &'static AllocErrorConfig) -> ! { + let config = active_config(default); + + // Guard the best-effort diagnostic path against accidental allocator + // re-entry on less common targets. A recursive failure must terminate + // immediately rather than recurse indefinitely. + if ALLOCATION_FAILURE_IN_PROGRESS.swap(true, Ordering::AcqRel) { + raw::exit(config.exit_code); + } + let program = load_text(&PROGRAM_PTR, &PROGRAM_LEN).unwrap_or(default_program(config)); + let program = program.strip_prefix("uu_").unwrap_or(program); + let message = load_text(&MESSAGE_PTR, &MESSAGE_LEN).unwrap_or(DEFAULT_MESSAGE); + + raw::write_stderr(program.as_bytes()); + raw::write_stderr(b": "); + raw::write_stderr(message.as_bytes()); + raw::write_stderr(b"\n"); + raw::exit(config.exit_code) +} + +/// Global allocator used by uutils executables. +/// +/// It delegates normal allocation to [`System`]. If the system allocator +/// returns null, it emits the configured diagnostic without allocating and +/// terminates with the configured status. +pub struct UuAllocator { + default_config: &'static AllocErrorConfig, +} + +impl UuAllocator { + /// Create an allocator with a process-startup fallback configuration. + pub const fn new(default_config: &'static AllocErrorConfig) -> Self { + Self { default_config } + } + + #[cold] + #[inline(never)] + fn allocation_failed(&self) -> ! { + allocation_failed(self.default_config) + } +} + +// SAFETY: every allocator operation is forwarded to System with the same +// arguments and contract. Null allocation results are converted into process +// termination before they can escape to the caller, except for zero-sized +// operations where System's result is valid and is preserved. +unsafe impl GlobalAlloc for UuAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + // SAFETY: the caller supplied a valid GlobalAlloc layout. + let ptr = unsafe { System.alloc(layout) }; + if ptr.is_null() && layout.size() != 0 { + self.allocation_failed(); + } + ptr + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + // SAFETY: the caller supplied a valid GlobalAlloc layout. + let ptr = unsafe { System.alloc_zeroed(layout) }; + if ptr.is_null() && layout.size() != 0 { + self.allocation_failed(); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + // SAFETY: the caller upholds GlobalAlloc::dealloc's contract. + unsafe { System.dealloc(ptr, layout) }; + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + // SAFETY: the caller upholds GlobalAlloc::realloc's contract. + let new_ptr = unsafe { System.realloc(ptr, layout, new_size) }; + if new_ptr.is_null() && new_size != 0 { + self.allocation_failed(); + } + new_ptr + } +} + +#[cfg(unix)] +mod raw { + use std::ffi::c_void; + + unsafe extern "C" { + fn write(fd: i32, buf: *const c_void, count: usize) -> isize; + fn _exit(status: i32) -> !; + } + + pub(super) fn write_stderr(mut bytes: &[u8]) { + while !bytes.is_empty() { + // SAFETY: bytes is valid for bytes.len() readable bytes. + let written = unsafe { write(2, bytes.as_ptr().cast(), bytes.len()) }; + if written <= 0 { + return; + } + bytes = &bytes[written as usize..]; + } + } + + pub(super) fn exit(code: u8) -> ! { + // SAFETY: _exit terminates the process immediately. + unsafe { _exit(i32::from(code)) } + } +} + +#[cfg(windows)] +mod raw { + use std::ffi::c_void; + use std::ptr; + + type Handle = *mut c_void; + + const STD_ERROR_HANDLE: u32 = -12_i32 as u32; + + #[link(name = "kernel32")] + unsafe extern "system" { + #[link_name = "GetStdHandle"] + fn get_std_handle(n_std_handle: u32) -> Handle; + #[link_name = "WriteFile"] + fn write_file( + file: Handle, + buffer: *const c_void, + bytes_to_write: u32, + bytes_written: *mut u32, + overlapped: *mut c_void, + ) -> i32; + #[link_name = "ExitProcess"] + fn exit_process(exit_code: u32) -> !; + } + + pub(super) fn write_stderr(mut bytes: &[u8]) { + // SAFETY: GetStdHandle accepts the documented STD_ERROR_HANDLE value. + let handle = unsafe { get_std_handle(STD_ERROR_HANDLE) }; + if handle.is_null() { + return; + } + + while !bytes.is_empty() { + let len = bytes.len().min(u32::MAX as usize); + let mut written = 0; + // SAFETY: buffer is valid for len bytes and written points to a u32. + let ok = unsafe { + write_file( + handle, + bytes.as_ptr().cast(), + len as u32, + &raw mut written, + ptr::null_mut(), + ) + }; + if ok == 0 || written == 0 { + return; + } + bytes = &bytes[written as usize..]; + } + } + + pub(super) fn exit(code: u8) -> ! { + // SAFETY: ExitProcess terminates the current process immediately. + unsafe { exit_process(u32::from(code)) } + } +} + +#[cfg(not(any(unix, windows)))] +mod raw { + use std::io::Write as _; + + pub(super) fn write_stderr(bytes: &[u8]) { + let _ = std::io::stderr().write_all(bytes); + } + + pub(super) fn exit(code: u8) -> ! { + std::process::exit(i32::from(code)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn default_program_removes_utility_prefix() { + let config = AllocErrorConfig::default_for("uu_tsort"); + assert_eq!(default_program(&config), "tsort"); + } + + #[test] + fn published_text_is_read_after_release() { + let ptr = AtomicPtr::new(ptr::null_mut()); + let len = AtomicUsize::new(0); + publish_text(&ptr, &len, "memory exhausted"); + assert_eq!(load_text(&ptr, &len), Some("memory exhausted")); + } +} diff --git a/src/uucore/src/lib/mods/clap_localization.rs b/src/uucore/src/lib/mods/clap_localization.rs index bd699793325..73e31879160 100644 --- a/src/uucore/src/lib/mods/clap_localization.rs +++ b/src/uucore/src/lib/mods/clap_localization.rs @@ -692,6 +692,7 @@ mod tests { let required_keys = [ "common-error", + "common-memory-exhausted", "common-usage", "common-tip", "common-help-suggestion", @@ -724,6 +725,7 @@ mod tests { assert_eq!(get_message("common-error"), "erreur"); assert_eq!(get_message("common-usage"), "Utilisation"); assert_eq!(get_message("common-tip"), "conseil"); + assert_eq!(get_message("common-memory-exhausted"), "mémoire épuisée"); } unsafe { diff --git a/src/uucore_procs/src/lib.rs b/src/uucore_procs/src/lib.rs index ff604c49cba..6fc5e8836ef 100644 --- a/src/uucore_procs/src/lib.rs +++ b/src/uucore_procs/src/lib.rs @@ -15,6 +15,26 @@ use quote::quote; //* ref: @@ //* ref: [path construction from LitStr](https://oschwald.github.io/maxminddb-rust/syn/struct.LitStr.html) @@ +fn alloc_error_exit_code(args: &str) -> Option { + args.split(',').find_map(|arg| { + let (name, value) = arg.split_once('=')?; + if name.trim() != "alloc_error_exit_code" { + return None; + } + + let code = value.trim().parse::().unwrap_or_else(|_| { + panic!("alloc_error_exit_code must be an integer between 1 and 255") + }); + + assert!( + code != 0, + "alloc_error_exit_code must be an integer between 1 and 255" + ); + + Some(code) + }) +} + /// A procedural macro to define the main function of a uutils binary. /// /// This macro handles: @@ -22,13 +42,31 @@ use quote::quote; /// - SIGPIPE restoration to default if parent didn't explicitly ignore it /// - Disabling Rust signal handlers for proper core dumps /// - Error handling and exit code management +/// - Per-utility allocation failure exit code configuration +/// +/// The allocation failure exit code defaults to 1. Utilities that require a +/// different status can use `#[uucore::main(alloc_error_exit_code = 2)]`. #[proc_macro_attribute] pub fn main(args: TokenStream, stream: TokenStream) -> TokenStream { let stream = proc_macro2::TokenStream::from(stream); + let args = args.to_string(); // Some utils e.g. true does not require signals - let signals = !args.to_string().contains("no_signals"); + let signals = !args.split(',').any(|arg| arg.trim() == "no_signals"); + let alloc_error_exit_code = alloc_error_exit_code(&args); + let alloc_error_config = if let Some(code) = alloc_error_exit_code { + quote!(uucore::allocation::AllocErrorConfig::new(env!("CARGO_PKG_NAME"), #code)) + } else { + quote!(uucore::allocation::AllocErrorConfig::default_for(env!( + "CARGO_PKG_NAME" + ))) + }; let new = quote!( + #[doc(hidden)] + /// Allocation-failure policy generated by `uucore::main`. + pub const UU_ALLOC_ERROR_CONFIG: uucore::allocation::AllocErrorConfig = + #alloc_error_config; + // Initialize SIGPIPE state capture at process startup (Unix only). // This must be at module level to set up the .init_array static that runs // before main() to capture whether SIGPIPE was ignored by the parent process.