src/descriptor.rs

// SPDX-FileCopyrightText: Copyright 2023 Arm Limited and/or its affiliates <open-source-office@arm.com>
// SPDX-License-Identifier: MIT OR Apache-2.0

//! Memory descriptor

use packed_struct::prelude::*;

use core::cell::UnsafeCell;
use core::ptr;

use crate::kernel_space::KernelSpace;
use crate::MemoryAttributesIndex;

use super::address::PhysicalAddress;

/// Memory shareability
#[derive(PrimitiveEnum_u8, Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum Shareability {
    #[default]
    NonShareable = 0b00,
    Outer = 0b10,
    Inner = 0b11,
}

/// Data access permission
#[derive(PrimitiveEnum_u8, Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum DataAccessPermissions {
    #[default]
    ReadWrite_None = 0b00,
    ReadWrite_ReadWrite = 0b01,
    ReadOnly_None = 0b10,
    ReadOnly_ReadOnly = 0b11,
}

/// Memory attributes
#[derive(PackedStruct, Clone, Debug, PartialEq, Eq, Default)]
#[packed_struct(size_bytes = "8", bit_numbering = "lsb0")]
pub struct Attributes {
    #[packed_field(bits = "54")]
    pub uxn: bool,
    #[packed_field(bits = "53")]
    pub pxn: bool,
    #[packed_field(bits = "52")]
    pub contiguous: bool,
    #[packed_field(bits = "11")]
    pub not_global: bool,
    #[packed_field(bits = "10")]
    pub access_flag: bool,
    #[packed_field(bits = "9..=8", ty = "enum")]
    pub shareability: Shareability,
    #[packed_field(bits = "7..=6", ty = "enum")]
    pub data_access_permissions: DataAccessPermissions,
    #[packed_field(bits = "5")]
    pub non_secure: bool,
    #[packed_field(bits = "4..=2", ty = "enum")]
    pub mem_attr_index: MemoryAttributesIndex,
}

impl From<Attributes> for u64 {
    fn from(attributes: Attributes) -> Self {
        u64::from_be_bytes(attributes.pack().unwrap())
    }
}

impl From<u64> for Attributes {
    fn from(bits: u64) -> Self {
        Self::unpack(&bits.to_be_bytes()).unwrap()
    }
}

/// Next level attributes
#[derive(PackedStruct, Clone, Debug, PartialEq, Eq, Default)]
#[packed_struct(size_bytes = "8", bit_numbering = "lsb0")]
pub struct NextLevelAttributes {
    #[packed_field(bits = "63")]
    ns_table: bool,
    #[packed_field(bits = "62..=61")]
    ap_table: Integer<u8, packed_bits::Bits<2>>,
    #[packed_field(bits = "60")]
    xn_table: bool,
    #[packed_field(bits = "59")]
    pxn_table: bool,
}

impl From<NextLevelAttributes> for u64 {
    fn from(attributes: NextLevelAttributes) -> Self {
        u64::from_be_bytes(attributes.pack().unwrap())
    }
}

impl From<u64> for NextLevelAttributes {
    fn from(bits: u64) -> Self {
        Self::unpack(&bits.to_be_bytes()).unwrap()
    }
}

/// Memory descriptor type
#[derive(PartialEq, Eq, Debug)]
pub enum DescriptorType {
    Invalid,
    Block,
    Table,
}

/// Memory descriptor of a memory translation table
#[repr(C)]
#[derive(Default)]
pub struct Descriptor {
    cell: UnsafeCell<u64>,
}

impl Descriptor {
    const ATTR_MASK: u64 = 0xfff8_0000_0000_0ffc;
    const DESCRIPTOR_TYPE_MASK: u64 = 0b11;
    pub const GRANULE_SIZES: [usize; 4] = [0, 0x4000_0000, 0x0020_0000, 0x0000_1000];
    const INVALID_DESCRIPTOR_VALUE: u64 = 0x0;
    const NEXT_ATTR_MASK: u64 = 0xf800_0000_0000_0000;
    const OA_MASK: u64 = 0x0000_ffff_ffff_f000;
    const TABLE_BIT: u64 = 0b10;
    const TABLE_ENTRY_COUNT: usize = 512;
    const TA_MASK: u64 = 0x0000_ffff_ffff_f000;
    const VALID_BIT: u64 = 0b01;

    /// Query descriptor type
    pub fn get_descriptor_type(&self, level: usize) -> DescriptorType {
        assert!(level <= 3);

        let desc_type_bits = self.get() & Self::DESCRIPTOR_TYPE_MASK;
        if desc_type_bits & Self::VALID_BIT != 0 {
            if level == 3 {
                assert_eq!(Self::TABLE_BIT, desc_type_bits & Self::TABLE_BIT);
                DescriptorType::Block
            } else if desc_type_bits & Self::TABLE_BIT != 0 {
                DescriptorType::Table
            } else {
                DescriptorType::Block
            }
        } else {
            DescriptorType::Invalid
        }
    }

    // Invalid descriptor functions

    /// Check if it is a valid descriptor
    pub fn is_valid(&self) -> bool {
        self.get() & Self::VALID_BIT != 0
    }

    // Block descriptor functions

    /// Set block descriptor
    pub fn set_block_descriptor(
        &mut self,
        level: usize,
        output_address: PhysicalAddress,
        attributes: Attributes,
    ) {
        let attr: u64 = attributes.into();

        assert!(level <= 3);
        assert!(self.get_descriptor_type(level) != DescriptorType::Table);
        assert_eq!(0, output_address.0 & !Self::get_oa_mask(level));
        assert_eq!(0, attr & !Self::ATTR_MASK);

        let table_bit = if level < 3 { 0 } else { Self::TABLE_BIT };

        self.set(Self::VALID_BIT | table_bit | output_address.0 as u64 | attr);
    }

    /// Get output address from the block descriptor
    pub fn get_block_output_address(&self, level: usize) -> PhysicalAddress {
        assert!(level <= 3);
        assert_eq!(DescriptorType::Block, self.get_descriptor_type(level));

        PhysicalAddress((self.get() & Self::OA_MASK) as usize)
    }

    /// Set the attributes of the block descriptor
    pub fn set_block_attributes(&mut self, level: usize, attributes: Attributes) {
        assert!(level <= 3);
        let attr: u64 = attributes.into();
        assert_eq!(0, attr & !Self::ATTR_MASK);
        assert_eq!(DescriptorType::Block, self.get_descriptor_type(level));

        self.modify(|d| (d & !Self::ATTR_MASK) | attr);
    }

    /// Get the attributes of the block descriptor
    pub fn get_block_attributes(&self, level: usize) -> Attributes {
        assert!(level <= 3);
        assert_eq!(DescriptorType::Block, self.get_descriptor_type(level));

        Attributes::from(self.get() & Self::ATTR_MASK)
    }

    /// Set block descriptor to invalid
    pub fn set_block_descriptor_to_invalid(&mut self, level: usize) {
        assert!(level <= 3);
        assert_eq!(DescriptorType::Block, self.get_descriptor_type(level));

        self.set(Self::INVALID_DESCRIPTOR_VALUE)
    }

    /// Set block or invalid descriptor to invalid
    pub fn set_block_or_invalid_descriptor_to_invalid(&mut self, level: usize) {
        assert!(level <= 3);
        assert!(DescriptorType::Table != self.get_descriptor_type(level));

        self.set(Self::INVALID_DESCRIPTOR_VALUE)
    }

    /// Set table descriptor
    ///
    /// Safety: The caller has to ensure that the passed next level table has the same life as
    /// the descriptor.
    pub unsafe fn set_table_descriptor(
        &mut self,
        level: usize,
        next_level_table: &mut [Descriptor],
        next_level_attributes: Option<NextLevelAttributes>,
    ) {
        assert!(level <= 2);
        assert_eq!(Self::TABLE_ENTRY_COUNT, next_level_table.len());
        assert!(self.get_descriptor_type(level) != DescriptorType::Table);

        let table_addr = KernelSpace::kernel_to_pa(next_level_table.as_ptr() as u64);
        assert_eq!(0, table_addr & !Self::TA_MASK);

        let mut raw_desc_value = Self::VALID_BIT | Self::TABLE_BIT | table_addr;

        if let Some(next_attr) = next_level_attributes {
            let next_attr_bits: u64 = next_attr.into();
            assert_eq!(0, next_attr_bits & !Self::NEXT_ATTR_MASK);
            raw_desc_value |= next_attr_bits;
        }

        self.set(raw_desc_value);
    }

    /// Get next level table
    ///
    /// Safety: The returned next level table is based on the address read from the descriptor.
    /// The caller has to ensure that no other references are being used of the table.
    pub unsafe fn get_next_level_table(&self, level: usize) -> &[Descriptor] {
        assert!(level <= 2);
        assert_eq!(DescriptorType::Table, self.get_descriptor_type(level));

        let table_address =
            KernelSpace::pa_to_kernel(self.get() & Self::TA_MASK) as *const Descriptor;
        core::slice::from_raw_parts(table_address, Self::TABLE_ENTRY_COUNT)
    }

    /// Get mutable next level table
    ///
    /// Safety: The returned next level table is based on the address read from the descriptor.
    /// The caller has to ensure that no other references are being used of the table.
    pub unsafe fn get_next_level_table_mut(&mut self, level: usize) -> &mut [Descriptor] {
        assert!(level <= 2);
        assert_eq!(DescriptorType::Table, self.get_descriptor_type(level));

        let table_address =
            KernelSpace::pa_to_kernel(self.get() & Self::TA_MASK) as *mut Descriptor;
        core::slice::from_raw_parts_mut(table_address, Self::TABLE_ENTRY_COUNT)
    }

    /// Get next level attributes
    pub fn get_next_level_attributes(&self, level: usize) -> NextLevelAttributes {
        assert!(level <= 2);
        assert_eq!(DescriptorType::Table, self.get_descriptor_type(level));

        NextLevelAttributes::from(self.get() & Self::NEXT_ATTR_MASK)
    }

    /// Set table descriptor to invalid
    ///
    /// **Unsafe:** The returned descriptor reference must be released by the caller, i.e. release
    /// to `PagePool`
    pub unsafe fn set_table_descriptor_to_invalid(&mut self, level: usize) -> &mut [Descriptor] {
        assert!(level <= 2);
        assert_eq!(DescriptorType::Table, self.get_descriptor_type(level));

        let table_address =
            KernelSpace::pa_to_kernel(self.get() & Self::TA_MASK) as *mut Descriptor;
        self.set(Self::INVALID_DESCRIPTOR_VALUE);
        core::slice::from_raw_parts_mut(table_address, Self::TABLE_ENTRY_COUNT)
    }

    /// Get raw descriptor value
    fn get(&self) -> u64 {
        unsafe { ptr::read_volatile(self.cell.get()) }
    }

    /// Set raw descriptor value
    fn set(&mut self, value: u64) {
        // Safety: The cell should point to a valid address and the assembly code is just a data
        // barrier.
        unsafe {
            ptr::write_volatile(self.cell.get(), value);
            #[cfg(target_arch = "aarch64")]
            core::arch::asm!("dsb nsh");
        }
    }

    /// Modify raw descriptor value
    fn modify<F>(&mut self, f: F)
    where
        F: Fn(u64) -> u64,
    {
        self.set(f(self.get()))
    }

    /// Get output address mask
    fn get_oa_mask(level: usize) -> usize {
        Self::OA_MASK as usize & !(Self::GRANULE_SIZES[level] - 1)
    }
}

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

    #[test]
    fn test_attributes() {
        let attributes = Attributes::default();
        assert_eq!(0u64, attributes.into());

        let attributes = Attributes {
            uxn: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 54, attributes.into());

        let attributes = Attributes {
            pxn: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 53, attributes.into());

        let attributes = Attributes {
            contiguous: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 52, attributes.into());

        let attributes = Attributes {
            not_global: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 11, attributes.into());

        let attributes = Attributes {
            access_flag: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 10, attributes.into());

        let attributes = Attributes {
            non_secure: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 5, attributes.into());

        let attributes = Attributes {
            mem_attr_index: MemoryAttributesIndex::Normal_IWBWA_OWBWA,
            ..Default::default()
        };
        assert_eq!(
            (MemoryAttributesIndex::Normal_IWBWA_OWBWA as u64) << 2,
            attributes.into()
        );

        let attributes: Attributes = 0.into();
        assert!(!attributes.uxn);
        assert!(!attributes.pxn);
        assert!(!attributes.contiguous);
        assert!(!attributes.not_global);
        assert!(!attributes.access_flag);
        assert_eq!(Shareability::NonShareable, attributes.shareability);
        assert_eq!(
            DataAccessPermissions::ReadWrite_None,
            attributes.data_access_permissions
        );
        assert!(!attributes.non_secure);
        assert_eq!(
            MemoryAttributesIndex::Device_nGnRnE,
            attributes.mem_attr_index
        );
    }

    #[test]
    fn test_next_level_attributes() {
        let next_level_attributes = NextLevelAttributes::default();
        assert_eq!(0u64, next_level_attributes.into());

        let next_level_attributes = NextLevelAttributes {
            ns_table: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 63, next_level_attributes.into());

        let next_level_attributes = NextLevelAttributes {
            ap_table: 3.into(),
            ..Default::default()
        };
        assert_eq!(3u64 << 61, next_level_attributes.into());

        let next_level_attributes = NextLevelAttributes {
            xn_table: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 60, next_level_attributes.into());

        let next_level_attributes = NextLevelAttributes {
            pxn_table: true,
            ..Default::default()
        };
        assert_eq!(1u64 << 59, next_level_attributes.into());

        let next_level_attributes: NextLevelAttributes = 0.into();
        assert!(!next_level_attributes.ns_table);
        assert_eq!(0u8, next_level_attributes.ap_table.into());
        assert!(!next_level_attributes.xn_table);
        assert!(!next_level_attributes.pxn_table);

        let next_level_attributes: NextLevelAttributes = u64::MAX.into();
        assert!(next_level_attributes.ns_table);
        assert_eq!(3u8, next_level_attributes.ap_table.into());
        assert!(next_level_attributes.xn_table);
        assert!(next_level_attributes.pxn_table);
    }

    #[test]
    fn test_descriptor_get_type() {
        let descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };
        assert_eq!(DescriptorType::Invalid, descriptor.get_descriptor_type(1));

        let descriptor = Descriptor {
            cell: UnsafeCell::new(1),
        };
        assert_eq!(DescriptorType::Block, descriptor.get_descriptor_type(1));

        let descriptor = Descriptor {
            cell: UnsafeCell::new(3),
        };
        assert_eq!(DescriptorType::Table, descriptor.get_descriptor_type(1));

        let descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };
        assert_eq!(DescriptorType::Invalid, descriptor.get_descriptor_type(3));

        let descriptor = Descriptor {
            cell: UnsafeCell::new(3),
        };
        assert_eq!(DescriptorType::Block, descriptor.get_descriptor_type(3));
    }

    #[test]
    fn test_descriptor_is_valid() {
        let descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };
        assert!(!descriptor.is_valid());

        let descriptor = Descriptor {
            cell: UnsafeCell::new(1),
        };
        assert!(descriptor.is_valid());
    }

    #[test]
    fn test_descriptor_set_block_to_block_again() {
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(1),
        };

        descriptor.set_block_descriptor(1, PhysicalAddress(0), Attributes::default());
        assert_eq!(0x1, descriptor.get());
    }

    #[test]
    #[should_panic]
    fn test_descriptor_set_block_invalid_oa() {
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        descriptor.set_block_descriptor(1, PhysicalAddress(1 << 63), Attributes::default());
    }

    #[test]
    fn test_descriptor_block() {
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        descriptor.set_block_descriptor(
            1,
            PhysicalAddress(0x0000000f_c0000000),
            Attributes {
                uxn: true,
                ..Default::default()
            },
        );
        assert_eq!(0x0040000f_c0000001, descriptor.get());

        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        descriptor.set_block_descriptor(
            3,
            PhysicalAddress(0x0000000f_fffff000),
            Attributes {
                uxn: true,
                ..Default::default()
            },
        );
        assert_eq!(0x0040000f_fffff003, descriptor.get());

        assert_eq!(
            PhysicalAddress(0x0000000f_fffff000),
            descriptor.get_block_output_address(3)
        );
        assert_eq!(
            Attributes {
                uxn: true,
                ..Default::default()
            },
            descriptor.get_block_attributes(3)
        );

        descriptor.set_block_attributes(
            3,
            Attributes {
                pxn: true,
                ..Default::default()
            },
        );
        assert_eq!(
            Attributes {
                pxn: true,
                ..Default::default()
            },
            descriptor.get_block_attributes(3)
        );
    }

    #[test]
    #[should_panic]
    fn test_descriptor_invalid_block_to_invalid() {
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        descriptor.set_block_descriptor_to_invalid(0);
    }

    #[test]
    fn test_descriptor_block_to_invalid() {
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(3),
        };

        descriptor.set_block_descriptor_to_invalid(3);
        assert_eq!(0, descriptor.get());
    }

    #[test]
    fn test_descriptor_block_or_invalid_block_to_invalid() {
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        descriptor.set_block_or_invalid_descriptor_to_invalid(0);
        assert_eq!(0, descriptor.get());

        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(3),
        };

        descriptor.set_block_or_invalid_descriptor_to_invalid(3);
        assert_eq!(0, descriptor.get());
    }

    #[test]
    #[should_panic]
    fn test_descriptor_level3_to_table() {
        let mut next_level_table = [Descriptor {
            cell: UnsafeCell::new(0),
        }];
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        unsafe {
            descriptor.set_table_descriptor(3, &mut next_level_table, None);
        }
    }

    #[test]
    fn test_descriptor_block_to_table() {
        let next_level_table =
            unsafe { core::slice::from_raw_parts_mut(0x1000 as *mut Descriptor, 512) };
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(1),
        };

        unsafe {
            descriptor.set_table_descriptor(0, next_level_table, None);
        }
        assert_eq!(0x1003, descriptor.get());
    }

    #[test]
    #[should_panic]
    fn test_descriptor_table_invalid_count() {
        let next_level_table =
            unsafe { core::slice::from_raw_parts_mut(0x800 as *mut Descriptor, 511) };
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        unsafe {
            descriptor.set_table_descriptor(0, next_level_table, None);
        }
    }

    #[test]
    #[should_panic]
    fn test_descriptor_table_non_aligned() {
        let next_level_table =
            unsafe { core::slice::from_raw_parts_mut(0x800 as *mut Descriptor, 512) };
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        unsafe {
            descriptor.set_table_descriptor(0, next_level_table, None);
        }
    }

    #[test]
    fn test_descriptor_table() {
        let next_level_table = unsafe {
            core::slice::from_raw_parts_mut(0x0000_000c_ba98_7000 as *mut Descriptor, 512)
        };
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        unsafe {
            descriptor.set_table_descriptor(0, next_level_table, None);
        }
        assert_eq!(0x0000_000c_ba98_7003, descriptor.get());
    }

    #[test]
    fn test_descriptor_table_next_level_attr() {
        const NEXT_LEVEL_ADDR: u64 = 0x0000_000c_ba98_7000;
        let next_level_table =
            unsafe { core::slice::from_raw_parts_mut(NEXT_LEVEL_ADDR as *mut Descriptor, 512) };
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(0),
        };

        unsafe {
            descriptor.set_table_descriptor(
                0,
                next_level_table,
                Some(NextLevelAttributes {
                    ns_table: true,
                    ..Default::default()
                }),
            );
        }
        assert_eq!(NEXT_LEVEL_ADDR | 0x8000_0000_0000_0003, descriptor.get());
    }

    #[test]
    fn test_descriptor_table_get_next_level_table() {
        const NEXT_LEVEL_ADDR: u64 = 0x0000_000c_ba98_7000;
        let descriptor = Descriptor {
            cell: UnsafeCell::new(NEXT_LEVEL_ADDR | 0x8000_0000_0000_0003),
        };
        assert_eq!(KernelSpace::pa_to_kernel(NEXT_LEVEL_ADDR), unsafe {
            descriptor.get_next_level_table(0).as_ptr() as u64
        });
    }

    #[test]
    fn test_descriptor_table_get_next_level_table_mut() {
        const NEXT_LEVEL_ADDR: u64 = 0x0000_000c_ba98_7000;
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(NEXT_LEVEL_ADDR | 0x8000_0000_0000_0003),
        };
        assert_eq!(KernelSpace::pa_to_kernel(NEXT_LEVEL_ADDR), unsafe {
            descriptor.get_next_level_table_mut(0).as_ptr() as *mut Descriptor as u64
        });
    }

    #[test]
    fn test_descriptor_table_get_next_level_attr() {
        const NEXT_LEVEL_ADDR: u64 = 0x0000_000c_ba98_7000;
        let descriptor = Descriptor {
            cell: UnsafeCell::new(NEXT_LEVEL_ADDR | 0x8000_0000_0000_0003),
        };
        assert_eq!(
            NextLevelAttributes {
                ns_table: true,
                ..Default::default()
            },
            descriptor.get_next_level_attributes(0)
        );
    }

    #[test]
    fn test_descriptor_table_set_to_invalid() {
        const NEXT_LEVEL_ADDR: u64 = 0x0000_000c_ba98_7000;
        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(NEXT_LEVEL_ADDR | 0x8000_0000_0000_0003),
        };
        assert_eq!(KernelSpace::pa_to_kernel(NEXT_LEVEL_ADDR), unsafe {
            descriptor.set_table_descriptor_to_invalid(0).as_ptr() as *mut Descriptor as u64
        });
        assert_eq!(0, descriptor.get());
    }

    #[test]
    fn test_descriptor_raw_interface() {
        let cell_value = 0x01234567_89abcdefu64;
        let cell_new_value = 0x12345678_9abcdef0u64;

        let mut descriptor = Descriptor {
            cell: UnsafeCell::new(cell_value),
        };

        assert_eq!(cell_value, descriptor.get());

        descriptor.set(cell_new_value);
        assert_eq!(cell_new_value, descriptor.get());

        descriptor.modify(|d| d + 1);
        assert_eq!(cell_new_value + 1, descriptor.get());
    }
}