Quelle  lib.rs

/*
 * Copyright (c) 2025 Google Inc. All rights reserved
 *
 * Permission is hereby granted, free of charge, to any person obtaining
 * a copy of this software and associated documentation files
 * (the "Software"), to deal in the Software without restriction,
 * including without limitation the rights to use, copy, modify, merge,
 * publish, distribute, sublicense, and/or sell copies of the Software,
 * and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#![no_std]

use libfdt::{FdtError, FdtNode};
use rust_support::arch::clean_cache_range;
use rust_support::vmm::vaddr_to_paddr;
use rust_support::Error as LkError;
use rust_support::{init::lk_init_level, LK_INIT_HOOK, SMP_MAX_CPUS};
use smccc::{psci, Hvc, Smc};

mod sys {
    #![allow(unused)]
    #![allow(non_camel_case_types)]
    include!(env!("BINDGEN_INC_FILE"));
}

/// Mask of relevant affinity fields inside the MPIDR.
///
/// We ignore all bits not included in this mask.
const MPIDR_MASK: u64 = 0xff_00_ff_ff_ff;

/// Mapping from LK logical CPU index to masked MPIDR value.
///
/// We initialize every element to [`u64::MAX`] which represents
/// the uninitialized value. Since this value has all bits set,
/// and the mask has some cleared, no valid masked MPIDR
/// will ever be equal to this default.
static mut CPU_LOGICAL_MAP: [u64; SMP_MAX_CPUS as usize] = [u64::MAX; SMP_MAX_CPUS as usize];

/// Function to get the current CPU number from the MPIDR.
///
/// This is a naked function because it gets called from assembly
/// very early during boot, before we have a stack or MMU.
///
/// ABI: same as the assembly function in `start.S`:
/// * Stores the result in `x0`
/// * Clobbers `x9-x10`
/// * Preserves all other registers
/// * Returns `SMP_MAX_CPUS` if the CPU could not be found
#[cfg(target_arch = "aarch64")]
#[no_mangle]
#[unsafe(naked)]
extern "C" fn arm64_curr_cpu_num() {
    core::arch::naked_asm!(
        // Load and mask the MPIDR of the current CPU
        "mrs    x9, mpidr_el1",
        "ldr    x10, ={mpidr_mask}",
        "and    x9, x9, x10",

        // if CPU_LOGICAL_MAP[0] == u64::MAX, then store the current MPIDR
        // there. This should only happen during the first call to this
        // function when the primary CPU first boots.
        "adrp   x10, {cpu_logical_map}",
        "add    x10, x10, #:lo12:{cpu_logical_map}",
        "ldr    x0, [x10]",
        "cmn    x0, #1",
        "b.eq   .Lduring_boot",

        // Iterate through all elements of CPU_LOGICAL_MAP
        // Registers in this loop:
        // * `x0` - current CPU index
        // * `x9` - masked value of MPIDR
        // * `x10` - value of CPU_LOGICAL_MAP[x0]
        //
        // TODO: this runs every time `arch_curr_cpu_num` is called from
        // Trusty kernel code, which is a lot so make it sub-linear time.
        // There are a few possible approaches:
        // * Binary tree of MPIDR bits in a binary heap layout,
        //   simpler to implement for logarithmic time, or
        // * Some sort of perfect hashing from MPIDR to array index,
        //   would provide constant time lookup but harder to implement.
        "mov    x0, #{smp_max_cpus}",
        ".Lmap_loop:",
        "sub    x0, x0, #1",
        "adrp   x10, {cpu_logical_map}",
        "add    x10, x10, #:lo12:{cpu_logical_map}",
        "ldr    x10, [x10, x0, lsl #3]",
        "cmp    x9, x10",
        "b.eq   .Ldone",
        "cmp    x0, #0",
        "b.ne   .Lmap_loop",

        // CPU not found in array, return SMP_MAX_CPUS
        "mov    x0, #{smp_max_cpus}",
        "b      .Ldone",

        // CPU_LOGICAL_MAP[0] = MPIDR, then return 0
        ".Lduring_boot:",
        "str    x9, [x10]",
        "mov    x0, xzr",

        ".Ldone:",
        "ret",

        mpidr_mask = const MPIDR_MASK,
        cpu_logical_map = sym CPU_LOGICAL_MAP,
        smp_max_cpus = const SMP_MAX_CPUS,
    );
}

// TODO: implement arm_curr_cpu_num for arm32.

#[no_mangle]
extern "C" fn arch_cpu_num_to_gic_affinities(cpu_num: usize) -> crate::sys::arm_gic_affinities {
    assert!(cpu_num < SMP_MAX_CPUS as usize);
    let cpu_mpidr = unsafe { CPU_LOGICAL_MAP[cpu_num] };
    assert!(cpu_mpidr != u64::MAX);

    crate::sys::arm_gic_affinities {
        aff3: (cpu_mpidr >> 32) as u8,
        aff2: (cpu_mpidr >> 16) as u8,
        aff1: (cpu_mpidr >> 8) as u8,
        aff0: cpu_mpidr as u8,
    }
}

/// Get the MPIDR of a CPU in the device tree by reading the `reg` property.
///
/// Returns:
/// * `Ok(Some(id))` if a valid CPU node could be found,
/// * `Ok(None)` if this node does not represent a CPU, or
/// * `Err` in case of error.
fn get_fdt_cpu_mpidr(cpu: &FdtNode) -> libfdt::Result<Option<u64>> {
    if !cpu.name()?.to_bytes().starts_with(b"cpu") {
        return Ok(None);
    }
    if cpu.device_type()? != Some(c"cpu") {
        return Ok(None);
    }

    let mut cpu_reg_iter = cpu.reg()?.ok_or(FdtError::NotFound)?;
    let cpu_reg = cpu_reg_iter.next().ok_or(FdtError::NotFound)?;

    Ok(Some(cpu_reg.addr))
}

#[cfg(target_arch = "aarch64")]
fn setup_dynamic_idle_thread_stacks(cpu: usize) -> Result<(), LkError> {
    use core::mem::ManuallyDrop;
    use rust_support::vmm::VmmPageArray;

    // Log2 alignment of stack pages. Use the same 2^4=16 alignment
    // from `start.S`. This has no effect since the alignment is smaller
    // than the page size.
    const STACK_ALIGN_LOG2: u8 = 4;
    let stack = VmmPageArray::new(
        c"idle thread stack",
        sys::DEFAULT_STACK_SIZE as usize,
        STACK_ALIGN_LOG2,
        0,
    )
    .map(ManuallyDrop::new)
    .inspect_err(|e| log::error!("Failed to allocate thread stack: {e}"))?;
    let stack_end = stack.ptr() as u64 + sys::DEFAULT_STACK_SIZE as u64;

    // Shadow stack log2 alignment from `start.S`.
    const SHADOW_STACK_ALIGN_LOG2: u8 = 3;
    let shadow_stack = VmmPageArray::new(
        c"idle thread shadow stack",
        sys::DEFAULT_SHADOW_STACK_SIZE as usize,
        SHADOW_STACK_ALIGN_LOG2,
        0,
    )
    .map(ManuallyDrop::new)
    .inspect_err(|e| log::error!("Failed to allocate thread shadow stack: {e}"))?;

    extern "C" {
        static mut arm64_idle_thread_stack_ptrs: [u64; SMP_MAX_CPUS as usize];
        static mut arm64_idle_thread_shadow_stack_ptrs: [u64; SMP_MAX_CPUS as usize];
    }
    // SAFETY: The array is defined in start.S and has `8 * SMP_MAX_CPUS` bytes.
    unsafe {
        arm64_idle_thread_stack_ptrs[cpu] = stack_end;
    }
    // SAFETY: The array is defined in start.S and has `8 * SMP_MAX_CPUS` bytes.
    unsafe {
        arm64_idle_thread_shadow_stack_ptrs[cpu] = shadow_stack.ptr() as u64;
    }

    Ok(())
}

#[cfg(target_arch = "arm")]
fn setup_dynamic_idle_thread_stacks(_cpu: usize) -> Result<(), LkError> {
    // TODO: implement this in `start.S` for arm
    Ok(())
}

/// An init hook that boots the secondary CPUs using PSCI.
extern "C" fn psci_boot_secondary_cpus_hook(_level: core::ffi::c_uint) {
    let fdt = match dtb_service::get_dtb() {
        Ok(fdt) => fdt,
        Err(e) => {
            log::error!("Error parsing FDT: {e}");
            return;
        }
    };

    let psci_node = match fdt.root().subnode(c"psci") {
        Ok(Some(n)) => n,
        Ok(None) => {
            log::error!("/psci node missing from FDT");
            return;
        }
        Err(e) => {
            log::error!("Error parsing /psci node in FDT: {e}");
            return;
        }
    };

    // TODO: move this into a common library that
    // other modules (like arm_ffa) can use
    let call_cpu_on = match psci_node.getprop_str(c"method") {
        Ok(Some(m)) if m == c"hvc" => psci::cpu_on::<Hvc>,
        Ok(Some(m)) if m == c"smc" => psci::cpu_on::<Smc>,
        Ok(Some(m)) => {
            log::error!("Unexpected PSCI method: {m:?}");
            return;
        }
        Ok(None) => {
            log::error!("PSCI method not found");
            return;
        }
        Err(e) => {
            log::error!("Error parsing PSCI method: {e}");
            return;
        }
    };
    // TODO: parse the CPU_ON SMC FID from the device tree;
    // this is optional for PSCI 0.2 and later, but we could still
    // parse and use the provided value for backward compatibility.

    let cpus_node = match fdt.root().subnode(c"cpus") {
        Ok(Some(n)) => n,
        Ok(None) => {
            log::error!("/cpus node missing from FDT");
            return;
        }
        Err(e) => {
            log::error!("Error parsing /cpus node in FDT: {e}");
            return;
        }
    };

    let cpus_subnodes = match cpus_node.subnodes() {
        Ok(n) => n,
        Err(e) => {
            log::error!("Failed to get /cpus subnodes in FDT: {e}");
            return;
        }
    };

    let start_ptr = &raw mut crate::sys::_start;
    // SAFETY: This function immutably converts a virtual address
    // to the corresponding physical address. It does not dereference
    // the given pointer, instead using it to look up the address
    // in the page tables.
    let start_paddr = unsafe { vaddr_to_paddr(start_ptr.cast()) };

    let mut found_cpu0 = false;
    // SAFETY: We write to `CPU_LOGICAL_MAP` during boot on the primary CPU
    // (either in this function or above in `arm64_curr_cpu_num`).
    let cpu0_mpidr = unsafe { CPU_LOGICAL_MAP[0] };
    assert!(cpu0_mpidr != u64::MAX);

    let mut next_cpu = 1;
    for cpu in cpus_subnodes {
        let fdt_cpu_mpidr = match get_fdt_cpu_mpidr(&cpu) {
            Ok(Some(id)) => id,
            Ok(None) => continue,
            Err(e) => {
                log::error!("Error reading CPU MPIDR: {e}");
                continue;
            }
        };

        let masked_cpu_mpidr = fdt_cpu_mpidr & MPIDR_MASK;
        if masked_cpu_mpidr == cpu0_mpidr {
            found_cpu0 = true;
            continue;
        }
        if next_cpu >= SMP_MAX_CPUS as usize {
            log::error!("Too many CPUs in device tree");
            if found_cpu0 {
                break;
            } else {
                // We still need to keep looking for CPU 0
                continue;
            }
        }

        if let Err(e) = setup_dynamic_idle_thread_stacks(next_cpu) {
            log::error!("Failed to set up dynamic idle thread stacks: {e}");
            continue;
        }

        // SAFETY: Elements `CPU_LOGICAL_MAP[1..]` are only written
        // by this function during its one initialization call,
        // and should never change after that. All secondary CPUs
        // should only read from that array using `arm64_curr_cpu_num`.
        unsafe {
            CPU_LOGICAL_MAP[next_cpu] = masked_cpu_mpidr;
            // Flush out the cache line of the map entry.
            // The callee includes a DMB SY barrier operation.
            //
            // We perform both the cache flush and barrier to ensure
            // that the value is written to memory before the target CPU starts
            // booting. The reason is that Trusty startup code reads this value
            // before the MMU has been enabled, and memory accesses at that
            // time are Device-nGnRnE (not in the cache).
            clean_cache_range(core::slice::from_ref(&CPU_LOGICAL_MAP[next_cpu]));
        };

        log::info!("Booting CPU {next_cpu}...");
        if let Err(e) = call_cpu_on(fdt_cpu_mpidr, start_paddr as u64, 0) {
            log::error!("PSCI error booting CPU {next_cpu}: {e}");
        }
        next_cpu += 1;
    }

    if !found_cpu0 {
        panic!("Device tree does not contain MPIDR for CPU 0: {cpu0_mpidr:#x}");
    }
}

LK_INIT_HOOK!(
    psci_boot_secondary_cpus,
    psci_boot_secondary_cpus_hook,
    lk_init_level::LK_INIT_LEVEL_ARCH
);