|
| 1 | +use core::{ |
| 2 | + mem::{self, MaybeUninit}, |
| 3 | + ops::Range, |
| 4 | + slice, |
| 5 | +}; |
| 6 | + |
| 7 | +use r_efi::{protocols::simple_text_output, system::RuntimeServices as RawRuntimeServices}; |
| 8 | +use uefi::{prelude::*, proto::console::text::Color, Completion}; |
| 9 | + |
| 10 | +static mut SYSTEM_TABLE: MaybeUninit<SystemTable<Boot>> = MaybeUninit::uninit(); |
| 11 | + |
| 12 | +pub fn system_table() -> &'static SystemTable<Boot> { |
| 13 | + unsafe { &*SYSTEM_TABLE.as_ptr() } |
| 14 | +} |
| 15 | + |
| 16 | +pub fn raw_runtime_services() -> &'static RawRuntimeServices { |
| 17 | + unsafe { &*(system_table().runtime_services() as *const _ as *const _) } |
| 18 | +} |
| 19 | + |
| 20 | +macro_rules! print { |
| 21 | + ($($arg:tt)*) => { { |
| 22 | + use ::core::fmt::Write; |
| 23 | + let _ = ::core::write!($crate::util::system_table().stdout(), $($arg)*); |
| 24 | + } } |
| 25 | +} |
| 26 | + |
| 27 | +macro_rules! println { |
| 28 | + ($($arg:tt)*) => { { |
| 29 | + use ::core::fmt::Write; |
| 30 | + let _ = ::core::writeln!($crate::util::system_table().stdout(), $($arg)*); |
| 31 | + } } |
| 32 | +} |
| 33 | + |
| 34 | +#[entry] |
| 35 | +fn efi_main(_image_handle: Handle, system_table: SystemTable<Boot>) -> Status { |
| 36 | + unsafe { SYSTEM_TABLE = MaybeUninit::new(system_table) }; |
| 37 | + |
| 38 | + main(); |
| 39 | + |
| 40 | + Status::LOAD_ERROR |
| 41 | +} |
| 42 | + |
| 43 | +fn main() { |
| 44 | + let stdout = system_table().stdout(); |
| 45 | + |
| 46 | + let (foreground, background) = unsafe { |
| 47 | + let raw_stdout = &*(stdout as *const _ as *const simple_text_output::Protocol); |
| 48 | + let mode = &*raw_stdout.mode; |
| 49 | + mem::transmute(( |
| 50 | + (mode.attribute & 0xF) as u8, |
| 51 | + (mode.attribute >> 4 & 0x7) as u8, |
| 52 | + )) |
| 53 | + }; |
| 54 | + |
| 55 | + match crate::main() { |
| 56 | + Status::SUCCESS => { |
| 57 | + let _ = stdout.set_color(Color::LightGreen, background); |
| 58 | + println!("╔══════════╗"); |
| 59 | + println!("║ Success! ║"); |
| 60 | + println!("╚══════════╝"); |
| 61 | + } |
| 62 | + status => { |
| 63 | + let _ = stdout.set_color(Color::LightRed, background); |
| 64 | + println!("[-] error: {:?}", status); |
| 65 | + } |
| 66 | + } |
| 67 | + |
| 68 | + let _ = stdout.set_color(Color::White, background); |
| 69 | + print!("Press any key to continue..."); |
| 70 | + let _ = stdout.set_color(foreground, background); |
| 71 | + |
| 72 | + let stdin = system_table().stdin(); |
| 73 | + let _ = system_table() |
| 74 | + .boot_services() |
| 75 | + .wait_for_event(&mut [stdin.wait_for_key_event()]); |
| 76 | + let _ = stdin.read_key(); |
| 77 | + |
| 78 | + println!(); |
| 79 | +} |
| 80 | + |
| 81 | +macro_rules! unwrap { |
| 82 | + ($expr:expr) => { |
| 83 | + $expr?.split().1 |
| 84 | + }; |
| 85 | +} |
| 86 | + |
| 87 | +static mut BUFFER: [u8; 4096] = [0; 4096]; |
| 88 | + |
| 89 | +pub fn region_containing(address: usize) -> uefi::Result<Range<usize>> { |
| 90 | + let (status, (_, descriptors)) = system_table() |
| 91 | + .boot_services() |
| 92 | + .memory_map(unsafe { &mut BUFFER })? |
| 93 | + .split(); |
| 94 | + |
| 95 | + let region = descriptors |
| 96 | + .map(|descriptor| { |
| 97 | + let start = descriptor.phys_start as usize; |
| 98 | + let end = start + descriptor.page_count as usize * 4096; |
| 99 | + |
| 100 | + start..end |
| 101 | + }) |
| 102 | + .find(|region| region.contains(&address)); |
| 103 | + |
| 104 | + match region { |
| 105 | + Some(region) => Ok(Completion::new(status, region)), |
| 106 | + None => Err(Status::NOT_FOUND.into()), |
| 107 | + } |
| 108 | +} |
| 109 | + |
| 110 | +pub unsafe fn range_to_slice(range: Range<usize>) -> &'static mut [u8] { |
| 111 | + slice::from_raw_parts_mut(range.start as _, range.len()) |
| 112 | +} |
| 113 | + |
| 114 | +pub fn search_for_contiguous(slice: &mut [u8], item: u8, count: usize) -> uefi::Result<&mut [u8]> { |
| 115 | + let mut current = 0; |
| 116 | + |
| 117 | + for (n, i) in slice.iter().enumerate() { |
| 118 | + if *i == item { |
| 119 | + current += 1; |
| 120 | + |
| 121 | + if current == count { |
| 122 | + let slice = &mut slice[n + 1 - count..n + 1]; |
| 123 | + |
| 124 | + return Ok(slice.into()); |
| 125 | + } |
| 126 | + } else if current != 0 { |
| 127 | + current = 0; |
| 128 | + } |
| 129 | + } |
| 130 | + |
| 131 | + Err(Status::NOT_FOUND.into()) |
| 132 | +} |
0 commit comments