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
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
//! Procedural macro for defining module startup functions.

extern crate proc_macro;
extern crate syn;
#[macro_use]
extern crate quote;

use proc_macro::TokenStream;

/// Marks a function as a library/executable constructor. This uses OS-specific
/// linker sections to call a specific function at load time.
///
/// Multiple startup functions are supported, but the invocation order is not
/// guaranteed.
///
/// # Examples
///
/// Print a startup message:
///
/// ```rust
/// # extern crate ctor;
/// # use ctor::*;
/// #[ctor]
/// fn foo() {
///   println!("Hello, world!");
/// }
///
/// # fn main() {
/// println!("main()");
/// # }
/// ```
///
/// Make changes to `static` variables:
///
/// ```rust
/// # extern crate ctor;
/// # use ctor::*;
/// # use std::sync::atomic::{AtomicBool, Ordering, ATOMIC_BOOL_INIT};
/// static INITED: AtomicBool = ATOMIC_BOOL_INIT;
///
/// #[ctor]
/// fn foo() {
///   INITED.store(true, Ordering::SeqCst);
/// }
/// ```
///
/// # Details
///
/// The `#[ctor]` macro makes use of linker sections to ensure that a
/// function is run at startup time.
///
/// The above example translates into the following Rust code (approximately):
///
///```rust
/// #[used]
/// #[cfg_attr(target_os = "linux", link_section = ".ctors")]
/// #[cfg_attr(target_os = "macos", link_section = "__DATA,__mod_init_func")]
/// #[cfg_attr(target_os = "windows", link_section = ".CRT$XCU")]
/// static foo: extern fn() = {
///   extern fn foo() { /* ... */ };
///   foo
/// };
/// ```
#[proc_macro_attribute]
pub fn ctor(_attribute: TokenStream, function: TokenStream) -> TokenStream {
    let function: syn::ItemFn = syn::parse_macro_input!(function);
    validate_item("ctor", &function);

    let syn::ItemFn {
        ident,
        unsafety,
        constness,
        abi,
        block,
        attrs,
        ..
    } = function;

    // Linux/ELF: https://www.exploit-db.com/papers/13234

    // Mac details: https://blog.timac.org/2016/0716-constructor-and-destructor-attributes/

    // Why .CRT$XCU on Windows? https://www.cnblogs.com/sunkang/archive/2011/05/24/2055635.html
    // 'I'=C init, 'C'=C++ init, 'P'=Pre-terminators and 'T'=Terminators

    let output = quote!(
        #[used]
        #[cfg_attr(target_os = "linux", link_section = ".ctors")]
        #[cfg_attr(target_os = "macos", link_section = "__DATA,__mod_init_func")]
        #[cfg_attr(target_os = "windows", link_section = ".CRT$XCU")]
        #(#attrs)*
        static #ident
        :
        extern #unsafety #abi #constness fn() =
        { extern #unsafety #abi #constness fn #ident() #block; #ident }
        ;
    );

    // eprintln!("{}", output);

    output.into()
}

/// Marks a function as a library/executable destructor. This uses OS-specific
/// linker sections to call a specific function at termination time.
///
/// Only available on Linux and OSX right now.
/// 
/// Multiple shutdown functions are supported, but the invocation order is not
/// guaranteed.
/// 
/// `sys_common::at_exit` is usually a better solution for shutdown handling, as
/// it allows you to use `stdout` in your handlers.
/// 
/// ```rust
/// # extern crate ctor;
/// # use ctor::*;
///
/// #[dtor]
/// fn shutdown() {
///   /* ... */
/// }
/// ```
#[cfg(not(target_os = "windows"))]
#[proc_macro_attribute]
pub fn dtor(_attribute: TokenStream, function: TokenStream) -> TokenStream {
    let function: syn::ItemFn = syn::parse_macro_input!(function);
    validate_item("dtor", &function);

    let syn::ItemFn {
        ident,
        unsafety,
        constness,
        abi,
        block,
        attrs,
        ..
    } = function;

    let output = quote!(
        #[used]
        #[cfg_attr(target_os = "linux", link_section = ".dtors")]
        #[cfg_attr(target_os = "macos", link_section = "__DATA,__mod_term_func")]
        #[cfg_attr(target_os = "windows", link_section = ".CRT$XTU")]
        #(#attrs)*
        static #ident
        :
        #unsafety extern #abi #constness fn() =
        { #unsafety extern #abi #constness fn #ident() #block; #ident }
        ;
    );

    // eprintln!("{}", output);

    output.into()
}

fn validate_item(typ: &str, item: &syn::ItemFn) {
    let syn::ItemFn {
        vis,
        decl,
        ..
    } = item;

    // Ensure that visibility modifier is not present
    match vis {
        syn::Visibility::Inherited => {}
        _ => panic!("#[{}] methods must not have visibility modifiers", typ),
    }

    // No parameters allowed
    if decl.inputs.len() > 0 {
        panic!("#[{}] methods may not have parameters", typ);
    }

    // No return type allowed
    match decl.output {
        syn::ReturnType::Default => {}
        _ => panic!("#[{}] methods must not have return types", typ),
    }
}