1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147

#![no_std]
//! Lightweight (zero dependency, proc_macro free) way to run code before main.
//!
//! This crate is the moral equivalent to the `ctor` crate, although the API is
//! completely different. The main reason for it's existence is:
//!
//! - Much faster to compile — no proc macros / syn / quote.
//!
//! - More obviously safe (avoids things I find dodgy, like `dtor`, ctors that
//!   initialize data, uses extern "C" function in function array called from C
//!   ...)
//!
//! - Try to handle untested unix platforms by assuming they support *at least*
//!   the `.ctors` section. This is in line with what clang does to compile c++
//!   static constructors.
//!
//! # Example
//!
//! ```
//! startup::on_startup! {
//!     println!("I'm running before main");
//! }
//! fn main() {
//!     println!("I'm inside main");
//! }
//! ```

/// Run some code automatically before the execution of `main`.
///
/// # Example
///
/// ```
/// startup::on_startup! {
///     println!("I'm running before main");
/// }
/// fn main() {
///     println!("I'm inside main");
/// }
/// ```
///
/// This outputs:
///
/// ```text
/// I'm running before main.
/// I'm inside main.
/// ```
///
/// # Caveats
///
/// - If your program is loaded as a dynamic library via dlopen/LoadLibrary,
///   it's not actually run "before main", but instead "when dlopen is called".
///   In practice, this doesn't matter for most use cases.
///
/// - This is on a best effort basis. There are known `rustc` bugs that will
///   prevent it from working. There are known platforms that don't support it
///   (wasm, maybe others). It is very important that your programs safety not
///   rely on this being called.
///
/// - The order two different `on_startup` invocations run in is totally
///   unspecified. Different platforms do wildly different things here. Do not
///   rely on one particular order. See also C++'s ["static initialization order
///   fiasco" (or problem)][static_init]
///
/// - Not all of the rust stdlib may be supported before main. It's best not to
///   call into it unless you're certain it will work.
///
/// [static_init]: https://isocpp.org/wiki/faq/ctors#static-init-order
#[macro_export]
macro_rules! on_startup {
    ($($tokens:tt)*) => {
        const _: () = {
            // pulled out and scoped to be unable to see the other defs because
            // of the issues around item-level hygene.
            extern "C" fn __init_function() {
                // Note: currently pointless, since even when loaded at runtime
                // via dlopen, panicing before main makes the stdlib abort.
                // However, if that ever changes in the future, we want to guard
                // against unwinding over an `extern "C"` boundary, so we force
                // a double-panic, which will trigger an abort (rather than have
                // any UB).
                let _guard = $crate::_private::PanicOnDrop;
                // Note: ensure we still forget the guard even if `$tokens` has
                // an explicit `return` in it somewhere.
                let _ = (|| -> () { $($tokens)* })();
                $crate::_private::forget(_guard);
            }
            {
                #[used]
                #[cfg_attr(
                    any(target_os = "macos", target_os = "ios", target_os = "tvos"),
                    link_section = "__DATA,__mod_init_func",
                )]
                // These definitely support .init_array
                #[cfg_attr(
                    any(
                        target_os = "linux",
                        target_os = "android",
                        target_os = "freebsd",
                        target_os = "netbsd",
                    ),
                    link_section = ".init_array"
                )]
                // Assume all other unixs support .ctors
                #[cfg_attr(all(
                    any(unix, all(target_os = "windows", target_env = "gnu")),
                    not(any(
                        target_os = "macos", target_os = "ios",
                        target_os = "tvos", target_os = "linux",
                        target_os = "android", target_os = "freebsd",
                        target_os = "netbsd",
                    ))
                ), link_section = ".ctors")]
                #[cfg_attr(all(windows, not(target_env = "gnu")), link_section = ".CRT$XCU")]
                static __CTOR: extern "C" fn() = __init_function;
            };
        };
    };
}

// Note: not part of the public api.
#[doc(hidden)]
pub mod _private {
    pub use core::mem::forget;
    pub struct PanicOnDrop;
    impl Drop for PanicOnDrop {
        #[cold]
        #[inline(never)]
        fn drop(&mut self) {
            panic!("Triggering abort via double panic (static initializer panicked).")
        }
    }
}

#[cfg(test)]
mod test {
    use core::sync::atomic::*;
    static VAL: AtomicU8 = AtomicU8::new(0);
    // do a few of them.
    on_startup! { VAL.fetch_add(1, Ordering::Relaxed); }
    on_startup! { VAL.fetch_add(2, Ordering::Relaxed); }
    on_startup! { VAL.fetch_add(4, Ordering::Relaxed); }

    #[test]
    fn smoke() {
        assert_eq!(VAL.load(Ordering::Relaxed), 1 + 2 + 4);
    }
}