RedBear-OS/src/main.rs

#![no_std]
#![cfg_attr(any(target_arch = "riscv64", target_os = "uefi"), no_main)]

extern crate alloc;

#[cfg(any(target_arch = "riscv64", target_os = "uefi"))]
#[macro_use]
extern crate uefi_std as std;

use alloc::{format, string::String, vec::Vec};
use core::{
    cmp,
    convert::TryFrom,
    fmt::{self, Write},
    mem, ptr, slice, str,
};
use redoxfs::{Disk, Node, TreeData};

use self::arch::{paging_create, paging_framebuffer};
use self::os::{Os, OsHwDesc, OsKey, OsMemoryEntry, OsMemoryKind, OsVideoMode};

#[macro_use]
mod os;

mod arch;
mod editor;
mod logger;
mod serial_16550;

const KIBI: usize = 1024;
const MIBI: usize = KIBI * KIBI;

//TODO: allocate this in a more reasonable manner
static mut AREAS: [OsMemoryEntry; 1024] = [OsMemoryEntry {
    base: 0,
    size: 0,
    kind: OsMemoryKind::Null,
}; 1024];
static mut AREAS_LEN: usize = 0;

pub fn area_add(area: OsMemoryEntry) {
    #[allow(static_mut_refs)]
    unsafe {
        for existing_area in &mut AREAS[0..AREAS_LEN] {
            if existing_area.kind == area.kind {
                if existing_area.base.unchecked_add(existing_area.size) == area.base {
                    existing_area.size += area.size;
                    return;
                }
                if area.base.unchecked_add(area.size) == existing_area.base {
                    existing_area.size += area.size;
                    existing_area.base = area.base;
                    return;
                }
            }
        }
        *AREAS.get_mut(AREAS_LEN).expect("AREAS overflowed!") = area;
        AREAS_LEN += 1;
    }
}

pub static mut KERNEL_64BIT: bool = false;

pub static mut LIVE_OPT: Option<(u64, &'static [u8])> = None;

struct SliceWriter<'a> {
    slice: &'a mut [u8],
    i: usize,
}

impl<'a> Write for SliceWriter<'a> {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        for b in s.bytes() {
            if let Some(slice_b) = self.slice.get_mut(self.i) {
                *slice_b = b;
                self.i += 1;
            } else {
                return Err(fmt::Error);
            }
        }
        Ok(())
    }
}

#[allow(dead_code)]
#[derive(Debug)]
#[repr(C, packed(8))]
pub struct KernelArgs {
    kernel_base: u64,
    kernel_size: u64,
    stack_base: u64,
    stack_size: u64,
    env_base: u64,
    env_size: u64,

    /// The base pointer to the saved RSDP.
    ///
    /// This field can be NULL, and if so, the system has not booted with UEFI or in some other way
    /// retrieved the RSDPs. The kernel or a userspace driver will thus try searching the BIOS
    /// memory instead. On UEFI systems, searching is not guaranteed to actually work though.
    acpi_rsdp_base: u64,
    /// The size of the RSDP region.
    acpi_rsdp_size: u64,

    areas_base: u64,
    areas_size: u64,

    bootstrap_base: u64,
    bootstrap_size: u64,
}

fn select_mode(
    os: &impl Os,
    output_i: usize,
    live: &mut bool,
    edit_env: &mut bool,
) -> Option<OsVideoMode> {
    const DEFAULT_WIDTH: u32 = 1280;
    const DEFAULT_HEIGHT: u32 = 720;
    const AUTOBOOT_SECONDS: usize = 5;

    let mut modes = Vec::new();
    for mode in os.video_modes(output_i) {
        let mut aspect_w = mode.width;
        let mut aspect_h = mode.height;
        for i in 2..cmp::min(aspect_w / 2, aspect_h / 2) {
            while aspect_w % i == 0 && aspect_h % i == 0 {
                aspect_w /= i;
                aspect_h /= i;
            }
        }

        modes.push((
            mode,
            format!(
                "{:>4}x{:<4} {:>3}:{:<3}",
                mode.width, mode.height, aspect_w, aspect_h
            ),
        ));
    }

    if modes.is_empty() {
        return None;
    }

    // Sort modes by pixel area, reversed
    modes.sort_by(|a, b| (b.0.width * b.0.height).cmp(&(a.0.width * a.0.height)));

    // Set selected based on Red Bear default resolution first, then best resolution fallback
    print!("Output {}", output_i);
    let mut selected = modes.first().map_or(0, |x| x.0.id);
    let mut selected_from_default = false;
    for (mode, _text) in modes.iter() {
        if mode.width == DEFAULT_WIDTH && mode.height == DEFAULT_HEIGHT {
            selected = mode.id;
            selected_from_default = true;
            print!(", default resolution: {}x{}", DEFAULT_WIDTH, DEFAULT_HEIGHT);
            break;
        }
    }
    if !selected_from_default {
        if let Some((best_width, best_height)) = os.best_resolution(output_i) {
            print!(", best resolution: {}x{}", best_width, best_height);
            for (mode, _text) in modes.iter() {
                if mode.width == best_width && mode.height == best_height {
                    selected = mode.id;
                    break;
                }
            }
        }
    }
    println!();

    println!("Arrow keys and enter select mode");
    let live_mode = os.get_text_position();
    if *live {
        println!("Press l to disable live mode");
    } else {
        println!("Press l to  enable live mode");
    }
    println!("Press e to edit boot environment");
    println!(
        "Autobooting default mode in {} seconds (press any key to cancel countdown)",
        AUTOBOOT_SECONDS
    );
    println!();
    print!(" ");

    let (off_x, off_y) = os.get_text_position();
    let rows = 12;
    let mut mode_opt = None;
    let countdown_y = off_y.saturating_sub(2);
    let mut countdown = AUTOBOOT_SECONDS;
    let mut countdown_active = true;
    while !modes.is_empty() {
        let mut row = 0;
        let mut col = 0;
        for (mode, text) in modes.iter() {
            if row >= rows {
                col += 1;
                row = 0;
            }

            os.set_text_position(off_x + col * 20, off_y + row);
            os.set_text_highlight(mode.id == selected);

            print!("{}", text);

            row += 1;
        }

        os.set_text_position(0, countdown_y);
        os.set_text_highlight(false);
        if countdown_active {
            println!(
                "Autobooting default mode in {} seconds (press any key to cancel countdown)",
                countdown
            );
        } else {
            println!("Manual mode selection active. Press Enter to boot selected mode.      ");
        }

        // Read keypress
        match if countdown_active {
            os.get_key_timeout(1000)
        } else {
            os.get_key()
        } {
            OsKey::Timeout => {
                if countdown_active {
                    if countdown == 0 {
                        if let Some(mode_i) = modes.iter().position(|x| x.0.id == selected) {
                            if let Some((mode, _text)) = modes.get(mode_i) {
                                mode_opt = Some(*mode);
                            }
                        }
                        break;
                    }
                    countdown = countdown.saturating_sub(1);
                }
            }
            OsKey::Left => {
                countdown_active = false;
                if let Some(mut mode_i) = modes.iter().position(|x| x.0.id == selected) {
                    if mode_i < rows {
                        while mode_i < modes.len() {
                            mode_i += rows;
                        }
                    }
                    mode_i -= rows;
                    if let Some(new) = modes.get(mode_i) {
                        selected = new.0.id;
                    }
                }
            }
            OsKey::Right => {
                countdown_active = false;
                if let Some(mut mode_i) = modes.iter().position(|x| x.0.id == selected) {
                    mode_i += rows;
                    if mode_i >= modes.len() {
                        mode_i %= rows;
                    }
                    if let Some(new) = modes.get(mode_i) {
                        selected = new.0.id;
                    }
                }
            }
            OsKey::Up => {
                countdown_active = false;
                if let Some(mut mode_i) = modes.iter().position(|x| x.0.id == selected) {
                    if mode_i % rows == 0 {
                        mode_i += rows;
                        if mode_i > modes.len() {
                            mode_i = modes.len();
                        }
                    }
                    mode_i -= 1;
                    if let Some(new) = modes.get(mode_i) {
                        selected = new.0.id;
                    }
                }
            }
            OsKey::Down => {
                countdown_active = false;
                if let Some(mut mode_i) = modes.iter().position(|x| x.0.id == selected) {
                    mode_i += 1;
                    if mode_i % rows == 0 {
                        mode_i -= rows;
                    }
                    if mode_i >= modes.len() {
                        mode_i = mode_i - mode_i % rows;
                    }
                    if let Some(new) = modes.get(mode_i) {
                        selected = new.0.id;
                    }
                }
            }
            OsKey::Enter => {
                countdown_active = false;
                if let Some(mode_i) = modes.iter().position(|x| x.0.id == selected) {
                    if let Some((mode, _text)) = modes.get(mode_i) {
                        mode_opt = Some(*mode);
                    }
                }
                break;
            }
            OsKey::Char('l') => {
                countdown_active = false;
                *live = !*live;
                os.set_text_position(live_mode.0, live_mode.1);
                if *live {
                    println!("Press l to disable live mode");
                } else {
                    println!("Press l to  enable live mode");
                }
            }
            OsKey::Char('e') => {
                countdown_active = false;
                if let Some(mode_i) = modes.iter().position(|x| x.0.id == selected) {
                    if let Some((mode, _text)) = modes.get(mode_i) {
                        *edit_env = true;
                        mode_opt = Some(*mode);
                    }
                }
                break;
            }
            OsKey::Other | OsKey::Backspace | OsKey::Delete | OsKey::Char(_) => {
                countdown_active = false;
            }
        }
    }

    os.set_text_position(0, off_y + rows);
    os.set_text_highlight(false);
    println!();

    mode_opt
}

fn redoxfs<O: Os>(os: &O) -> (redoxfs::FileSystem<O::D>, Option<&'static [u8]>) {
    let attempts = 10;
    for attempt in 0..=attempts {
        let mut password_opt = None;
        if attempt > 0 {
            print!("\rRedoxFS password ({}/{}): ", attempt, attempts);

            let mut password = String::new();

            loop {
                match os.get_key() {
                    OsKey::Backspace | OsKey::Delete => {
                        if !password.is_empty() {
                            print!("\x08 \x08");
                            password.pop();
                        }
                    }
                    OsKey::Char(c) => {
                        print!("*");
                        password.push(c)
                    }
                    OsKey::Enter => break,
                    _ => (),
                }
            }

            // Erase password information
            while os.get_text_position().0 > 0 {
                print!("\x08 \x08");
            }

            if !password.is_empty() {
                password_opt = Some(password);
            }
        }
        match os.filesystem(password_opt.as_ref().map(|x| x.as_bytes())) {
            Ok(fs) => {
                return (
                    fs,
                    password_opt.map(|password| {
                        // Copy password to page aligned memory
                        let password_size = password.len();
                        let password_base = os.alloc_zeroed_page_aligned(password_size);

                        area_add(OsMemoryEntry {
                            base: password_base as u64,
                            size: password_size as u64,
                            kind: OsMemoryKind::Reserved,
                        });

                        unsafe {
                            ptr::copy(password.as_ptr(), password_base, password_size);
                            slice::from_raw_parts(password_base, password_size)
                        }
                    }),
                );
            }
            Err(err) => match err.errno {
                // Incorrect password, try again
                syscall::ENOKEY => (),
                _ => {
                    panic!("Failed to open RedoxFS: {}", err);
                }
            },
        }
    }
    panic!("RedoxFS out of unlock attempts");
}

#[derive(PartialEq)]
enum Filetype {
    Elf,
    Initfs,
}
fn load_to_memory<O: Os>(
    os: &O,
    fs: &mut redoxfs::FileSystem<O::D>,
    path: &str,
    filetype: Filetype,
) -> &'static mut [u8] {
    fs.tx(|tx| {
        let mut node = None;
        for component in path.split('/') {
            node = Some(
                tx.find_node(
                    node.map_or(redoxfs::TreePtr::root(), |node: TreeData<Node>| node.ptr()),
                    component,
                )
                .unwrap_or_else(|err| panic!("Failed to find {component}: {err}")),
            );
        }
        let node = node.unwrap();

        let size = node.data().size();

        print!("{}: 0/{} MiB", path, size / MIBI as u64);

        let ptr = os.alloc_zeroed_page_aligned(size as usize);
        if ptr.is_null() {
            panic!("Failed to allocate memory for {}", path);
        }

        let slice = unsafe { slice::from_raw_parts_mut(ptr, size as usize) };

        let mut i = 0;
        for chunk in slice.chunks_mut(MIBI) {
            print!("\r{}: {}/{} MiB", path, i / MIBI as u64, size / MIBI as u64);
            i += tx
                .read_node_inner(&node, i, chunk)
                .unwrap_or_else(|err| panic!("Failed to read `{}` file: {}", path, err))
                as u64;
        }
        println!("\r{}: {}/{} MiB", path, i / MIBI as u64, size / MIBI as u64);

        if filetype == Filetype::Elf {
            let magic = &slice[..4];
            if magic != b"\x7FELF" {
                panic!("{} has invalid magic number {:#X?}", path, magic);
            }
        } else if filetype == Filetype::Initfs {
            let magic = &slice[..8];
            if magic != b"RedoxFtw" {
                panic!("{} has invalid magic number {:#X?}", path, magic);
            }
        }

        Ok(slice)
    })
    .unwrap_or_else(|err| {
        panic!(
            "RedoxFS transaction failed while loading `{}`: {}",
            path, err
        )
    })
}

fn elf_entry(data: &[u8]) -> (u64, bool) {
    match (data[4], data[5]) {
        // 32-bit, little endian
        (1, 1) => (
            u32::from_le_bytes(
                <[u8; 4]>::try_from(&data[0x18..0x18 + 4]).expect("conversion cannot fail"),
            ) as u64,
            false,
        ),
        // 32-bit, big endian
        (1, 2) => (
            u32::from_be_bytes(
                <[u8; 4]>::try_from(&data[0x18..0x18 + 4]).expect("conversion cannot fail"),
            ) as u64,
            false,
        ),
        // 64-bit, little endian
        (2, 1) => (
            u64::from_le_bytes(
                <[u8; 8]>::try_from(&data[0x18..0x18 + 8]).expect("conversion cannot fail"),
            ),
            true,
        ),
        // 64-bit, big endian
        (2, 2) => (
            u64::from_be_bytes(
                <[u8; 8]>::try_from(&data[0x18..0x18 + 8]).expect("conversion cannot fail"),
            ),
            true,
        ),
        (ei_class, ei_data) => {
            panic!("Unsupported ELF EI_CLASS {} EI_DATA {}", ei_class, ei_data);
        }
    }
}

fn main(os: &impl Os) -> (usize, u64, KernelArgs) {
    println!(
        "Redox OS Bootloader {} on {}",
        env!("CARGO_PKG_VERSION"),
        os.name()
    );

    let hwdesc = os.hwdesc();
    println!("Hardware descriptor: {:x?}", hwdesc);
    let (acpi_rsdp_base, acpi_rsdp_size) = match hwdesc {
        OsHwDesc::Acpi(base, size) => (base, size),
        OsHwDesc::DeviceTree(base, size) => (base, size),
        OsHwDesc::NotFound => (0, 0),
    };

    let (mut fs, password_opt) = redoxfs(os);

    print!("RedoxFS ");
    for i in 0..fs.header.uuid().len() {
        if i == 4 || i == 6 || i == 8 || i == 10 {
            print!("-");
        }

        print!("{:>02x}", fs.header.uuid()[i]);
    }
    println!(": {} MiB", fs.header.size() / MIBI as u64);
    println!();

    let mut mode_opts = Vec::new();
    let mut live = cfg!(feature = "live");
    let mut edit_env = false;
    for output_i in 0..os.video_outputs() {
        if output_i > 0 {
            os.clear_text();
        }
        mode_opts.push(select_mode(os, output_i, &mut live, &mut edit_env));
    }

    let stack_size = 128 * KIBI;
    let stack_base = os.alloc_zeroed_page_aligned(stack_size);
    if stack_base.is_null() {
        panic!("Failed to allocate memory for stack");
    }

    let live_opt = if live {
        let size = fs.header.size();

        print!("live: 0/{} MiB", size / MIBI as u64);

        let live_size = match usize::try_from(size) {
            Ok(live_size) => live_size,
            Err(_) => {
                println!("\rlive: disabled (image too large for bootloader address space)");
                live = false;
                0
            }
        };

        let ptr = if live {
            os.alloc_zeroed_page_aligned(live_size)
        } else {
            ptr::null_mut()
        };
        if live && ptr.is_null() {
            println!(
                "\rlive: disabled (unable to allocate {} MiB upfront)",
                size / MIBI as u64
            );
            live = false;
        }

        let live = if live {
            Some(unsafe { slice::from_raw_parts_mut(ptr, live_size) })
        } else {
            println!("Continuing without live preload");
            None
        };

        if let Some(live) = live {
            let mut i = 0;
            for chunk in live.chunks_mut(MIBI) {
                print!("\rlive: {}/{} MiB", i / MIBI as u64, size / MIBI as u64);
                i += unsafe {
                    fs.disk
                        .read_at(fs.block + i / redoxfs::BLOCK_SIZE, chunk)
                        .expect("Failed to read live disk") as u64
                };
            }
            println!("\rlive: {}/{} MiB", i / MIBI as u64, size / MIBI as u64);

            println!("Switching to live disk");
            unsafe {
                LIVE_OPT = Some((fs.block, slice::from_raw_parts_mut(ptr, live_size)));
            }

            area_add(OsMemoryEntry {
                base: live.as_ptr() as u64,
                size: live.len() as u64,
                kind: OsMemoryKind::Reserved,
            });

            Some(live)
        } else {
            None
        }
    } else {
        None
    };

    let (kernel, kernel_entry) = {
        let kernel = load_to_memory(os, &mut fs, "usr/lib/boot/kernel", Filetype::Elf);
        let (kernel_entry, kernel_64bit) = elf_entry(kernel);
        unsafe {
            KERNEL_64BIT = kernel_64bit;
        }
        (kernel, kernel_entry)
    };

    let (bootstrap_size, bootstrap_base) = {
        let initfs_slice = load_to_memory(os, &mut fs, "usr/lib/boot/initfs", Filetype::Initfs);

        let memory = unsafe {
            let total_size = initfs_slice.len().next_multiple_of(4096);
            let ptr = os.alloc_zeroed_page_aligned(total_size);
            assert!(!ptr.is_null(), "failed to allocate bootstrap+initfs memory");
            core::slice::from_raw_parts_mut(ptr, total_size)
        };
        memory[..initfs_slice.len()].copy_from_slice(initfs_slice);

        (memory.len() as u64, memory.as_mut_ptr() as u64)
    };

    let page_phys = unsafe { paging_create(os, kernel.as_ptr() as u64, kernel.len() as u64) }
        .expect("Failed to set up paging");

    let max_env_size = 64 * KIBI;
    let mut env_size = max_env_size;
    let env_base = os.alloc_zeroed_page_aligned(env_size);
    if env_base.is_null() {
        panic!("Failed to allocate memory for stack");
    }

    {
        let mut w = SliceWriter {
            slice: unsafe { slice::from_raw_parts_mut(env_base, max_env_size) },
            i: 0,
        };

        match hwdesc {
            OsHwDesc::Acpi(addr, size) => {
                writeln!(w, "RSDP_ADDR={addr:016x}").unwrap();
                writeln!(w, "RSDP_SIZE={size:016x}").unwrap();
            }
            OsHwDesc::DeviceTree(addr, size) => {
                writeln!(w, "DTB_ADDR={addr:016x}").unwrap();
                writeln!(w, "DTB_SIZE={size:016x}").unwrap();
            }
            OsHwDesc::NotFound => {}
        }

        if let Some(live) = live_opt {
            writeln!(w, "DISK_LIVE_ADDR={:016x}", live.as_ptr() as usize).unwrap();
            writeln!(w, "DISK_LIVE_SIZE={:016x}", live.len()).unwrap();
            writeln!(w, "REDOXFS_BLOCK={:016x}", 0).unwrap();
        } else {
            writeln!(w, "REDOXFS_BLOCK={:016x}", fs.block).unwrap();
        }
        write!(w, "REDOXFS_UUID=").unwrap();
        for i in 0..fs.header.uuid().len() {
            if i == 4 || i == 6 || i == 8 || i == 10 {
                write!(w, "-").unwrap();
            }

            write!(w, "{:>02x}", fs.header.uuid()[i]).unwrap();
        }
        writeln!(w).unwrap();
        if let Some(password) = password_opt {
            writeln!(
                w,
                "REDOXFS_PASSWORD_ADDR={:016x}",
                password.as_ptr() as usize
            )
            .unwrap();
            writeln!(w, "REDOXFS_PASSWORD_SIZE={:016x}", password.len()).unwrap();
        }

        #[cfg(target_arch = "riscv64")]
        {
            let boot_hartid = os::efi_get_boot_hartid()
                .expect("Could not retrieve boot hart id from EFI implementation!");
            writeln!(w, "BOOT_HART_ID={:016x}", boot_hartid).unwrap();
        }
        if edit_env {
            editor::edit_env(os, env_base, &mut w.i, max_env_size);
        }

        for output_i in 0..os.video_outputs() {
            if let Some(mut mode) = mode_opts[output_i] {
                // Set mode to get updated values
                os.set_video_mode(output_i, &mut mode);

                if output_i == 0 {
                    let virt = unsafe {
                        paging_framebuffer(
                            os,
                            page_phys,
                            mode.base,
                            (mode.stride * mode.height * 4) as u64,
                        )
                    }
                    .expect("Failed to map framebuffer");

                    writeln!(w, "FRAMEBUFFER_ADDR={:016x}", mode.base).unwrap();
                    writeln!(w, "FRAMEBUFFER_VIRT={virt:016x}").unwrap();
                    writeln!(w, "FRAMEBUFFER_WIDTH={:016x}", mode.width).unwrap();
                    writeln!(w, "FRAMEBUFFER_HEIGHT={:016x}", mode.height).unwrap();
                    writeln!(w, "FRAMEBUFFER_STRIDE={:016x}", mode.stride).unwrap();
                } else {
                    writeln!(
                        w,
                        "FRAMEBUFFER{}={:#x},{},{},{}",
                        output_i, mode.base, mode.width, mode.height, mode.stride,
                    )
                    .unwrap();
                }
            }
        }
        env_size = w.i;
    }

    #[allow(static_mut_refs)]
    (
        page_phys,
        kernel_entry,
        KernelArgs {
            kernel_base: kernel.as_ptr() as u64,
            kernel_size: kernel.len() as u64,
            stack_base: stack_base as u64,
            stack_size: stack_size as u64,
            env_base: env_base as u64,
            env_size: env_size as u64,
            acpi_rsdp_base,
            acpi_rsdp_size,
            areas_base: unsafe { AREAS.as_ptr() as u64 },
            areas_size: unsafe { (AREAS.len() * mem::size_of::<OsMemoryEntry>()) as u64 },
            bootstrap_base,
            bootstrap_size,
        },
    )
}