Demilade Sonuga's blog

All posts

Refactoring II

2022-12-03

Unsafe is everywhere in our code and that is not the Rust way of doing things. Rather than scattering unsafe here and there, we create safe abstractions which internally use unsafe code, whose safety has been manually verified.

What we'll be doing now is creating safe abstractions to work with UEFI.

We start from here:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // Getting the pointer to the Boot Services from the System Table
    let boot_services = unsafe { (*sys_table).boot_services };
    
    // ... Others
}

Our code begins with dereferencing the sys_table raw pointer to get the boot_services. The problem with raw pointers is that Rust can't verify that what we're interpreting the bits as is what we're supposed to be interpreting them as. But we know for sure that it is correct. We've been able to manually verify that it is correct from the UEFI spec. Now, that we know it is correct, we can make our code safer by using a reference to SystemTable rather than a raw pointer:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // Using a reference instead of a raw pointer
    let sys_table = unsafe { &*sys_table };
    // Getting the pointer to the Boot Services from the System Table
    let boot_services = sys_table.boot_services;

    // ... Others
}

Although there is still unsafety here when we dereference the raw pointer (we'll deal with this later), from that line on, we can be sure that the bits behind sys_table will surely be an actual SystemTable instance and not a misinterpretation. This is because the &*sys_table dereferences the sys_table raw pointer (whose correctness we've manually verified) and immediately obtains a reference to those bits behind the pointer.

Just in case you didn't remember, the reason we can assign to sys_table the way we did here is because of variable overshadowing. The let sys_table = unsafe { &*sys_table } creates a new variable whose name is sys_table and the resulting value of unsafe { &*sys_table } is assigned to that new variable.

We now refactor the rest of our code to reflect the safeness of our sys_table:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // ... Others
    if locate_gop_status != STATUS_SUCCESS {
        // DELETED: let simple_text_output = unsafe { (*sys_table).simple_text_output };
        let simple_text_output = sys_table.simple_text_output;
        pre_graphics_print_str(simple_text_output, "Failed to locate GOP\n");
        loop {}
    }

    // ... Others

    for mode_number in 0..max_mode {
        // ... Others
        if query_status != STATUS_SUCCESS {
            // DELETED: let simple_text_output = unsafe { (*sys_table).simple_text_output };
            let simple_text_output = sys_table.simple_text_output;
            pre_graphics_print_str(simple_text_output, "Failed to locate GOP\n");
            loop {}
        }

        // ... Others
        if mode_number == max_mode - 1 {
            // DELETED: let simple_text_output = unsafe { (*sys_table).simple_text_output };
            let simple_text_output = sys_table.simple_text_output;
            pre_graphics_print_str(simple_text_output, "Failed to locate GOP\n");
            loop {}
        }
    }

    // ... Others
    if set_mode_status != STATUS_SUCCESS {
        // DELETED: let simple_text_output = unsafe { (*sys_table).simple_text_output };
        let simple_text_output = sys_tablesimple_text_output;
        pre_graphics_print_str(simple_text_output, "Failed to set the desired mode\n");
        loop {}
    }
    
    // ... Others
}

Our next instances of unsafety:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // ... Others

    let boot_services = sys_table.boot_services;

    // ... Others

    let locate_gop_status =
        unsafe { ((*boot_services).locate_protocol)(guid_ptr, registration, gop_ptr) };    

    // ... Others
    let gop = gop as *mut GraphicsOutput;
    let mode = unsafe { (*gop).mode };
    let max_mode = unsafe { (*mode).max_mode };

    // ... Others

    for mode_number in 0..max_mode {
        // ... Others
        let query_mode = unsafe { (*gop).query_mode };

        // ... Others

        let horizontal_resolution = unsafe { (*mode).horizontal_resolution };
        let vertical_resolution = unsafe { (*mode).vertical_resolution };
        let pixel_format = unsafe { (*mode).pixel_format };

        // ... Others
    }
    
    let set_mode_status = unsafe { ((*gop).set_mode)(gop, desired_mode) };

    // ... Others

    let framebuffer_base = unsafe { (*mode).framebuffer_base };

    // ... Others
}

Looking at that let boot_services = sys_table.boot_services;, this line returns a pointer to the BootServices instance which we'll inevitably have to dereference to use. Rather than return a pointer for dereferencing anywhere, we can create a safe function to insulate us from this potentially dangerous operation:

#[repr(C)]
struct SystemTable {
    unneeded: [u8; 60],
    simple_text_output: *mut SimpleTextOutput,
    unneeded2: [u8; 24],
    boot_services: *const BootServices,
}

// NEW:
impl SystemTable {
    // Returns a reference to the Boot Services instance in the System Table
    fn boot_services(&self) -> &BootServices {
        unsafe { &*self.boot_services }
    }
}

This way, instead of having to use the pointer directly and unsafely dereferencing it here and there, the unsafe operation of dereferencing the raw pointer is simply and clearly done correctly in this safe function. This function then returns a reference, which we can use safely.

Refactoring to use this safe function:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // ... Others

    // DELETED: let boot_services = sys_table.boot_services;
    // Getting a reference to the Boot Services from the System Table
    let boot_services = sys_table.boot_services();

    // ... Others

    let locate_gop_status = (boot_services.locate_protocol)(guid_ptr, registration, gop_ptr);
    
    // ... Others
}

As for our gop pointer:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // ... Others

    let gop = gop as *mut GraphicsOutput;

    // NEW:
    // Getting a reference to the GOP instance
    let gop = unsafe { &*gop };
    
    // ... Others
}

With our gop reference, we now refactor to reflect the safeness of gop:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // ... Others
    let gop = gop as *mut GraphicsOutput;
    let gop = unsafe { &*gop };
    // DELETED: let mode = unsafe { (*gop).mode };
    let mode = gop.mode;
    let max_mode = unsafe { (*mode).max_mode };

    // ... Others

    for mode_number in 0..max_mode {
        // ... Others

        // DELETED: let query_mode = unsafe { (*gop).query_mode };
        let query_mode = gop.query_mode;

        // ... Others
    }
    
    // DELETED: let set_mode_status = unsafe { ((*gop).set_mode)(gop, desired_mode) };
    let set_mode_status = (gop.set_mode)(gop, desired_mode);

    // ... Others

    let framebuffer_base = unsafe { (*mode).framebuffer_base };

    // ... Others
}

At this point, for our code to work, there's one more thing we need to change:

#[repr(C)]
struct GraphicsOutput {
    query_mode: extern "efiapi" fn(
        this: *mut GraphicsOutput,
        mode_number: u32,
        size_of_info: *const usize,
        info: *mut *const GraphicsModeInfo,
    ) -> Status,
    set_mode: extern "efiapi" fn(this: *mut GraphicsOutput, mode_number: u32) -> Status,
    unneeded: [u8; 8],
    mode: *const GraphicsMode,
}

This is the current definition of our GraphicsOutput. The functions take mutable pointers to GraphicsOutput as their first arguments. In our code, we're no longer dealing with raw mutable pointers to GraphicsOutput. We're now using references to it. we rectify this like so:

#[repr(C)]
struct GraphicsOutput {
    query_mode: extern "efiapi" fn(
        // DELETED: this: *mut GraphicsOutput,
        // NEW:
        // A reference to the `GraphicsOutput` instance
        this: &GraphicsOutput,
        mode_number: u32,
        size_of_info: *const usize,
        info: *mut *const GraphicsModeInfo,
    ) -> Status,
    // DELETED: set_mode: extern "efiapi" fn(this: *mut GraphicsOutput, mode_number: u32) -> Status,
    // NEW:
    set_mode: extern "efiapi" fn(this: &GraphicsOutput, mode_number: u32) -> Status,
    unneeded: [u8; 8],
    mode: *const GraphicsMode,
}

We can do this because references are also just pointers.

As for the mode pointer, we dereference it to get information from the GraphicsMode instance. This is because we retrieve a pointer to GraphicsMode from the GOP, rather than a reference to it:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // ... Others

    let mode = gop.mode;

    // ... Others
}

The let mode = gop.mode assigns a raw pointer to a GraphicsMode instance to mode. This raw pointer will inevitably have to be dereferenced to get the information behind it. Rather than getting a raw pointer and dereferencing everywhere, we can do the raw pointer dereferencing once in a small area in our code, within a safe function that then returns a reference to GraphicsMode instead of a raw pointer:

#[repr(C)]
struct GraphicsOutput {
    query_mode: extern "efiapi" fn(
        this: &GraphicsOutput,
        mode_number: u32,
        size_of_info: *const usize,
        info: *mut *const GraphicsModeInfo,
    ) -> Status,
    set_mode: extern "efiapi" fn(this: &GraphicsOutput, mode_number: u32) -> Status,
    unneeded: [u8; 8],
    mode: *const GraphicsMode,
}

// NEW:
impl GraphicsOutput {
    // Retrieves a reference to the GraphicsMode
    fn mode(&self) -> &GraphicsMode {
        unsafe { &*self.mode }
    }
}

Refactoring to use this function:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // ... Others

    // DELETED: let mode = gop.mode;
    // NEW:
    let mode = gop.mode();

    // ... Others
}

And to reflect the new safety of our mode:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // ... Others

    let mode = gop.mode();

    // DELETED: let max_mode = unsafe { (*mode).max_mode };
    // NEW:
    let max_mode = mode.max_mode;

    // ... Others

    // DELETED: let framebuffer_base = unsafe { (*mode).framebuffer_base };
    // NEW:
    let framebuffer_base = mode.framebuffer_base;
}

Within our for loop, we have a load of unsafe too:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // ... Others

    for mode_number in 0..max_mode {
        // ... Others

        let horizontal_resolution = unsafe { (*mode).horizontal_resolution };
        let vertical_resolution = unsafe { (*mode).vertical_resolution };
        let pixel_format = unsafe { (*mode).pixel_format };

        // ... Others
    }

    // ... Others
}

We resolve this:

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    // ... Others

    for mode_number in 0..max_mode {
        // ... Others
        
        // NEW:
        let mode = unsafe { &*mode };

        /* DELETED:
        let horizontal_resolution = unsafe { (*mode).horizontal_resolution };
        let vertical_resolution = unsafe { (*mode).vertical_resolution };
        let pixel_format = unsafe { (*mode).pixel_format };
        */

        let horizontal_resolution = mode.horizontal_resolution;
        let vertical_resolution = mode.vertical_resolution;
        let pixel_format = mode.pixel_format;

        // ... Others
    }

    // ... Others
}

At this point, we've seriously reduced the amount of unsafety in our code, but we aren't done yet.

Take Away

Code till now:

Directory view:

blasterball/
| .cargo/
| | config.toml
| src/
| | font.rs
| | main.rs
| .gitignore
| Cargo.lock
| Cargo.toml

main.rs contents

// ... Before efi_main

#[no_mangle]
extern "efiapi" fn efi_main(
    handle: *const core::ffi::c_void,
    sys_table: *mut SystemTable,
) -> usize {
    let sys_table = unsafe { &*sys_table };
    let boot_services = sys_table.boot_services();
    let mut gop: *mut core::ffi::c_void = core::ptr::null_mut();
    let guid_ptr = &GOP_GUID as *const Guid;
    let registration = core::ptr::null_mut();
    let gop_ptr = &mut gop as *mut _;
    let locate_gop_status = (boot_services.locate_protocol)(guid_ptr, registration, gop_ptr);

    if locate_gop_status != STATUS_SUCCESS {
        let simple_text_output = sys_table.simple_text_output;
        pre_graphics_print_str(simple_text_output, "Failed to locate GOP\n");
        loop {}
    }

    let gop = gop as *mut GraphicsOutput;
    let gop = unsafe { &*gop };
    let mode = gop.mode();
    let max_mode = mode.max_mode;
    let mut desired_mode = 0;
    for mode_number in 0..max_mode {
        let size_of_info = core::mem::size_of::<GraphicsModeInfo>();
        let mut mode: *const GraphicsModeInfo = core::ptr::null_mut();
        let query_mode = gop.query_mode;
        let query_status = (query_mode)(
            gop,
            mode_number,
            &size_of_info as *const _,
            &mut mode as *mut _,
        );
        if query_status != STATUS_SUCCESS {
            let simple_text_output = sys_table.simple_text_output;
            pre_graphics_print_str(simple_text_output, "Failed to locate GOP\n");
            loop {}
        }

        let mode = unsafe { &*mode };
        let horizontal_resolution = mode.horizontal_resolution;
        let vertical_resolution = mode.vertical_resolution;
        let pixel_format = mode.pixel_format;
        if horizontal_resolution == 640
            && vertical_resolution == 480
            && pixel_format == PixelFormat::BlueGreenRedReserved
        {
            desired_mode = mode_number;
            break;
        }
        if mode_number == max_mode - 1 {
            let simple_text_output = sys_table.simple_text_output;
            pre_graphics_print_str(simple_text_output, "Failed to locate GOP\n");
            loop {}
        }
    }

    let set_mode_status = (gop.set_mode)(gop, desired_mode);

    if set_mode_status != STATUS_SUCCESS {
        let simple_text_output = sys_table.simple_text_output;
        pre_graphics_print_str(simple_text_output, "Failed to set the desired mode\n");
        loop {}
    }

    let framebuffer_base = mode.framebuffer_base;

    let screen = framebuffer_base as *mut Screen;

    print_str(screen, "Hello World!");

    0
}

// ... Others

#[repr(C)]
struct SystemTable {
    unneeded: [u8; 60],
    simple_text_output: *mut SimpleTextOutput,
    unneeded2: [u8; 24],
    boot_services: *const BootServices,
}

impl SystemTable {
    fn boot_services(&self) -> &BootServices {
        unsafe { &*self.boot_services }
    }
}

// ... Others

#[repr(C)]
struct GraphicsOutput {
    query_mode: extern "efiapi" fn(
        this: &GraphicsOutput,
        mode_number: u32,
        size_of_info: *const usize,
        info: *mut *const GraphicsModeInfo,
    ) -> Status,
    set_mode: extern "efiapi" fn(this: &GraphicsOutput, mode_number: u32) -> Status,
    unneeded: [u8; 8],
    mode: *const GraphicsMode,
}

impl GraphicsOutput {
    fn mode(&self) -> &GraphicsMode {
        unsafe { &*self.mode }
    }
}

// ... Others

In the Next Post

We'll continue refactoring our code