path: root/kernel/src/device_tree.rs

//! Support for the DeviceTree format.

use crate::collections::stack_linked_list::StackLinkedList;
use bstr::BStr;
use contracts::requires;
use core::{fmt, iter, slice};
use either::Either;

/// A reference to a flattened DeviceTree (DTB) in memory.
#[derive(Debug)]
pub struct FlattenedDeviceTree<'dt> {
    header: &'dt FdtHeader,
    struct_block: &'dt [u32],
    strings_block: &'dt BStr,
    memrsv_block: &'dt [FdtMemRsv],
}

impl<'dt> FlattenedDeviceTree<'dt> {
    /// Looks for a DeviceTree at the given address, and returns it if it looks valid.
    ///
    /// # Safety
    ///
    /// - `ptr` must point to a spec-compliant flattened DeviceTree.
    /// - The memory contained by the flattened DeviceTree must be valid for the duration of
    ///   lifetime `'dt`.
    /// - The memory contained by the flattened DeviceTree must not be mutated for the duration of
    ///   lifetime `'dt`.
    #[requires((ptr as *const FdtHeader).is_aligned())]
    pub unsafe fn from_ptr(ptr: *const u8) -> Result<FlattenedDeviceTree<'dt>, DeviceTreeError> {
        // Check that the header appears to point to a valid flattened DeviceTree.
        let header: &'dt FdtHeader = &*(ptr as *const FdtHeader);
        let magic = u32::from_be(header.magic);
        if magic != 0xd00dfeed {
            return Err(DeviceTreeError::BadMagic(magic));
        }
        let version = u32::from_be(header.version);
        let last_comp_version = u32::from_be(header.last_comp_version);
        if last_comp_version > 17 {
            return Err(DeviceTreeError::IncompatibleVersion(
                version,
                last_comp_version,
            ));
        }

        // Get pointers to each block.
        let off_dt_struct = u32::from_be(header.off_dt_struct) as usize;
        let size_dt_struct = u32::from_be(header.size_dt_struct) as usize;
        let off_dt_strings = u32::from_be(header.off_dt_strings) as usize;
        let size_dt_strings = u32::from_be(header.size_dt_strings) as usize;
        let off_mem_rsvmap = u32::from_be(header.off_mem_rsvmap) as usize;

        // Check that the structure block has an aligned size.
        if (size_dt_struct & 0b11) != 0 {
            return Err(DeviceTreeError::InvalidDeviceTree);
        }

        // Extract the structure and strings blocks.
        let struct_block: &[u32] =
            slice::from_raw_parts(ptr.add(off_dt_struct).cast(), size_dt_struct / 4);
        let strings_block = BStr::new(slice::from_raw_parts(
            ptr.add(off_dt_strings),
            size_dt_strings,
        ));

        // Read memory reservations until the terminating one is found, then construct the block of
        // the appropriate length.
        let mut memrsv_count = 0;
        let memrsv_ptr: *const FdtMemRsv = ptr.add(off_mem_rsvmap).cast();
        let memrsv_block = loop {
            let memrsv = *memrsv_ptr.add(memrsv_count);

            // We can skip the endian conversion, since we're just testing against zero.
            if memrsv == (FdtMemRsv { addr: 0, size: 0 }) {
                break slice::from_raw_parts(memrsv_ptr, memrsv_count);
            }

            memrsv_count += 1;
        };

        Ok(FlattenedDeviceTree {
            header,
            struct_block,
            strings_block,
            memrsv_block,
        })
    }

    /// Returns the string at the given offset of the strings block.
    pub fn get_string(&self, offset: u32) -> Result<&BStr, DeviceTreeError> {
        let out = &self.strings_block[offset as usize..];
        let i = out
            .iter()
            .position(|&b| b == b'\0')
            .ok_or(DeviceTreeError::StringMissingNulTerminator(offset))?;
        Ok(&out[..i])
    }

    /// Iterates over the properties stored in the DeviceTree, only allocating memory on the stack.
    pub fn for_each_property<'iter: 'dt, E>(
        &'iter self,
        mut func: impl for<'a> FnMut(FdtNodePath<'dt, 'a>, &'dt BStr, &'dt BStr) -> Result<(), E>,
    ) -> Result<(), Either<DeviceTreeError, E>> {
        for_each_property(
            &mut self.struct_events().peekable(),
            &mut func,
            StackLinkedList::NIL,
        )
    }

    /// Returns an iterator over the events in the flattened DeviceTree's structure block.
    pub fn struct_events<'iter: 'dt>(
        &'iter self,
    ) -> impl 'iter + Iterator<Item = Result<FdtStructEvent<'dt>, DeviceTreeError>> {
        let mut ptr = self.struct_block;
        iter::from_fn(move || loop {
            break match u32::from_be(ptr[0]) {
                // FDT_BEGIN_NODE
                0x00000001 => {
                    // Save a pointer to the start of the extra data.
                    let name_ptr = (&ptr[1]) as *const u32 as *const u8;

                    // Look for a null terminator.
                    //
                    // SAFETY: This method can only be called when the FlattenedDeviceTree was
                    // constructed from a valid flattened DeviceTree.
                    let mut name_len = 0;
                    while unsafe { *name_ptr.add(name_len) } != b'\0' {
                        name_len += 1;
                    }

                    // Create the name as a BStr.
                    //
                    // SAFETY: Well, we already accessed this memory above when finding the null
                    // terminator... But yes, this being valid is guaranteed by this being a
                    // flattened DeviceTree.
                    let name = BStr::new(unsafe { slice::from_raw_parts(name_ptr, name_len) });

                    // Advance the pointer.
                    let extra_data_count = (name_len + 4) >> 2;
                    ptr = &ptr[1 + extra_data_count..];

                    // Return the event.
                    Some(Ok(FdtStructEvent::BeginNode(name)))
                }
                // FDT_END_NODE
                0x00000002 => {
                    // Advance the pointer.
                    ptr = &ptr[1..];

                    // Return the event.
                    Some(Ok(FdtStructEvent::EndNode))
                }
                // FDT_PROP
                0x00000003 => {
                    // Get the length of the property data and the offset of the name in the
                    // strings block.
                    let len = u32::from_be(ptr[1]) as usize;
                    let name_off = u32::from_be(ptr[2]);

                    // Get the property data as a BStr.
                    //
                    // SAFETY: This is valid by the requirements of a flattened DeviceTree.
                    let data = BStr::new(unsafe {
                        slice::from_raw_parts(&ptr[3] as *const u32 as *const u8, len)
                    });

                    // Get the property name as a BStr.
                    let name = match self.get_string(name_off) {
                        Ok(name) => name,
                        Err(err) => break Some(Err(err)),
                    };

                    // Advance the pointer.
                    let data_count = (len + 3) >> 2;
                    ptr = &ptr[3 + data_count..];

                    // Return the event.
                    Some(Ok(FdtStructEvent::Prop(name, data)))
                }
                // FDT_NOP
                0x00000004 => {
                    ptr = &ptr[1..];
                    continue;
                }
                // FDT_END
                0x00000009 => None,
                token_type => Some(Err(DeviceTreeError::InvalidTokenType(token_type))),
            };
        })
    }
}

/// An error encountered when reading a DeviceTree.
#[derive(Debug)]
pub enum DeviceTreeError {
    BadMagic(u32),
    IncompatibleVersion(u32, u32),
    InvalidDeviceTree,
    InvalidTokenType(u32),
    StringMissingNulTerminator(u32),

    UnexpectedEndOfStructBlock,
    UnexpectedEvent,
}

fn for_each_property<'dt, E>(
    events: &mut iter::Peekable<impl Iterator<Item = Result<FdtStructEvent<'dt>, DeviceTreeError>>>,
    func: &mut impl for<'a> FnMut(FdtNodePath<'dt, 'a>, &'dt BStr, &'dt BStr) -> Result<(), E>,
    node: StackLinkedList<&'dt BStr>,
) -> Result<(), Either<DeviceTreeError, E>> {
    // Read the first event, which should be the BeginNode starting the node we're trying to read.
    let event = events
        .next()
        .ok_or(Either::Left(DeviceTreeError::UnexpectedEndOfStructBlock))?
        .map_err(Either::Left)?;
    let FdtStructEvent::BeginNode(node_name) = event else {
        return Err(Either::Left(DeviceTreeError::UnexpectedEvent));
    };
    let node = node.cons(node_name);

    // Parse properties and subnodes until we get to the end.
    loop {
        if matches!(events.peek(), Some(Ok(FdtStructEvent::BeginNode(_)))) {
            for_each_property(events, func, node)?;
        } else {
            let event = events
                .next()
                .ok_or(Either::Left(DeviceTreeError::UnexpectedEndOfStructBlock))?
                .map_err(Either::Left)?;
            match event {
                FdtStructEvent::BeginNode(_) => {
                    return Err(Either::Left(DeviceTreeError::UnexpectedEvent))
                }
                FdtStructEvent::EndNode => break Ok(()),
                FdtStructEvent::Prop(prop, value) => {
                    func(FdtNodePath(node), prop, value).map_err(Either::Right)?
                }
            }
        }
    }
}

/// The flattened DeviceTree header.
///
/// All fields are big-endian.
#[derive(Debug)]
#[repr(C)]
struct FdtHeader {
    magic: u32,
    totalsize: u32,
    off_dt_struct: u32,
    off_dt_strings: u32,
    off_mem_rsvmap: u32,
    version: u32,
    last_comp_version: u32,
    boot_cpuid_phys: u32,
    size_dt_strings: u32,
    size_dt_struct: u32,
}

/// A memory reservation from the appropriate block of the DeviceTree.
///
/// All fields are big-endian.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[repr(C)]
struct FdtMemRsv {
    addr: u64,
    size: u64,
}

/// An event returned from iterating over the structure block of a flattened DeviceTree.
#[derive(Clone, Copy, Debug)]
pub enum FdtStructEvent<'dt> {
    BeginNode(&'dt BStr),
    EndNode,
    Prop(&'dt BStr, &'dt BStr),
}

/// The path to a node. Note that the list this contains is in _reverse_ order.
#[derive(Debug, Eq, PartialEq)]
pub struct FdtNodePath<'dt, 'list>(pub StackLinkedList<'list, &'dt BStr>);

impl<'dt, 'list> fmt::Display for FdtNodePath<'dt, 'list> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match (self.0).0 {
            Some((hd, tl)) => write!(fmt, "{}{}/", FdtNodePath(*tl), hd),
            None => Ok(()),
        }
    }
}

impl<'dt, 'list, U> PartialEq<[U]> for FdtNodePath<'dt, 'list>
where
    &'dt BStr: PartialEq<U>,
{
    fn eq(&self, other: &[U]) -> bool {
        self.0.iter().eq(other.iter().rev())
    }
}

impl<'dt, 'list, U, const N: usize> PartialEq<[U; N]> for FdtNodePath<'dt, 'list>
where
    &'dt BStr: PartialEq<U>,
{
    fn eq(&self, other: &[U; N]) -> bool {
        self.0.iter().eq(other.iter().rev())
    }
}