use std::fmt;
use std::io::{Read, Seek, SeekFrom};
use vm_memory::{Address, ByteValued, Bytes, GuestAddress, GuestMemory, GuestUsize, ReadVolatile};
use crate::loader::{Error as KernelLoaderError, KernelLoader, KernelLoaderResult, Result};
pub struct PE;
#[derive(Debug, PartialEq, Eq)]
pub enum Error {
SeekImageEnd,
SeekImageHeader,
SeekDtbStart,
SeekDtbEnd,
DtbTooBig,
ReadKernelImage,
ReadImageHeader,
ReadDtbImage,
InvalidImage,
InvalidImageMagicNumber,
InvalidBaseAddrAlignment,
}
impl fmt::Display for Error {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let desc = match self {
Error::SeekImageEnd => "unable to seek Image end",
Error::SeekImageHeader => "unable to seek Image header",
Error::ReadImageHeader => "unable to read Image header",
Error::ReadDtbImage => "unable to read DTB image",
Error::SeekDtbStart => "unable to seek DTB start",
Error::SeekDtbEnd => "unable to seek DTB end",
Error::InvalidImage => "invalid Image",
Error::InvalidImageMagicNumber => "invalid Image magic number",
Error::DtbTooBig => "device tree image too big",
Error::ReadKernelImage => "unable to read kernel image",
Error::InvalidBaseAddrAlignment => "base address not aligned to 2 MB",
};
write!(f, "PE Kernel Loader: {}", desc)
}
}
impl std::error::Error for Error {}
#[cfg(target_arch = "aarch64")]
#[repr(C)]
#[derive(Debug, Copy, Clone, Default)]
#[allow(missing_docs)]
pub struct arm64_image_header {
pub code0: u32,
pub code1: u32,
pub text_offset: u64,
pub image_size: u64,
pub flags: u64,
pub res2: u64,
pub res3: u64,
pub res4: u64,
pub magic: u32,
pub res5: u32,
}
#[cfg(target_arch = "aarch64")]
unsafe impl ByteValued for arm64_image_header {}
#[cfg(target_arch = "riscv64")]
#[repr(C)]
#[derive(Debug, Copy, Clone, Default)]
#[allow(missing_docs)]
pub struct riscv64_image_header {
pub code0: u32,
pub code1: u32,
pub text_offset: u64,
pub image_size: u64,
pub flags: u64,
pub version: u32,
pub res1: u32,
pub res2: u64,
pub magic: u64,
pub magic2: u32,
pub res3: u32,
}
#[cfg(target_arch = "riscv64")]
unsafe impl ByteValued for riscv64_image_header {}
impl KernelLoader for PE {
fn load<F, M: GuestMemory>(
guest_mem: &M,
kernel_offset: Option<GuestAddress>,
kernel_image: &mut F,
_highmem_start_address: Option<GuestAddress>,
) -> Result<KernelLoaderResult>
where
F: ReadVolatile + Read + Seek,
{
let kernel_size = kernel_image
.seek(SeekFrom::End(0))
.map_err(|_| Error::SeekImageEnd)? as usize;
#[cfg(target_arch = "aarch64")]
let mut image_header: arm64_image_header = Default::default();
#[cfg(target_arch = "riscv64")]
let mut image_header: riscv64_image_header = Default::default();
kernel_image.rewind().map_err(|_| Error::SeekImageHeader)?;
kernel_image
.read_exact(image_header.as_mut_slice())
.map_err(|_| Error::ReadImageHeader)?;
#[cfg(target_arch = "aarch64")]
if u32::from_le(image_header.magic) != 0x644d_5241 {
return Err(Error::InvalidImageMagicNumber.into());
}
#[cfg(target_arch = "riscv64")]
if u32::from_le(image_header.magic2) != 0x0543_5352 {
return Err(Error::InvalidImageMagicNumber.into());
}
#[cfg(target_arch = "aarch64")]
let text_offset = if u64::from_le(image_header.image_size) == 0 {
0x80000
} else {
u64::from_le(image_header.text_offset)
};
#[cfg(target_arch = "riscv64")]
let text_offset = u64::from_le(image_header.text_offset);
if let Some(kernel_offset) = kernel_offset {
if kernel_offset.raw_value() % 0x0020_0000 != 0 {
return Err(Error::InvalidBaseAddrAlignment.into());
}
}
let mem_offset = kernel_offset
.unwrap_or(GuestAddress(0))
.checked_add(text_offset)
.ok_or(Error::InvalidImage)?;
let mut loader_result = KernelLoaderResult {
kernel_load: mem_offset,
..Default::default()
};
kernel_image.rewind().map_err(|_| Error::SeekImageHeader)?;
guest_mem
.read_exact_volatile_from(mem_offset, kernel_image, kernel_size)
.map_err(|_| Error::ReadKernelImage)?;
loader_result.kernel_end = mem_offset
.raw_value()
.checked_add(kernel_size as GuestUsize)
.ok_or(KernelLoaderError::MemoryOverflow)?;
Ok(loader_result)
}
}
#[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
pub fn load_dtb<F, M: GuestMemory>(
guest_mem: &M,
guest_addr: GuestAddress,
dtb_image: &mut F,
) -> Result<()>
where
F: ReadVolatile + Read + Seek,
{
let dtb_size = dtb_image
.seek(SeekFrom::End(0))
.map_err(|_| Error::SeekDtbEnd)? as usize;
if dtb_size > 0x200000 {
return Err(Error::DtbTooBig.into());
}
dtb_image.rewind().map_err(|_| Error::SeekDtbStart)?;
guest_mem
.read_exact_volatile_from(guest_addr, dtb_image, dtb_size)
.map_err(|_| Error::ReadDtbImage.into())
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
use vm_memory::{Address, GuestAddress};
type GuestMemoryMmap = vm_memory::GuestMemoryMmap<()>;
const MEM_SIZE: u64 = 0x100_0000;
fn create_guest_mem() -> GuestMemoryMmap {
GuestMemoryMmap::from_ranges(&[(GuestAddress(0x0), (MEM_SIZE as usize))]).unwrap()
}
fn make_image_bin() -> Vec<u8> {
let mut v = Vec::new();
#[cfg(target_arch = "aarch64")]
v.extend_from_slice(include_bytes!("test_arm64_image.bin"));
#[cfg(target_arch = "riscv64")]
v.extend_from_slice(include_bytes!("test_riscv64_image.bin"));
v
}
#[cfg(target_arch = "aarch64")]
#[test]
fn load_arm64_image() {
let gm = create_guest_mem();
let mut image = make_image_bin();
let kernel_addr = GuestAddress(0x200000);
let loader_result =
PE::load(&gm, Some(kernel_addr), &mut Cursor::new(&image), None).unwrap();
assert_eq!(loader_result.kernel_load.raw_value(), 0x280000);
assert_eq!(loader_result.kernel_end, 0x281000);
let kernel_offset = GuestAddress(0x0030_0000);
let loader_result = PE::load(&gm, Some(kernel_offset), &mut Cursor::new(&image), None);
assert_eq!(
loader_result,
Err(KernelLoaderError::Pe(Error::InvalidBaseAddrAlignment))
);
image[0x39] = 0x0;
let loader_result = PE::load(&gm, Some(kernel_addr), &mut Cursor::new(&image), None);
assert_eq!(
loader_result,
Err(KernelLoaderError::Pe(Error::InvalidImageMagicNumber))
);
}
#[cfg(target_arch = "riscv64")]
#[test]
fn load_riscv64_image() {
let gm = create_guest_mem();
let mut image = make_image_bin();
let kernel_addr = GuestAddress(0x400000);
let loader_result =
PE::load(&gm, Some(kernel_addr), &mut Cursor::new(&image), None).unwrap();
assert_eq!(loader_result.kernel_load.raw_value(), 0x600000);
assert_eq!(loader_result.kernel_end, 0x601000);
let kernel_offset = GuestAddress(0x0030_0000);
let loader_result = PE::load(&gm, Some(kernel_offset), &mut Cursor::new(&image), None);
assert_eq!(
loader_result,
Err(KernelLoaderError::Pe(Error::InvalidBaseAddrAlignment))
);
image[0x38] = 0x0;
let loader_result = PE::load(&gm, Some(kernel_addr), &mut Cursor::new(&image), None);
assert_eq!(
loader_result,
Err(KernelLoaderError::Pe(Error::InvalidImageMagicNumber))
);
}
}