path: root/kernel/src/console.rs

//! The console subsystem, used for the kernel to log things.

use crate::util::likely;
use contracts::{ensures, invariant};
use core::fmt::Write;
use log::Log;
use spin::Mutex;

/// The console singleton.
static CONSOLE: Console = Console(Mutex::new(ConsoleInner::new()));

/// A function pointer that may be set by the platform-specific initialization
/// code **before** hart0_boot is called. It must never be modified after that
/// point.
///
/// If present, this function is called after every write to the console with
/// the buffer contents, after which the console buffer is cleared.
#[no_mangle]
static mut CONSOLE_STRICT_FLUSH: Option<unsafe extern "C" fn(*const u8, usize)> = None;

/// Initializes the console subsystem. This must be done before any logging is
/// performed, and must
/// only be called once.
pub fn init() {
    log::set_logger(&CONSOLE).expect("failed to set logger");
    log::set_max_level(log::LevelFilter::Trace);
}

/// The point of interaction with the console. This is a singleton; see
/// `CONSOLE`.
struct Console(Mutex<ConsoleInner<4096>>);

impl Log for Console {
    fn enabled(&self, _metadata: &log::Metadata) -> bool {
        true
    }

    fn log(&self, record: &log::Record) {
        if !self.enabled(record.metadata()) {
            return;
        }

        let mut inner = self.0.lock();
        let mut body = |line| {
            let level = match record.level() {
                log::Level::Error => "\x1b[1;31mERR\x1b[0m",
                log::Level::Warn => "\x1b[1;33mWRN\x1b[0m",
                log::Level::Info => "\x1b[1;36mINF\x1b[0m",
                log::Level::Debug => "\x1b[1;35mDBG\x1b[0m",
                log::Level::Trace => "TRC",
            };
            let file = record.file().unwrap_or("???");
            let args = record.args();

            let result = if args.as_str() == Some("") {
                // A silly convenience, but don't write the extra space if we're
                // not going to write anything anyways.
                writeln!(inner, "[{level}][{file}:{line}]")
            } else {
                writeln!(inner, "[{level}][{file}:{line}] {args}")
            };

            // UNWRAP: Since the fmt::Write impl for ConsoleInner has a
            // contract promising that it won't ever return an error, this
            // should be unreachable.
            result.unwrap();
        };

        // Some contortions to avoid running afoul of the lifetime requirements
        // of format_args...
        match record.line() {
            Some(line) => body(format_args!("{line}")),
            None => body(format_args!("???")),
        }

        // Check for a strict flush handler.
        //
        // SAFETY: We document the function that should be called, as well as
        // the requirement that the function pointer cannot be concurrently
        // modified.
        unsafe {
            if let Some(console_strict_flush) = CONSOLE_STRICT_FLUSH {
                let buffer = inner.buffer();
                console_strict_flush(buffer.as_ptr(), buffer.len());
                inner.clear_buffer();
            }
        }
    }

    fn flush(&self) {}
}

/// The internals of the console. This is put behind a lock.
struct ConsoleInner<const BUFFER_SIZE: usize> {
    /// The buffer that gets written to. Wraps on overflow.
    buffer: [u8; BUFFER_SIZE],

    /// The number of bytes that have actually been written to.
    len: usize,

    /// A flag for whether the log buffer has overflowed since the last time it
    /// was cleared.
    has_overflowed: bool,
}

#[invariant(self.len <= BUFFER_SIZE)]
impl<const BUFFER_SIZE: usize> ConsoleInner<BUFFER_SIZE> {
    const fn new() -> ConsoleInner<BUFFER_SIZE> {
        ConsoleInner {
            buffer: [0; BUFFER_SIZE],
            len: 0,
            has_overflowed: false,
        }
    }

    /// Returns a reference to the filled portion of the buffer.
    fn buffer<'a>(&'a self) -> &'a [u8] {
        &self.buffer[..self.len]
    }

    /// Marks the buffer as cleared.
    fn clear_buffer(&mut self) {
        self.len = 0;
        if self.has_overflowed {
            self.has_overflowed = false;
            log::warn!("console buffer overflowed");
        }
    }

    /// Writes bytes to the buffer, wrapping and setting the flag on overflow.
    fn write(&mut self, bytes: &[u8]) {
        // Check if there's enough room for the contents of the buffer.
        let remaining = &mut self.buffer[self.len..];
        if likely(bytes.len() <= remaining.len()) {
            // If there is, copy the whole buffer in.
            remaining[..bytes.len()].copy_from_slice(bytes);
            self.len += bytes.len();
        } else if bytes.len() >= BUFFER_SIZE {
            // If not, and the bytes to write would fill up the entire buffer,
            // just write over the whole thing. (This frees us from handling
            // this case below.)
            self.buffer
                .copy_from_slice(&bytes[bytes.len() - BUFFER_SIZE..]);
            self.len = BUFFER_SIZE;
            self.has_overflowed = true;
        } else {
            // If there wasn't enough space, but there still will be enough
            // space for some of the current contents, calculate how many bytes
            // will be kept.
            //
            // This won't overflow (or even be zero), because we tested for
            // that above.
            let bytes_to_keep = BUFFER_SIZE - bytes.len();

            // Copy those bytes down to the start of the buffer.
            self.buffer
                .copy_within(self.len - bytes_to_keep..self.len, 0);

            // Copy the new bytes to the end of the buffer.
            self.buffer[bytes_to_keep..].copy_from_slice(bytes);

            // Update the length and mark the buffer as having overflowed.
            self.len = BUFFER_SIZE;
            self.has_overflowed = true;
        }
    }
}

impl<const BUFFER_SIZE: usize> Write for ConsoleInner<BUFFER_SIZE> {
    #[ensures(ret.is_ok())]
    fn write_str(&mut self, s: &str) -> core::fmt::Result {
        self.write(s.as_bytes());
        Ok(())
    }
}

#[test]
fn console_handles_overflow() {
    let mut console = ConsoleInner::<8>::new();
    assert_eq!(&console.buffer, b"\0\0\0\0\0\0\0\0");
    assert_eq!(console.len, 0);
    assert_eq!(console.has_overflowed, false);

    console.write(b"Hello");
    assert_eq!(&console.buffer, b"Hello\0\0\0");
    assert_eq!(console.len, 5);
    assert_eq!(console.has_overflowed, false);

    console.write(b", ");
    assert_eq!(&console.buffer, b"Hello, \0");
    assert_eq!(console.len, 7);
    assert_eq!(console.has_overflowed, false);

    console.write(b"world!");
    assert_eq!(&console.buffer, b", world!");
    assert_eq!(console.len, 8);
    assert_eq!(console.has_overflowed, true);
}

#[test]
fn console_handles_overflow_larger_than_buffer() {
    let mut console = ConsoleInner::<8>::new();
    assert_eq!(&console.buffer, b"\0\0\0\0\0\0\0\0");
    assert_eq!(console.len, 0);
    assert_eq!(console.has_overflowed, false);

    console.write("Hello, world!".as_bytes());
    assert_eq!(&console.buffer, b", world!");
    assert_eq!(console.len, 8);
    assert_eq!(console.has_overflowed, true);
}