preemptive_threads/security/

aslr.rs

//! Address Space Layout Randomization (ASLR) implementation.

use crate::errors::ThreadError;
use crate::security::{SecurityConfig, crypto_rng::secure_random_u64};
use portable_atomic::{AtomicU64, AtomicUsize, Ordering};
use alloc::{vec, vec::Vec};

/// ASLR implementation for thread stacks and memory layout.
pub struct AslrManager {
    /// Randomization statistics
    randomizations_applied: AtomicUsize,
    /// Entropy used for randomization
    entropy_consumed: AtomicU64,
    /// Virtual address space layout
    address_space_layout: AddressSpaceLayout,
}

impl AslrManager {
    pub fn new() -> Self {
        Self {
            randomizations_applied: AtomicUsize::new(0),
            entropy_consumed: AtomicU64::new(0),
            address_space_layout: AddressSpaceLayout::detect(),
        }
    }
    
    /// Randomize stack allocation address.
    pub fn randomize_stack_address(
        &self,
        base_address: usize,
        stack_size: usize,
    ) -> Result<usize, ThreadError> {
        let entropy_bits = self.address_space_layout.available_entropy_bits();
        if entropy_bits == 0 {
            return Ok(base_address); // No randomization available
        }
        
        // Generate random offset within available address space
        let max_offset = self.address_space_layout.max_stack_offset(stack_size);
        let random_offset = self.generate_random_offset(max_offset)?;
        
        // Apply offset with proper alignment
        let page_size = self.address_space_layout.page_size;
        let aligned_offset = (random_offset / page_size) * page_size;
        
        let randomized_address = base_address.wrapping_add(aligned_offset);
        
        // Verify address is in valid range
        if !self.address_space_layout.is_valid_stack_address(randomized_address, stack_size) {
            return Err(ThreadError::MemoryError());
        }
        
        self.randomizations_applied.fetch_add(1, Ordering::Relaxed);
        Ok(randomized_address)
    }
    
    /// Randomize heap allocation address.
    pub fn randomize_heap_address(
        &self,
        base_address: usize,
        allocation_size: usize,
    ) -> Result<usize, ThreadError> {
        let max_offset = self.address_space_layout.max_heap_offset(allocation_size);
        let random_offset = self.generate_random_offset(max_offset)?;
        
        // Apply alignment requirements
        let align_size = core::mem::align_of::<usize>();
        let aligned_offset = (random_offset / align_size) * align_size;
        
        let randomized_address = base_address.wrapping_add(aligned_offset);
        
        if !self.address_space_layout.is_valid_heap_address(randomized_address, allocation_size) {
            return Err(ThreadError::MemoryError());
        }
        
        self.randomizations_applied.fetch_add(1, Ordering::Relaxed);
        Ok(randomized_address)
    }
    
    /// Generate guard gap between memory regions.
    pub fn generate_guard_gap(&self) -> Result<usize, ThreadError> {
        let min_gap = self.address_space_layout.page_size;
        let max_gap = min_gap * 16; // Up to 16 pages
        
        let random_pages = self.generate_random_offset(16)? + 1;
        Ok(random_pages * min_gap)
    }
    
    /// Generate random memory layout for thread.
    pub fn generate_thread_layout(&self) -> Result<ThreadMemoryLayout, ThreadError> {
        let stack_base = self.address_space_layout.stack_base_address();
        let heap_base = self.address_space_layout.heap_base_address();
        
        // Randomize stack position
        let stack_size = 1024 * 1024; // 1MB default stack
        let randomized_stack = self.randomize_stack_address(stack_base, stack_size)?;
        
        // Randomize heap position
        let heap_size = 64 * 1024 * 1024; // 64MB heap region
        let randomized_heap = self.randomize_heap_address(heap_base, heap_size)?;
        
        // Add guard gaps
        let stack_guard_gap = self.generate_guard_gap()?;
        let heap_guard_gap = self.generate_guard_gap()?;
        
        Ok(ThreadMemoryLayout {
            stack_base: randomized_stack,
            stack_size,
            stack_guard_gap,
            heap_base: randomized_heap,
            heap_size,
            heap_guard_gap,
            randomization_entropy: self.entropy_consumed.load(Ordering::Relaxed),
        })
    }
    
    /// Generate cryptographically secure random offset.
    fn generate_random_offset(&self, max_offset: usize) -> Result<usize, ThreadError> {
        if max_offset == 0 {
            return Ok(0);
        }
        
        let random_value = secure_random_u64()?;
        self.entropy_consumed.fetch_add(8, Ordering::Relaxed);
        
        Ok((random_value as usize) % max_offset)
    }
}

/// Address space layout detection and management.
#[derive(Debug, Clone)]
pub struct AddressSpaceLayout {
    /// Page size for alignment
    pub page_size: usize,
    /// Available address space size
    pub address_space_size: usize,
    /// Stack region boundaries
    pub stack_region: MemoryRegion,
    /// Heap region boundaries
    pub heap_region: MemoryRegion,
    /// Architecture-specific constraints
    pub arch_constraints: ArchConstraints,
}

impl AddressSpaceLayout {
    /// Detect current system address space layout.
    pub fn detect() -> Self {
        Self {
            page_size: detect_page_size(),
            address_space_size: detect_address_space_size(),
            stack_region: detect_stack_region(),
            heap_region: detect_heap_region(),
            arch_constraints: ArchConstraints::detect(),
        }
    }
    
    /// Get available entropy bits for randomization.
    pub fn available_entropy_bits(&self) -> u32 {
        #[cfg(target_pointer_width = "64")]
        {
            // 64-bit systems typically have more address space for randomization
            match self.arch_constraints.arch {
                Architecture::X86_64 => 28, // ~256GB randomization space
                Architecture::Aarch64 => 32, // ~4TB randomization space
                Architecture::Riscv64 => 30, // ~1TB randomization space
                _ => 20, // Conservative default
            }
        }
        
        #[cfg(target_pointer_width = "32")]
        {
            // 32-bit systems have limited address space
            16 // ~64MB randomization space
        }
    }
    
    /// Get maximum stack offset for randomization.
    pub fn max_stack_offset(&self, stack_size: usize) -> usize {
        let available_space = self.stack_region.size.saturating_sub(stack_size);
        available_space / 2 // Use half available space for randomization
    }
    
    /// Get maximum heap offset for randomization.
    pub fn max_heap_offset(&self, heap_size: usize) -> usize {
        let available_space = self.heap_region.size.saturating_sub(heap_size);
        available_space / 4 // Use quarter available space for randomization
    }
    
    /// Check if stack address is valid.
    pub fn is_valid_stack_address(&self, address: usize, size: usize) -> bool {
        address >= self.stack_region.start &&
        address + size <= self.stack_region.start + self.stack_region.size
    }
    
    /// Check if heap address is valid.
    pub fn is_valid_heap_address(&self, address: usize, size: usize) -> bool {
        address >= self.heap_region.start &&
        address + size <= self.heap_region.start + self.heap_region.size
    }
    
    /// Get stack base address for allocation.
    pub fn stack_base_address(&self) -> usize {
        self.stack_region.start + (self.stack_region.size / 4)
    }
    
    /// Get heap base address for allocation.
    pub fn heap_base_address(&self) -> usize {
        self.heap_region.start + (self.heap_region.size / 8)
    }
}

/// Memory region descriptor.
#[derive(Debug, Clone)]
pub struct MemoryRegion {
    pub start: usize,
    pub size: usize,
}

/// Architecture-specific constraints for ASLR.
#[derive(Debug, Clone)]
pub struct ArchConstraints {
    pub arch: Architecture,
    pub min_alignment: usize,
    pub max_randomization: usize,
    pub forbidden_ranges: Vec<MemoryRegion>,
}

impl ArchConstraints {
    pub fn detect() -> Self {
        #[cfg(feature = "x86_64")]
        {
            return Self {
                arch: Architecture::X86_64,
                min_alignment: 4096, // Page aligned
                max_randomization: 1 << 28, // 256MB
                forbidden_ranges: vec![
                    // Kernel space
                    MemoryRegion { start: 0xFFFF800000000000, size: usize::MAX },
                ],
            };
        }
        
        #[cfg(feature = "arm64")]
        {
            Self {
                arch: Architecture::Aarch64,
                min_alignment: 4096,
                max_randomization: 1 << 32, // 4GB
                forbidden_ranges: vec![
                    // Kernel space (simplified)
                    MemoryRegion { start: 0xFFFF000000000000, size: usize::MAX },
                ],
            }
        }
        
        #[cfg(not(any(feature = "x86_64", feature = "arm64")))]
        {
            Self {
                arch: Architecture::Generic,
                min_alignment: core::mem::size_of::<usize>(),
                max_randomization: 1 << 20, // 1MB conservative
                forbidden_ranges: Vec::new(),
            }
        }
    }
}

/// Supported architectures for ASLR.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Architecture {
    X86_64,
    Aarch64,
    Riscv64,
    Generic,
}

/// Thread-specific memory layout with randomization.
#[derive(Debug, Clone)]
pub struct ThreadMemoryLayout {
    pub stack_base: usize,
    pub stack_size: usize,
    pub stack_guard_gap: usize,
    pub heap_base: usize,
    pub heap_size: usize,
    pub heap_guard_gap: usize,
    pub randomization_entropy: u64,
}

impl ThreadMemoryLayout {
    /// Create stack with randomized address.
    pub fn create_randomized_stack(&self) -> Result<RandomizedStack, ThreadError> {
        // In a real implementation, this would allocate memory at the randomized address
        Ok(RandomizedStack {
            layout: self.clone(),
            actual_address: self.stack_base,
            entropy_used: 64, // bits of entropy used
        })
    }
}

/// Stack with randomized allocation address.
pub struct RandomizedStack {
    pub layout: ThreadMemoryLayout,
    pub actual_address: usize,
    pub entropy_used: u32,
}

impl RandomizedStack {
    /// Get stack bounds with guard gaps.
    pub fn bounds(&self) -> (usize, usize) {
        let start = self.actual_address + self.layout.stack_guard_gap;
        let end = start + self.layout.stack_size - self.layout.stack_guard_gap;
        (start, end)
    }
    
    /// Check if address is within this stack.
    pub fn contains(&self, address: usize) -> bool {
        let (start, end) = self.bounds();
        address >= start && address < end
    }
}

/// Global ASLR manager instance.
static mut ASLR_MANAGER: Option<AslrManager> = None;

/// System detection functions.

fn detect_page_size() -> usize {
    #[cfg(target_os = "linux")]
    {
        // In real implementation, would use sysconf(_SC_PAGESIZE)
        4096
    }
    
    #[cfg(not(target_os = "linux"))]
    {
        4096 // Common page size
    }
}

fn detect_address_space_size() -> usize {
    #[cfg(target_pointer_width = "64")]
    {
        // 48-bit address space is common on x86_64
        1usize << 48
    }
    
    #[cfg(target_pointer_width = "32")]
    {
        1usize << 32
    }
}

fn detect_stack_region() -> MemoryRegion {
    #[cfg(target_pointer_width = "64")]
    {
        MemoryRegion {
            start: 0x7F0000000000, // Typical user stack region on Linux x86_64
            size: 0x10000000000,   // 1TB region
        }
    }
    
    #[cfg(target_pointer_width = "32")]
    {
        MemoryRegion {
            start: 0xC0000000,  // 3GB mark
            size: 0x40000000,   // 1GB region
        }
    }
}

fn detect_heap_region() -> MemoryRegion {
    #[cfg(target_pointer_width = "64")]
    {
        MemoryRegion {
            start: 0x100000000,    // 4GB mark
            size: 0x7EF000000000,  // Large heap region
        }
    }
    
    #[cfg(target_pointer_width = "32")]
    {
        MemoryRegion {
            start: 0x08000000,  // 128MB mark
            size: 0xB8000000,   // ~3GB region
        }
    }
}

/// ASLR statistics.
#[derive(Debug, Clone)]
pub struct AslrStats {
    pub randomizations_applied: usize,
    pub entropy_consumed: u64,
    pub entropy_bits_available: u32,
    pub aslr_enabled: bool,
}

/// Initialize ASLR subsystem.
pub fn init_aslr(_config: SecurityConfig) -> Result<(), ThreadError> {
    unsafe {
        ASLR_MANAGER = Some(AslrManager::new());
    }
    
    let stats = get_aslr_stats();
    // ASLR initialized with entropy bits available from stats
    
    Ok(())
}

/// Create randomized thread memory layout.
pub fn create_randomized_layout() -> Result<ThreadMemoryLayout, ThreadError> {
    unsafe {
        match &ASLR_MANAGER {
            Some(manager) => manager.generate_thread_layout(),
            None => Err(ThreadError::InvalidState()),
        }
    }
}

/// Randomize stack allocation address.
pub fn randomize_stack_address(base: usize, size: usize) -> Result<usize, ThreadError> {
    unsafe {
        match &ASLR_MANAGER {
            Some(manager) => manager.randomize_stack_address(base, size),
            None => Ok(base), // No randomization if not initialized
        }
    }
}

/// Get ASLR statistics.
pub fn get_aslr_stats() -> AslrStats {
    unsafe {
        match &ASLR_MANAGER {
            Some(manager) => AslrStats {
                randomizations_applied: manager.randomizations_applied.load(Ordering::Relaxed),
                entropy_consumed: manager.entropy_consumed.load(Ordering::Relaxed),
                entropy_bits_available: manager.address_space_layout.available_entropy_bits(),
                aslr_enabled: true,
            },
            None => AslrStats {
                randomizations_applied: 0,
                entropy_consumed: 0,
                entropy_bits_available: 0,
                aslr_enabled: false,
            },
        }
    }
}