use core::fmt::Debug;
use patina::pi::dxe_services;
use r_efi::efi;
use crate::error;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Error {
InvalidStateTransition,
BlockOutsideRange,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MemoryBlock {
Unallocated(dxe_services::MemorySpaceDescriptor),
Allocated(dxe_services::MemorySpaceDescriptor),
}
pub enum StateTransition {
Add(dxe_services::GcdMemoryType, u64, u64),
Remove,
Allocate(efi::Handle, Option<efi::Handle>),
AllocateRespectingOwnership(efi::Handle, Option<efi::Handle>),
Free,
FreePreservingOwnership,
SetAttributes(u64),
SetCapabilities(u64),
}
impl Debug for StateTransition {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
StateTransition::Add(memory_type, capabilities, attributes) => f
.debug_struct("Add")
.field("memory_type", memory_type)
.field("capabilities", &format_args!("{capabilities:#X}"))
.field("attributes", &format_args!("{attributes:#X}"))
.finish(),
StateTransition::Remove => f.debug_struct("Remove").finish(),
StateTransition::Allocate(image_handle, device_handle) => f
.debug_struct("Allocate")
.field("image_handle", image_handle)
.field("device_handle", device_handle)
.finish(),
StateTransition::AllocateRespectingOwnership(image_handle, device_handle) => f
.debug_struct("AllocateRespectingOwnership")
.field("image_handle", image_handle)
.field("device_handle", device_handle)
.finish(),
StateTransition::Free => f.debug_struct("Free").finish(),
StateTransition::FreePreservingOwnership => f.debug_struct("FreePreservingOwnership").finish(),
StateTransition::SetAttributes(attributes) => {
f.debug_struct("SetAttributes").field("attributes", &format_args!("{attributes:#X}")).finish()
}
StateTransition::SetCapabilities(capabilities) => {
f.debug_struct("SetCapabilities").field("capabilities", &format_args!("{capabilities:#X}")).finish()
}
}
}
}
#[derive(Debug)]
pub enum MemoryBlockSplit<'a> {
Same(&'a mut MemoryBlock),
Before(&'a mut MemoryBlock, MemoryBlock),
After(&'a mut MemoryBlock, MemoryBlock),
Middle(&'a mut MemoryBlock, MemoryBlock, MemoryBlock),
}
impl MemoryBlock {
pub fn merge(&mut self, other: &mut MemoryBlock) -> bool {
if self.is_same_state(other) && self.end() == other.start() {
self.as_mut().length += other.as_ref().length;
other.as_mut().length = 0;
true
} else {
false
}
}
pub fn split(&mut self, base_address: usize, len: usize) -> Result<MemoryBlockSplit<'_>, Error> {
let start = base_address;
let end = base_address + len;
if !(self.start() <= start && start < end && end <= self.end()) {
return Err(Error::BlockOutsideRange);
}
if self.start() == start && end == self.end() {
return Ok(MemoryBlockSplit::Same(self));
}
if self.start() == start && end < self.end() {
let mut next = MemoryBlock::clone(self);
self.as_mut().base_address = base_address as u64;
self.as_mut().length = len as u64;
next.as_mut().base_address = end as u64;
next.as_mut().length -= len as u64;
return Ok(MemoryBlockSplit::Before(self, next));
}
if self.start() < start && end == self.end() {
let mut next = MemoryBlock::clone(self);
self.as_mut().length -= len as u64;
next.as_mut().base_address = base_address as u64;
next.as_mut().length = len as u64;
return Ok(MemoryBlockSplit::After(self, next));
}
if self.start() < start && end < self.end() {
let mut next = MemoryBlock::clone(self);
let mut last = MemoryBlock::clone(self);
self.as_mut().length = (start - self.start()) as u64;
next.as_mut().base_address = base_address as u64;
next.as_mut().length = len as u64;
last.as_mut().length = (last.end() - end) as u64;
last.as_mut().base_address = end as u64;
return Ok(MemoryBlockSplit::Middle(self, next, last));
}
unreachable!()
}
pub fn split_state_transition(
&mut self,
base_address: usize,
len: usize,
transition: StateTransition,
) -> Result<MemoryBlockSplit<'_>, Error> {
let mut split = self.split(base_address, len)?;
match &mut split {
MemoryBlockSplit::Same(mb) => {
mb.state_transition(transition)?;
}
MemoryBlockSplit::Before(mb, next) => {
if let Err(e) = mb.state_transition(transition) {
mb.merge(next);
error!(e);
}
}
MemoryBlockSplit::After(prev, mb) => {
if let Err(e) = mb.state_transition(transition) {
prev.merge(mb);
error!(e)
}
}
MemoryBlockSplit::Middle(prev, mb, next) => {
if let Err(e) = mb.state_transition(transition) {
mb.merge(next);
prev.merge(mb);
error!(e)
}
}
}
Ok(split)
}
pub fn is_same_state(&self, other: &MemoryBlock) -> bool {
matches!((self, other),
(MemoryBlock::Unallocated(self_desc), MemoryBlock::Unallocated(other_desc)) |
(MemoryBlock::Allocated(self_desc), MemoryBlock::Allocated(other_desc))
if self_desc.memory_type == other_desc.memory_type
&& self_desc.attributes == other_desc.attributes
&& self_desc.capabilities == other_desc.capabilities
&& self_desc.device_handle == other_desc.device_handle
&& self_desc.image_handle == other_desc.image_handle
)
}
pub fn state_transition(&mut self, transition: StateTransition) -> Result<(), Error> {
match transition {
StateTransition::Add(memory_type, capabilities, attributes) => {
self.add_transition(memory_type, capabilities, attributes)
}
StateTransition::Remove => self.remove_transition(),
StateTransition::Allocate(image_handle, device_handle) => {
self.allocate_transition(image_handle, device_handle, false)
}
StateTransition::AllocateRespectingOwnership(image_handle, device_handle) => {
self.allocate_transition(image_handle, device_handle, true)
}
StateTransition::Free => self.free_transition(false),
StateTransition::FreePreservingOwnership => self.free_transition(true),
StateTransition::SetAttributes(attributes) => self.attribute_transition(attributes),
StateTransition::SetCapabilities(capabilities) => self.capabilities_transition(capabilities),
}
}
pub fn add_transition(
&mut self,
memory_type: dxe_services::GcdMemoryType,
capabilities: u64,
attributes: u64,
) -> Result<(), Error> {
match self {
Self::Unallocated(md)
if md.memory_type == dxe_services::GcdMemoryType::NonExistent
&& memory_type != dxe_services::GcdMemoryType::NonExistent =>
{
md.memory_type = memory_type;
md.capabilities = capabilities;
md.attributes = attributes;
Ok(())
}
_ => Err(Error::InvalidStateTransition),
}
}
pub fn remove_transition(&mut self) -> Result<(), Error> {
match self {
Self::Unallocated(md) if md.memory_type != dxe_services::GcdMemoryType::NonExistent => {
md.memory_type = dxe_services::GcdMemoryType::NonExistent;
md.capabilities = 0;
Ok(())
}
_ => Err(Error::InvalidStateTransition),
}
}
pub fn allocate_transition(
&mut self,
image_handle: efi::Handle,
device_handle: Option<efi::Handle>,
respect_ownership: bool,
) -> Result<(), Error> {
match self {
Self::Unallocated(md)
if !matches!(
md.memory_type,
dxe_services::GcdMemoryType::NonExistent | dxe_services::GcdMemoryType::Unaccepted
) =>
{
if respect_ownership && !(md.image_handle == 0 as efi::Handle || md.image_handle == image_handle) {
Err(Error::InvalidStateTransition)?;
}
md.image_handle = image_handle;
if let Some(device_handle) = device_handle {
md.device_handle = device_handle;
}
*self = Self::Allocated(*md);
Ok(())
}
_ => Err(Error::InvalidStateTransition),
}
}
pub fn free_transition(&mut self, preserve_ownership: bool) -> Result<(), Error> {
match self {
Self::Allocated(md) if md.memory_type != dxe_services::GcdMemoryType::NonExistent => {
if !preserve_ownership {
md.image_handle = 0 as efi::Handle;
}
md.device_handle = 0 as efi::Handle;
*self = Self::Unallocated(*md);
Ok(())
}
_ => Err(Error::InvalidStateTransition),
}
}
pub fn attribute_transition(&mut self, attributes: u64) -> Result<(), Error> {
match self {
Self::Allocated(md) | Self::Unallocated(md)
if md.memory_type != dxe_services::GcdMemoryType::NonExistent =>
{
if (md.capabilities | attributes) != md.capabilities {
Err(Error::InvalidStateTransition)
} else {
md.attributes = attributes;
Ok(())
}
}
_ => Err(Error::InvalidStateTransition),
}
}
pub fn capabilities_transition(&mut self, capabilities: u64) -> Result<(), Error> {
match self {
Self::Allocated(md) | Self::Unallocated(md)
if md.memory_type != dxe_services::GcdMemoryType::NonExistent =>
{
if (capabilities & md.attributes) != md.attributes {
Err(Error::InvalidStateTransition)
} else {
md.capabilities = capabilities;
Ok(())
}
}
_ => Err(Error::InvalidStateTransition),
}
}
pub fn start(&self) -> usize {
self.as_ref().base_address as usize
}
pub fn end(&self) -> usize {
(self.as_ref().base_address + self.as_ref().length) as usize
}
pub fn len(&self) -> usize {
self.as_ref().length as usize
}
#[allow(dead_code)]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
impl AsRef<dxe_services::MemorySpaceDescriptor> for MemoryBlock {
fn as_ref(&self) -> &dxe_services::MemorySpaceDescriptor {
match self {
MemoryBlock::Unallocated(msd) | MemoryBlock::Allocated(msd) => msd,
}
}
}
impl AsMut<dxe_services::MemorySpaceDescriptor> for MemoryBlock {
fn as_mut(&mut self) -> &mut dxe_services::MemorySpaceDescriptor {
match self {
MemoryBlock::Unallocated(msd) | MemoryBlock::Allocated(msd) => msd,
}
}
}
#[cfg(test)]
mod memory_block_tests {
use super::*;
use dxe_services::{GcdMemoryType, MemorySpaceDescriptor};
#[test]
fn test_transition_types() {
let block = MemoryBlock::Unallocated(MemorySpaceDescriptor {
base_address: 0,
length: 0,
memory_type: GcdMemoryType::NonExistent,
attributes: 0,
capabilities: 0,
device_handle: 0 as efi::Handle,
image_handle: 0 as efi::Handle,
});
let mut b1 = block;
b1.state_transition(StateTransition::Add(GcdMemoryType::MemoryMappedIo, 0, 0)).unwrap();
assert_eq!(b1.as_ref().memory_type, GcdMemoryType::MemoryMappedIo);
let mut b2 = b1;
b2.state_transition(StateTransition::Remove).unwrap();
assert_eq!(b2.as_ref().memory_type, GcdMemoryType::NonExistent);
let mut b3 = block;
b3.as_mut().memory_type = GcdMemoryType::MemoryMappedIo;
b3.state_transition(StateTransition::Allocate(0 as efi::Handle, None)).unwrap();
match b3 {
MemoryBlock::Allocated(md) => {
assert_eq!(md.image_handle, 0 as efi::Handle);
assert_eq!(md.device_handle, 0 as efi::Handle);
}
_ => panic!("Expected Allocated"),
}
let mut b4 = b3;
b4.state_transition(StateTransition::Free).unwrap();
match b4 {
MemoryBlock::Unallocated(md) => {
assert_eq!(md.image_handle, 0 as efi::Handle);
assert_eq!(md.device_handle, 0 as efi::Handle);
}
_ => panic!("Expected Unallocated"),
}
let mut b5 = block;
b5.as_mut().memory_type = GcdMemoryType::MemoryMappedIo;
b5.as_mut().attributes = 0b11;
b5.as_mut().capabilities = 0b111;
assert!(b5.state_transition(StateTransition::SetAttributes(0b1111)).is_err());
b5.state_transition(StateTransition::SetCapabilities(0b1111)).unwrap();
assert_eq!(b5.as_ref().capabilities, 0b1111);
b5.as_mut().memory_type = GcdMemoryType::MemoryMappedIo;
b5.state_transition(StateTransition::SetAttributes(0b1111)).unwrap();
assert_eq!(b5.as_ref().attributes, 0b1111);
assert!(b5.state_transition(StateTransition::SetCapabilities(0b1011)).is_err());
let mut b7 = block;
assert!(b7.state_transition(StateTransition::SetAttributes(0b1111)).is_err());
assert!(b7.state_transition(StateTransition::SetCapabilities(0b1111)).is_err());
}
}