ruvix_boot/

stages.rs

//! Boot stage implementations (ADR-087 Section 9.1).
//!
//! The kernel boot follows a five-stage process:
//!
//! | Stage | Name | Description |
//! |-------|------|-------------|
//! | **0** | Hardware Init | Platform-specific initialization (mocked in Phase A) |
//! | **1** | RVF Verify | Parse manifest + ML-DSA-65 signature verification |
//! | **2** | Object Create | Create root task, regions, queues, witness log |
//! | **3** | Component Mount | Mount components + distribute capabilities |
//! | **4** | First Attestation | Boot attestation to witness log |

use crate::manifest::RvfManifest;
use crate::signature::SignatureVerifier;
use crate::witness_log::{WitnessLog, WitnessLogConfig};
use crate::capability_distribution::CapabilityDistribution;
use crate::attestation::BootAttestation;
use ruvix_types::{KernelError, RegionHandle, TaskHandle};
use ruvix_cap::BootCapabilitySet;

/// Stage 0: Hardware initialization.
///
/// In Phase A (Linux-hosted), this is a no-op.
/// In Phase B (bare metal), this would:
/// - Initialize CPU/MMU
/// - Set up interrupt vectors
/// - Initialize memory controller
/// - Enable UART for diagnostics
#[derive(Debug, Clone)]
pub struct Stage0Hardware {
    /// Whether hardware initialization completed.
    pub initialized: bool,

    /// Platform identifier.
    pub platform_id: u64,

    /// Total physical memory in bytes.
    pub physical_memory_bytes: u64,

    /// CPU frequency in Hz.
    pub cpu_frequency_hz: u64,
}

impl Stage0Hardware {
    /// Creates a new Stage 0 handler.
    #[must_use]
    pub fn new() -> Self {
        Self {
            initialized: false,
            platform_id: 0,
            physical_memory_bytes: 0,
            cpu_frequency_hz: 0,
        }
    }

    /// Executes Stage 0 hardware initialization.
    ///
    /// # Phase A (Linux-hosted)
    ///
    /// Returns mock hardware information.
    ///
    /// # Phase B (bare metal)
    ///
    /// Performs actual hardware initialization.
    pub fn execute(&mut self) -> Result<(), KernelError> {
        #[cfg(feature = "verbose")]
        eprintln!("Stage 0: Hardware initialization (Phase A mock)");

        // Phase A: Mock hardware info
        self.platform_id = 0x5255_5649_585F_4131; // "RUVIX_A1"
        self.physical_memory_bytes = 1024 * 1024 * 1024; // 1 GiB mock
        self.cpu_frequency_hz = 1_000_000_000; // 1 GHz mock

        self.initialized = true;

        #[cfg(feature = "verbose")]
        eprintln!(
            "  Platform: 0x{:016X}, Memory: {} MiB, CPU: {} MHz",
            self.platform_id,
            self.physical_memory_bytes / (1024 * 1024),
            self.cpu_frequency_hz / 1_000_000
        );

        Ok(())
    }
}

impl Default for Stage0Hardware {
    fn default() -> Self {
        Self::new()
    }
}

/// Stage 1: RVF manifest parse + signature verification.
///
/// **SECURITY CRITICAL (SEC-001)**: Signature verification failure
/// causes immediate PANIC with no fallback boot path.
pub struct Stage1Verify {
    /// The verified manifest (set after successful verification).
    pub manifest: Option<RvfManifest>,

    /// Signature verifier.
    verifier: Option<SignatureVerifier>,
}

impl Stage1Verify {
    /// Creates a new Stage 1 handler.
    #[must_use]
    pub fn new() -> Self {
        Self {
            manifest: None,
            verifier: None,
        }
    }

    /// Sets the public key for signature verification.
    pub fn set_public_key(&mut self, public_key: &[u8]) {
        self.verifier = Some(SignatureVerifier::new(public_key));
    }

    /// Executes Stage 1 manifest verification.
    ///
    /// # Panics
    ///
    /// Panics if signature verification fails (SEC-001).
    ///
    /// # Errors
    ///
    /// Returns `KernelError::InvalidManifest` if manifest parsing fails.
    pub fn execute(&mut self, manifest_bytes: &[u8], signature: &[u8]) -> Result<(), KernelError> {
        #[cfg(feature = "verbose")]
        eprintln!("Stage 1: RVF manifest verification");

        // Verify signature (panics on failure per SEC-001)
        let verifier = self.verifier.as_ref().ok_or(KernelError::InternalError)?;
        verifier.verify_boot_signature(manifest_bytes, signature);

        #[cfg(feature = "verbose")]
        eprintln!("  Signature verified (ML-DSA-65)");

        // Parse manifest
        let manifest = RvfManifest::parse(manifest_bytes)?;

        #[cfg(feature = "verbose")]
        eprintln!(
            "  Manifest parsed: {} components, {} regions",
            manifest.component_graph.component_count(),
            manifest.memory_schema.region_count()
        );

        // Validate manifest
        if !manifest.validate() {
            return Err(KernelError::InvalidManifest);
        }

        #[cfg(feature = "verbose")]
        eprintln!("  Manifest validated");

        self.manifest = Some(manifest);
        Ok(())
    }
}

impl Default for Stage1Verify {
    fn default() -> Self {
        Self::new()
    }
}

/// Stage 2: Kernel object creation.
///
/// Creates:
/// - Root task
/// - Memory regions per manifest schema
/// - Queues per manifest wiring
/// - Witness log region
pub struct Stage2Create {
    /// Root task handle.
    pub root_task: Option<TaskHandle>,

    /// Created region handles.
    pub regions: [Option<RegionHandle>; 256],
    /// Number of active regions in the array.
    pub region_count: usize,

    /// Witness log handle.
    pub witness_log: Option<WitnessLog>,

    /// Initial capability set for root task.
    pub boot_capabilities: Option<BootCapabilitySet>,
}

impl Stage2Create {
    /// Creates a new Stage 2 handler.
    #[must_use]
    pub fn new() -> Self {
        Self {
            root_task: None,
            regions: [None; 256],
            region_count: 0,
            witness_log: None,
            boot_capabilities: None,
        }
    }

    /// Executes Stage 2 kernel object creation.
    ///
    /// # Errors
    ///
    /// Returns an error if object creation fails (e.g., out of memory).
    pub fn execute(
        &mut self,
        manifest: &RvfManifest,
        _physical_memory_bytes: u64,
    ) -> Result<(), KernelError> {
        #[cfg(feature = "verbose")]
        eprintln!("Stage 2: Kernel object creation");

        // Create root task
        let root_task_id = 1u32;
        self.root_task = Some(TaskHandle::new(root_task_id, 0));

        #[cfg(feature = "verbose")]
        eprintln!("  Created root task: {:?}", self.root_task);

        // Create regions per manifest
        self.create_regions(manifest)?;

        // Create witness log
        let witness_config = WitnessLogConfig::from_policy(&manifest.witness_log_policy);
        self.witness_log = Some(WitnessLog::new(witness_config));

        #[cfg(feature = "verbose")]
        eprintln!("  Created witness log");

        // Create initial capability set
        let mut boot_caps = BootCapabilitySet::new();
        boot_caps.set_root_task(root_task_id as u64);
        boot_caps.add_memory_region(0x1000, 0); // Physical memory
        boot_caps.set_witness_log(0x2000);
        boot_caps.set_timer(0x3000);
        boot_caps.set_interrupt_queue(0x4000, u64::MAX);

        // Add RVF package capability
        boot_caps.set_rvf_package(
            0x5000,
            u64::from_le_bytes(manifest.content_hash[0..8].try_into().unwrap()),
        );

        #[cfg(feature = "verbose")]
        eprintln!("  Created {} initial capabilities", boot_caps.total_count());

        self.boot_capabilities = Some(boot_caps);

        Ok(())
    }

    fn create_regions(&mut self, manifest: &RvfManifest) -> Result<(), KernelError> {
        let count = manifest.memory_schema.region_count();

        #[cfg(feature = "verbose")]
        eprintln!("  Creating {} regions from manifest", count);

        // Phase A: Create mock region handles
        for i in 0..count {
            if i >= 256 {
                return Err(KernelError::LimitExceeded);
            }

            self.regions[i] = Some(RegionHandle::new(0x1000 + i as u32, 0));
            self.region_count += 1;
        }

        Ok(())
    }
}

impl Default for Stage2Create {
    fn default() -> Self {
        Self::new()
    }
}

/// Stage 3: Component mount + capability distribution.
///
/// This stage:
/// - Mounts WASM components
/// - Distributes capabilities per manifest
/// - Connects queues per wiring
/// - Spawns initial tasks per WIT entry points
/// - **Drops root task to minimum capability set (SEC-001)**
pub struct Stage3Mount {
    /// Number of components mounted.
    pub components_mounted: usize,

    /// Number of queues connected.
    pub queues_connected: usize,

    /// Number of tasks spawned.
    pub tasks_spawned: usize,

    /// Capability distribution result.
    pub capability_distribution: Option<CapabilityDistribution>,
}

impl Stage3Mount {
    /// Creates a new Stage 3 handler.
    #[must_use]
    pub fn new() -> Self {
        Self {
            components_mounted: 0,
            queues_connected: 0,
            tasks_spawned: 0,
            capability_distribution: None,
        }
    }

    /// Executes Stage 3 component mounting.
    ///
    /// # SEC-001 Compliance
    ///
    /// After this stage, the root task's capability set is reduced to
    /// the minimum required for operation (no longer has full physical
    /// memory access).
    pub fn execute(
        &mut self,
        manifest: &RvfManifest,
        boot_capabilities: &BootCapabilitySet,
    ) -> Result<(), KernelError> {
        #[cfg(feature = "verbose")]
        eprintln!("Stage 3: Component mount + capability distribution");

        // Mount components
        self.mount_components(manifest)?;

        // Connect queues
        self.connect_queues(manifest)?;

        // Distribute capabilities
        self.distribute_capabilities(manifest, boot_capabilities)?;

        // SEC-001: Drop root task to minimum capability set
        self.drop_root_capabilities()?;

        Ok(())
    }

    fn mount_components(&mut self, manifest: &RvfManifest) -> Result<(), KernelError> {
        let count = manifest.component_graph.component_count();

        #[cfg(feature = "verbose")]
        eprintln!("  Mounting {} components", count);

        // Phase A: Mock component mounting
        self.components_mounted = count;

        Ok(())
    }

    fn connect_queues(&mut self, manifest: &RvfManifest) -> Result<(), KernelError> {
        let count = manifest.component_graph.wiring_count();

        #[cfg(feature = "verbose")]
        eprintln!("  Connecting {} queues", count);

        // Phase A: Mock queue connection
        self.queues_connected = count;

        Ok(())
    }

    fn distribute_capabilities(
        &mut self,
        manifest: &RvfManifest,
        boot_capabilities: &BootCapabilitySet,
    ) -> Result<(), KernelError> {
        #[cfg(feature = "verbose")]
        eprintln!("  Distributing capabilities per manifest");

        let distribution = CapabilityDistribution::from_manifest(manifest, boot_capabilities)?;
        self.capability_distribution = Some(distribution);

        Ok(())
    }

    fn drop_root_capabilities(&mut self) -> Result<(), KernelError> {
        #[cfg(feature = "verbose")]
        eprintln!("  SEC-001: Dropping root task to minimum capability set");

        // In a real implementation, this would:
        // 1. Identify the minimum capability set needed by root task
        // 2. Revoke all other capabilities from root task
        // 3. Ensure root task cannot re-escalate

        // Phase A: Just mark that we did it
        if let Some(ref mut dist) = self.capability_distribution {
            dist.root_dropped_to_minimum = true;
        }

        Ok(())
    }
}

impl Default for Stage3Mount {
    fn default() -> Self {
        Self::new()
    }
}

/// Stage 4: First attestation (boot attestation to witness log).
///
/// Records the initial boot attestation containing:
/// - RVF package hash
/// - Capability table hash
/// - Region layout hash
/// - Timestamp
pub struct Stage4Attest {
    /// Boot attestation entry.
    pub attestation: Option<BootAttestation>,

    /// Whether attestation was written to witness log.
    pub attested: bool,
}

impl Stage4Attest {
    /// Creates a new Stage 4 handler.
    #[must_use]
    pub fn new() -> Self {
        Self {
            attestation: None,
            attested: false,
        }
    }

    /// Executes Stage 4 boot attestation.
    pub fn execute(
        &mut self,
        manifest: &RvfManifest,
        witness_log: &mut WitnessLog,
        boot_capabilities: &BootCapabilitySet,
    ) -> Result<(), KernelError> {
        #[cfg(feature = "verbose")]
        eprintln!("Stage 4: First attestation");

        // Create boot attestation
        let attestation = BootAttestation::new(
            manifest.content_hash,
            Self::hash_capability_table(boot_capabilities),
            Self::hash_region_layout(manifest),
            Self::get_timestamp(),
        );

        #[cfg(feature = "verbose")]
        eprintln!("  Boot attestation created");

        // Write to witness log
        witness_log.append_boot_attestation(&attestation)?;

        #[cfg(feature = "verbose")]
        eprintln!("  Boot attestation written to witness log");

        self.attestation = Some(attestation);
        self.attested = true;

        Ok(())
    }

    fn hash_capability_table(caps: &BootCapabilitySet) -> [u8; 32] {
        use sha2::{Sha256, Digest};

        let mut hasher = Sha256::new();

        for cap in caps.iter() {
            hasher.update(&cap.object_id.to_le_bytes());
            hasher.update(&[cap.object_type as u8]);
            hasher.update(&cap.rights.bits().to_le_bytes());
            hasher.update(&cap.badge.to_le_bytes());
        }

        let result = hasher.finalize();
        let mut hash = [0u8; 32];
        hash.copy_from_slice(&result);
        hash
    }

    fn hash_region_layout(manifest: &RvfManifest) -> [u8; 32] {
        use sha2::{Sha256, Digest};

        let mut hasher = Sha256::new();

        hasher.update(&manifest.memory_schema.total_memory_required.to_le_bytes());
        hasher.update(&(manifest.memory_schema.region_count() as u32).to_le_bytes());

        let result = hasher.finalize();
        let mut hash = [0u8; 32];
        hash.copy_from_slice(&result);
        hash
    }

    fn get_timestamp() -> u64 {
        #[cfg(feature = "std")]
        {
            use std::time::{SystemTime, UNIX_EPOCH};
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .map(|d| d.as_nanos() as u64)
                .unwrap_or(0)
        }
        #[cfg(not(feature = "std"))]
        {
            0 // No timestamp in no_std without a timer
        }
    }
}

impl Default for Stage4Attest {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_stage0_hardware_init() {
        let mut stage = Stage0Hardware::new();
        assert!(!stage.initialized);

        stage.execute().unwrap();

        assert!(stage.initialized);
        assert!(stage.physical_memory_bytes > 0);
        assert!(stage.cpu_frequency_hz > 0);
    }

    #[test]
    fn test_stage1_needs_public_key() {
        let mut stage = Stage1Verify::new();
        let manifest = b"test";
        let signature = [0u8; 3309];

        // Should fail without public key
        let result = stage.execute(manifest, &signature);
        assert!(result.is_err());
    }

    #[test]
    fn test_stage2_creates_objects() {
        let mut stage = Stage2Create::new();

        // Create minimal valid manifest
        let mut manifest_bytes = vec![0u8; 100];
        manifest_bytes[0..4].copy_from_slice(b"RVF1");
        manifest_bytes[4..6].copy_from_slice(&1u16.to_le_bytes()); // major
        manifest_bytes[6..8].copy_from_slice(&0u16.to_le_bytes()); // minor

        let manifest = RvfManifest::parse(&manifest_bytes).unwrap();
        let physical_memory = 1024 * 1024 * 1024;

        stage.execute(&manifest, physical_memory).unwrap();

        assert!(stage.root_task.is_some());
        assert!(stage.witness_log.is_some());
        assert!(stage.boot_capabilities.is_some());
    }

    #[test]
    fn test_stage3_sec001_capability_drop() {
        let mut stage = Stage3Mount::new();

        // Create minimal manifest
        let mut manifest_bytes = vec![0u8; 100];
        manifest_bytes[0..4].copy_from_slice(b"RVF1");
        manifest_bytes[4..6].copy_from_slice(&1u16.to_le_bytes());

        let manifest = RvfManifest::parse(&manifest_bytes).unwrap();
        let boot_caps = BootCapabilitySet::minimal(1);

        stage.execute(&manifest, &boot_caps).unwrap();

        // Verify SEC-001: root capability drop occurred
        assert!(stage.capability_distribution.is_some());
        assert!(stage.capability_distribution.as_ref().unwrap().root_dropped_to_minimum);
    }

    #[test]
    fn test_stage4_attestation() {
        let mut stage = Stage4Attest::new();

        // Create minimal manifest
        let mut manifest_bytes = vec![0u8; 100];
        manifest_bytes[0..4].copy_from_slice(b"RVF1");
        manifest_bytes[4..6].copy_from_slice(&1u16.to_le_bytes());

        let manifest = RvfManifest::parse(&manifest_bytes).unwrap();
        let boot_caps = BootCapabilitySet::minimal(1);
        let mut witness_log = WitnessLog::new(WitnessLogConfig::default());

        stage.execute(&manifest, &mut witness_log, &boot_caps).unwrap();

        assert!(stage.attested);
        assert!(stage.attestation.is_some());
    }
}