stackaroo/

lib.rs

//! # stackaroo
//!
//! A highly unsafe library for swapping out the OS-provided thread stack with custom stacks.
//!
//! ## Overview
//!
//! `stackaroo` provides platform-specific assembly implementations (x86_64 and AArch64) to perform
//! direct stack pointer manipulation, allowing you to execute functions on arbitrarily large custom
//! stacks that exceed OS-provided stack limitations.
//!
//! ## Safety
//!
//! This library is **extremely unsafe** and should only be used by people who understand the
//! implications of manual stack management. It's primarily intended for research, kernel development,
//! and testing scenarios where you need to bypass OS stack limitations.
//!
//! ## Supported Platforms
//!
//! - x86_64 (Intel/AMD 64-bit)
//! - AArch64 (ARM 64-bit)
//!
//! ## Features
//!
//! - `std` (default): Enables standard library support (implies `alloc`)
//! - `tls` (default): Enables thread-local storage for thread-safe stack swapping (requires `std`)
//! - `alloc`: Enables heap allocation support (required for `swap_to_heap`)
//!
//! ## FFI Bindings
//!
//! For C/C++ interoperability, see the `stackaroo-ffi` crate in this workspace.
//! It provides C-compatible bindings and generates a `stackaroo.h` header file.
//!
//! ## Concurrency
//!
//! - With `tls` feature: One swap per thread (uses thread-local storage, thread-safe)
//! - Without `tls` feature: One swap globally (uses static variables, not thread-safe)
//! - Does not support recursive/nested stack swaps
//!
//! ## Examples
//!
//! ### Basic usage with heap-allocated stack
//!
//! ```no_run
//! use stackaroo::swap_to_heap;
//!
//! fn deep_recursion(depth: usize) {
//!     let large_array = [0u8; 1024 * 1024]; // 1MB per frame
//!     std::hint::black_box(&large_array);   // Don't let it be compiled-out
//!     if depth > 0 {
//!         deep_recursion(depth - 1);
//!     }
//! }
//!
//! fn main() {
//!     unsafe {
//!         swap_to_heap(|_: &mut ()| deep_recursion(1024), None, 4 << 30)
//!     }.expect("Stack swap failed");
//! }
//! ```
//!
//! ### Usage with static stack
//!
//! ```no_run
//! use stackaroo::swap_to_static;
//!
//! static mut STACK: [u8; 1 << 26] = [0; 1 << 26]; // 64MB global-backed stack
//!
//! fn compute(arg: &mut u32) {
//!     *arg = *arg * 2 + 1;
//! }
//!
//! fn main() {
//!     unsafe {
//!         let mut value = 100u32;
//!         swap_to_static(compute, Some(&mut value), &mut STACK).unwrap();
//!         println!("Result: {}", value);
//!     }
//! }
//! ```
//!
//! ### Usage with argument passing
//!
//! ```no_run
//! use stackaroo::swap_to_heap;
//!
//! struct Args {
//!     input: u64,
//!     output: u64,
//! }
//!
//! fn fibonacci(args: &mut Args) {
//!     fn fib(n: u64) -> u64 {
//!         if n <= 1 { return n; }
//!         fib(n - 1) + fib(n - 2)
//!     }
//!     args.output = fib(args.input);
//! }
//!
//! fn main() {
//!     unsafe {
//!         let mut args = Args { input: 40, output: 0 };
//!         swap_to_heap(fibonacci, Some(&mut args), 1 << 28).unwrap();
//!         println!("Fibonacci({}) = {}", args.input, args.output);
//!     }
//! }
//! ```

#![cfg_attr(not(feature = "std"), no_std)]

#[cfg(feature = "alloc")]
extern crate alloc;

use core::ffi::c_void;
use core::result::Result;
use core::{ptr, sync::atomic};

pub(crate) mod helpers;
use helpers::*;

mod arch;

/// Error types that can occur during stack swap operations.
#[derive(Debug, PartialEq, PartialOrd, Eq, Ord)]
pub enum Error {
    /// The provided stack pointer is not properly aligned.
    ///
    /// Stack pointers must be aligned to the platform's requirements
    /// (typically 16 bytes on x86_64 and AArch64).
    StackPtrNotAligned,

    /// A stack swap is already in progress on the current thread (with `tls` feature)
    /// or globally (without `tls` feature).
    ///
    /// Nested or concurrent stack swaps are not supported. You must complete
    /// the current stack swap before initiating another one.
    StackSwapInProgress,
}

/// Swaps to a custom stack at a specific memory address, executes a function, then returns.
///
/// This is the low-level function that powers both `swap_to_heap` and
/// `swap_to_static`. It allows you to provide an arbitrary stack pointer,
/// giving you complete control over stack placement.
///
/// # Type Parameters
///
/// * `T` - The type of the argument passed to the callout function. Can be any type.
///
/// # Parameters
///
/// * `callout` - The function to execute on the new stack. It receives a mutable reference
///   to the argument.
/// * `arg` - An optional mutable reference to pass to the callout function. Use `None` if
///   no argument is needed.
/// * `new_stack_top` - A raw pointer to the top of the new stack (stacks grow downward).
///   Must be properly aligned (typically 16 bytes) and point to valid, writable memory.
///
/// # Returns
///
/// * `Ok(())` - The stack swap completed successfully.
/// * `Err(Error::StackPtrNotAligned)` - The provided stack pointer is not properly aligned.
/// * `Err(Error::StackSwapInProgress)` - Another stack swap is already in progress.
///
/// # Safety
///
/// This function is **extremely unsafe** and requires careful attention:
///
/// * `new_stack_top` must point to valid, writable memory with sufficient size
/// * The memory region must be properly aligned (typically 16 bytes for x86_64/AArch64)
/// * The stack must be large enough for the callout function's needs, including all
///   local variables, function calls, and recursion
/// * The callout function must not attempt to use stack references from before the swap
/// * The callout function must not unwind (panic) past the swap point
/// * Nested stack swaps are not supported and will return an error
/// * With the `tls` feature, this is thread-safe but only one swap per thread is allowed
/// * Without the `tls` feature, only one swap can be active globally
/// * The caller is responsible for ensuring the memory remains valid for the entire duration
/// * Stack memory must not be freed or deallocated while the swap is active
///
/// # Platform-Specific Details
///
/// * **x86_64**: Uses `rsp` register for stack pointer manipulation
/// * **AArch64**: Uses `sp` register for stack pointer manipulation
/// * Stacks grow downward, so `new_stack_top` should point to the highest address
///
/// # Examples
///
/// ```no_run
/// use stackaroo::swap_to;
/// use std::alloc::{alloc, dealloc, Layout};
///
/// unsafe {
///     // Allocate 1MB of memory for the stack
///     let layout = Layout::from_size_align(1 << 20, 16).unwrap();
///     let stack_bottom = alloc(layout);
///     let stack_top = stack_bottom.add(1 << 20);
///     
///     // Execute function on custom stack
///     let mut value = 42u32;
///     swap_to(
///         |v: &mut u32| { *v *= 2; },
///         Some(&mut value),
///         stack_top as *mut core::ffi::c_void
///     ).unwrap();
///     
///     assert_eq!(value, 84);
///     
///     // Clean up
///     dealloc(stack_bottom, layout);
/// }
/// ```
///
/// # Advanced Example: Memory-Mapped Stack
///
/// ```no_run
/// # #[cfg(unix)]
/// # {
/// use stackaroo::swap_to;
///
/// unsafe {
///     // Map 4MB of memory for a custom stack (Unix example)
///     let size = 4 * 1024 * 1024;
///     let addr = libc::mmap(
///         std::ptr::null_mut(),
///         size,
///         libc::PROT_READ | libc::PROT_WRITE,
///         libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
///         -1,
///         0,
///     );
///     
///     if addr == libc::MAP_FAILED {
///         panic!("mmap failed");
///     }
///     
///     let stack_top = (addr as *mut u8).add(size);
///     
///     swap_to(
///         |_: &mut ()| println!("Running on memory-mapped stack!"),
///         None,
///         stack_top as *mut core::ffi::c_void
///     ).unwrap();
///     
///     libc::munmap(addr, size);
/// }
/// # }
/// ```
///
/// # Notes
///
/// * Most users should prefer `swap_to_heap` or `swap_to_static` instead
/// * This function uses compiler memory fences to ensure proper ordering of operations
/// * The implementation uses platform-specific inline assembly for stack manipulation
/// * Consider using this function when you need custom memory management (e.g., memory-mapped stacks)
#[inline(never)]
pub unsafe fn swap_to<T>(
    callout: fn(&mut T) -> (),
    arg: Option<&mut T>,
    new_stack_top: *mut c_void,
) -> Result<(), Error> {
    if !new_stack_top.is_aligned() {
        return Err(Error::StackPtrNotAligned);
    }

    if !get_old_rsp().is_null() || !get_callout_arg().is_null() {
        return Err(Error::StackSwapInProgress);
    }

    // Preserve the callout function and argument in global variables
    set_callout(core::mem::transmute(callout));

    // Preserve the argument pointer in a global variable
    if let Some(arg) = arg {
        set_callout_arg(arg as *mut T as *mut c_void);
    } else {
        set_callout_arg(ptr::null_mut());
    }

    // Memory fence to ensure globals are written before stack switch
    atomic::compiler_fence(atomic::Ordering::SeqCst);

    // Save old stack and switch to new stack using platform-specific implementation
    arch::swap(get_old_rsp_ptr(), new_stack_top);

    // IMPORTANT NOTE:
    // If by any chance the compiler inserts here instructions referencing the stack,
    // this thing will burn and die.

    // Memory fence to ensure global variables are written before function call
    atomic::compiler_fence(atomic::Ordering::SeqCst);

    // Call the callout function on the new stack, with an argument stored in a global variables
    get_callout()(get_callout_arg());

    // Memory fence to ensure function call completes before stack restoration
    atomic::compiler_fence(atomic::Ordering::SeqCst);

    // Restore old stack using platform-specific implementation
    arch::restore(get_old_rsp_ptr());

    // IMPORTANT NOTE:
    // It's safe again for the compiler to insert stack-referencing instructions here.

    // Memory fence to ensure stack is restored before clearing globals
    atomic::compiler_fence(atomic::Ordering::SeqCst);

    set_old_rsp(ptr::null_mut());
    set_callout(dummy_callout);
    set_callout_arg(ptr::null_mut());

    Ok(())
}

/// Swaps to a heap-allocated stack, executes a function, then returns to the original stack.
///
/// This function allocates a new stack on the heap with the specified size, switches to it,
/// executes the provided callout function, and then restores the original stack.
///
/// # Type Parameters
///
/// * `T` - The type of the argument passed to the callout function. Can be any type.
///
/// # Parameters
///
/// * `callout` - The function to execute on the new stack. It receives a mutable reference
///   to the argument.
/// * `arg` - An optional mutable reference to pass to the callout function. Use `None` if
///   no argument is needed.
/// * `stack_size` - The size in bytes of the stack to allocate. Must be large enough for
///   the callout function's needs. Common sizes range from 1MB to 4GB.
///
/// # Returns
///
/// * `Ok(())` - The stack swap completed successfully.
/// * `Err(Error::StackPtrNotAligned)` - The calculated stack pointer is not properly aligned.
/// * `Err(Error::StackSwapInProgress)` - Another stack swap is already in progress.
///
/// # Safety
///
/// This function is highly unsafe for several reasons:
///
/// * The callout function must not attempt to use stack references from before the swap
/// * The stack size must be sufficient for the callout function's needs, including all
///   recursive calls and stack allocations
/// * Nested stack swaps are not supported and will return an error
/// * The callout function must not unwind (panic) past the swap point
/// * With the `tls` feature, this is thread-safe but only one swap per thread is allowed
/// * Without the `tls` feature, only one swap can be active globally
///
/// # Examples
///
/// ```no_run
/// use stackaroo::swap_to_heap;
///
/// // Example 1: Without argument
/// fn simple_function(_: &mut ()) {
///     println!("Running on custom stack!");
/// }
///
/// unsafe {
///     swap_to_heap(simple_function, None, 1 << 20).unwrap();
/// }
///
/// // Example 2: With argument
/// fn compute(value: &mut u32) {
///     *value = *value * 2;
/// }
///
/// unsafe {
///     let mut x = 42;
///     swap_to_heap(compute, Some(&mut x), 1 << 20).unwrap();
///     assert_eq!(x, 84);
/// }
/// ```
///
/// # Panics
///
/// If the heap allocation for the stack fails (out of memory).
///
/// # Feature Requirements
///
/// Requires the `alloc` feature to be enabled (enabled by default with `std`).
#[inline(never)]
#[cfg(feature = "alloc")]
pub unsafe fn swap_to_heap<T>(
    callout: fn(&mut T) -> (),
    arg: Option<&mut T>,
    stack_size: usize,
) -> Result<(), Error> {
    use alloc::vec;
    let mut new_stack = vec![0u8; stack_size];
    let new_stack_top = new_stack.as_mut_ptr().add(new_stack.len());
    swap_to(callout, arg, new_stack_top as *mut c_void)
}

/// Swaps to a static stack, executes a function, then returns to the original stack.
///
/// This function switches to a pre-allocated static stack buffer, executes the provided
/// callout function, and then restores the original stack. Unlike `swap_to_heap`,
/// this doesn't perform any allocation and works in `no_std` environments.
///
/// # Type Parameters
///
/// * `T` - The type of the argument passed to the callout function. Can be any type.
///
/// # Parameters
///
/// * `callout` - The function to execute on the new stack. It receives a mutable reference
///   to the argument.
/// * `arg` - An optional mutable reference to pass to the callout function. Use `None` if
///   no argument is needed.
/// * `stack` - A mutable reference to a static byte array to use as the stack. The array
///   must have a `'static` lifetime and be large enough for the callout function's needs.
///
/// # Returns
///
/// * `Ok(())` - The stack swap completed successfully.
/// * `Err(Error::StackPtrNotAligned)` - The calculated stack pointer is not properly aligned.
/// * `Err(Error::StackSwapInProgress)` - Another stack swap is already in progress.
///
/// # Safety
///
/// This function is highly unsafe for several reasons:
///
/// * The callout function must not attempt to use stack references from before the swap
/// * The static buffer must be large enough for the callout function's needs, including
///   all recursive calls and stack allocations
/// * Nested stack swaps are not supported and will return an error
/// * The callout function must not unwind (panic) past the swap point
/// * With the `tls` feature, this is thread-safe but only one swap per thread is allowed
/// * Without the `tls` feature, only one swap can be active globally
/// * The static buffer must not be used concurrently by multiple threads
///
/// # Examples
///
/// ```no_run
/// use stackaroo::swap_to_static;
///
/// static mut MY_STACK: [u8; 1 << 26] = [0; 1 << 26]; // 64MB
///
/// fn compute(value: &mut u32) {
///     *value = *value * 2 + 1;
/// }
///
/// unsafe {
///     let mut x = 100;
///     swap_to_static(compute, Some(&mut x), &mut MY_STACK).unwrap();
///     assert_eq!(x, 201);
/// }
/// ```
///
/// # Notes
///
/// * Stack grows downward, so the function uses the end of the buffer as the stack top
/// * The buffer is not initialized or cleared between uses
/// * This function works in both `std` and `no_std` environments
#[inline(never)]
pub unsafe fn swap_to_static<T>(
    callout: fn(&mut T) -> (),
    arg: Option<&mut T>,
    stack: &'static mut [u8],
) -> Result<(), Error> {
    let stack_ptr = stack.as_ptr() as *mut u8;
    let stack_top = stack_ptr.add(stack.len());
    swap_to(callout, arg, stack_top as *mut c_void)
}

#[cfg(test)]
#[allow(static_mut_refs)]
mod tests {
    use super::*;

    fn deep_recursion(depth: usize) {
        let large_array = [0u8; 1 << 20]; // 1 MB per stack frame
        core::hint::black_box(&large_array);
        if depth == 0 {
            return;
        }
        deep_recursion(depth - 1);
    }

    fn callout(_: &mut ()) {
        deep_recursion(1024); // Should consume over 1GB of stack
    }

    #[test]
    fn test_stack_swap() {
        // Run on 4GB stack
        unsafe { swap_to_heap(callout, None, 1 << 32) }.unwrap();
    }

    mod fibbonacci {
        use super::*;

        #[derive(Debug)]
        struct Args {
            n: u64,
            result: u64,
        }

        fn fibonacci_callout(arg: &mut Args) {
            // Deep recursive Fibonacci calculation that would overflow the regular stack
            fn fibonacci(n: u64) -> u64 {
                if n <= 1 {
                    return n;
                }
                fibonacci(n - 1) + fibonacci(n - 2)
            }
            arg.result = fibonacci(arg.n);
        }

        #[test]
        fn test_fibonacci() {
            const STACK_SIZE: usize = 1 << 28;
            unsafe {
                let mut args = Args { n: 35, result: 0 };
                swap_to_heap(fibonacci_callout, Some(&mut args), STACK_SIZE).unwrap();
                assert_eq!(args.result, 9227465);

                let mut args = Args { n: 40, result: 0 };
                swap_to_heap(fibonacci_callout, Some(&mut args), STACK_SIZE).unwrap();
                assert_eq!(args.result, 102334155);
            }
        }
    }

    #[test]
    fn test_no_concurrent_stack_swaps() {
        const STACK_SIZE: usize = 1 << 28;

        fn attempt_nested_swap(_: &mut ()) {
            // Try to initiate another stack swap while one is already in progress
            let result = unsafe { swap_to_heap(|_: &mut ()| {}, None, STACK_SIZE) };

            // This should fail with StackSwapInProgress error
            assert_eq!(result, Err(Error::StackSwapInProgress));
        }

        // Initiate the first stack swap which will attempt a nested swap
        let result = unsafe { swap_to_heap(attempt_nested_swap, None, STACK_SIZE) };

        // The outer stack swap should succeed
        assert!(result.is_ok());
    }

    mod global {
        use super::*;

        const STACK_SIZE: usize = 1 << 26;
        static mut GLOBAL_STACK: [u8; STACK_SIZE] = [255; STACK_SIZE];

        #[test]
        fn test_swap_to_with_global_static() {
            fn global_stack_callout(arg: &mut u32) {
                // Simple computation on the global static stack
                *arg = *arg * 2 + 1;
                // Add to our result for verification
                *arg += 42;
            }

            unsafe {
                let mut arg = 100u32;
                swap_to_static(global_stack_callout, Some(&mut arg), &mut GLOBAL_STACK).unwrap();
                // Expected: (100 * 2 + 1) + 42 = 243
                assert_eq!(arg, 243);
            }
        }
    }

    #[cfg(feature = "tls")]
    #[test]
    fn test_thread_safety() {
        use std::sync::{Arc, Barrier};
        use std::thread;

        fn simple_callout(_: &mut ()) {
            // Just a simple function that runs on the new stack
        }

        let barrier = Arc::new(Barrier::new(4));
        let mut handles = vec![];

        // Spawn multiple threads that all use stack swapping simultaneously
        for _ in 0..4 {
            let barrier_clone = Arc::clone(&barrier);
            let handle = thread::spawn(move || {
                barrier_clone.wait(); // Synchronize thread start

                // Each thread should be able to use stack swapping independently
                // This would fail with regular static variables due to race conditions
                unsafe {
                    swap_to_heap(simple_callout, None, 1 << 28).unwrap();
                }

                true // Return success
            });
            handles.push(handle);
        }

        // Wait for all threads to complete
        for handle in handles {
            let result = handle.join().unwrap();
            assert!(result); // All threads should complete successfully
        }
    }
}