|  | 
|  | 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