yggdrasil/userspace/strace/src/format.rs

use std::fmt::Write;

use runtime::{abi::SyscallFunction, abi_lib::SyscallRegister, rt::io::{OpenOptions, SeekFrom}};

use crate::tracer::{CommandTracer, SyscallTrace};

enum Arg {
    Fd,
    OptFd,
    OptPid,
    InByteSlice,
    OutByteSlice,
    OpenOptions,
    Hex,
    Dec,
    Oct,
    Bool,
    ThinPointer,
    U64Dec,
    Seek,
    Todo,
}

#[derive(Clone, Copy)]
enum Result {
    ResUnit,
    ResFd,
    ResPid,
    ResDec,
    ResHex,
    ResU32,
    None,
}

#[derive(Clone, Copy)]
struct Formatter {
    fun: &'static str,
    args: &'static [Arg],
    res: Result
}

impl Formatter {
    const EMPTY: Self = Self {
        fun: "",
        args: &[],
        res: Result::None
    };
}

macro_rules! formatter_list {
    (
        [
            $($syscall:ident => $name:ident ($($arg:ident),* $(,)?) $(-> $result:ident)?),+ $(,)?
        ]
    ) => {
        const FORMATTERS: [Option<Formatter>; SYSCALL_MAX] = const {
            let mut array = [None; SYSCALL_MAX];

            $(
                array[SyscallFunction::$syscall as usize] = Some(Formatter {
                    fun: stringify!($name),
                    args: &[$(Arg::$arg),*],
                    $(res: Result::$result,)?
                    ..Formatter::EMPTY
                });
            )+

            array
        };
    };
}

// TODO somehow sync this up with abi definition
const SYSCALL_MAX: usize = 256;
formatter_list!([
    // System
    GetRandom => get_random(Todo),
    GetClock => get_clock(Todo) -> ResUnit,
    Mount => mount(Todo) -> ResUnit,
    Unmount => unmount(Todo) -> ResUnit,
    LoadModule => load_module(OutByteSlice) -> ResUnit,
    FilesystemControl => filesystem_control(OptFd, Hex, Todo) -> ResDec,
    GetSystemInfo => get_system_info(Hex, Todo) -> ResDec,

    // Memory management
    MapMemory => map_memory(ThinPointer, Hex, Todo) -> ResHex,
    UnmapMemory => unmap_memory(ThinPointer, Hex) -> ResUnit,

    // Process/thread
    SpawnThread => spawn_thread(Todo) -> ResU32,
    WaitThread => wait_thread(Dec, Todo) -> ResUnit,
    GetThreadOption => get_thread_option(Hex, Todo) -> ResDec,
    SetThreadOption => set_thread_option(Hex, Todo) -> ResUnit,
    GetProcessOption => get_process_option(OptPid, Hex, Todo) -> ResDec,
    SetProcessOption => set_process_option(OptPid, Hex, Todo) -> ResUnit,

    Nanosleep => nanosleep(Todo) -> ResUnit,
    ExitSignal => exit_signal(Todo),
    SendSignal => send_signal(Dec, Todo) -> ResUnit,

    Mutex => mutex(ThinPointer, Todo) -> ResUnit,
    StartSession => start_session() -> ResUnit,

    // I/O
    Open => open(OptFd, OutByteSlice, OpenOptions, Oct) -> ResFd,
    Close => close(Fd) -> ResUnit,
    Read => read(Fd, InByteSlice) -> ResDec,
    Write => write(Fd, OutByteSlice) -> ResDec,
    Seek => seek(Fd, Seek, Todo) -> ResUnit,
    Truncate => truncate(Fd, U64Dec) -> ResUnit,
    Fsync => fsync(Fd, Todo) -> ResUnit,
    ReadAt => read_at(Fd, U64Dec, InByteSlice) -> ResDec,
    WriteAt => write_at(Fd, U64Dec, OutByteSlice) -> ResDec,
    GetFileOption => get_file_option(Fd, Hex, Todo) -> ResDec,
    SetFileOption => set_file_option(Fd, Hex, Todo) -> ResUnit,
    OpenDirectory => open_directory(OptFd, OutByteSlice) -> ResFd,
    ReadDirectoryEntries => read_directory_entries(Fd, Todo) -> ResDec,
    CreateDirectory => create_directory(OptFd, OutByteSlice, Oct) -> ResUnit,
    CreateSymlink => create_symlink(OptFd, OutByteSlice, OutByteSlice) -> ResUnit,
    ReadLink => read_link(OptFd, OutByteSlice, InByteSlice) -> ResDec,
    Remove => remove(OptFd, OutByteSlice, Todo) -> ResUnit,
    Rename => rename(Todo) -> ResUnit,
    CloneFd => clone_fd(Fd, OptFd) -> ResFd,
    UpdateMetadata => update_metadata(OptFd, OutByteSlice, Todo) -> ResUnit,
    GetMetadata => get_metadata(OptFd, OutByteSlice, Todo, Bool) -> ResUnit,
    DeviceRequest => device_request(Fd, Hex, Todo) -> ResDec,

    // Misc I/O
    CreateTimer => create_timer(Todo) -> ResFd,
    CreatePid => create_pid(Todo) -> ResFd,
    CreatePty => create_pty(Todo) -> ResFd,
    CreateSharedMemory => create_shared_memory(Hex) -> ResFd,
    CreatePipe => create_pipe(Todo) -> ResFd,
    PollChannelWait => poll_channel_wait(Fd, Todo) -> ResUnit,
    PollChannelControl => poll_channel_control(Fd, Todo) -> ResUnit,

    // Network
    CreateSocket => create_socket(Todo) -> ResFd,
    Bind => bind(Fd, Todo) -> ResUnit,
    Listen => listen(Fd) -> ResUnit,
    Connect => connect(Fd, Todo) -> ResUnit,
    Accept => accept(Fd, Todo) -> ResFd,
    Shutdown => shutdown(Fd, Todo) -> ResUnit,
    SendTo => send_to(Fd, OutByteSlice, Todo) -> ResDec,
    ReceiveFrom => receive_from(Fd, InByteSlice, Todo) -> ResDec,
    GetSocketOption => get_socket_option(Fd, Hex, Todo) -> ResDec,
    SetSocketOption => set_socket_option(Fd, Hex, Todo) -> ResUnit,
    SendMessage => send_message(Fd, Todo) -> ResDec,
    ReceiveMessage => receive_message(Fd, Todo) -> ResDec,

    // C compat
    Fork => fork() -> ResPid,
    Execve => execve(Todo) -> ResUnit,

    // Debugging
    DebugTrace => debug_trace(Dec, OutByteSlice),
    DebugControl => debug_control(Dec, Hex, Todo) -> ResDec,
]);

impl Formatter {
    fn print(&self, tracer: &mut CommandTracer, args: &[usize], result: usize) {
        print!("{}(", self.fun);
        let mut arg_index = 0;
        for (i, arg) in self.args.iter().enumerate() {
            if i != 0 {
                print!(", ");
            }

            let Some((arg, len)) = arg.format(tracer, &args[arg_index..], result) else {
                print!("<???>");
                break;
            };
            arg_index += len;

            print!("{arg}")
        }
        print!(")");
        if let Some(result) = self.res.format(result) {
            print!(" -> {result}");
        }
        println!();
    }
}

const MAX_BYTES: usize = 32;

fn format_byte_slice(bytes: &[u8], real_len: usize, allow_ascii: bool) -> String {
    let mut result = String::new();
    if bytes.is_empty() {
        return "<empty>".into();
    }
    if bytes[0].is_ascii_graphic() && allow_ascii {
        // Format as ASCII string
        write!(result, "\"").ok();
        for &byte in bytes {
            let ch = byte.escape_ascii();
            write!(result, "{ch}").ok();
        }
        if bytes.len() < real_len {
            write!(result, "...\" +{}", real_len - bytes.len()).ok();
        } else {
            write!(result, "\"").ok();
        }
    } else {
        write!(result, "[").ok();
        for (i, &byte) in bytes.iter().enumerate() {
            if i != 0 {
                write!(result, ", ").ok();
            }
            write!(result, "{byte:#04X}").ok();
        }
        if bytes.len() < real_len {
            write!(result, "... +{}", real_len - bytes.len()).ok();
        }
        write!(result, "]").ok();
    }
    result
}

fn format_open_options(opts: OpenOptions) -> String {
    const OPTS: &[(&str, OpenOptions)] = &[
        ("READ", OpenOptions::READ),
        ("WRITE", OpenOptions::WRITE),
        ("APPEND", OpenOptions::APPEND),
        ("CREATE", OpenOptions::CREATE),
        ("CREATE_EXCL", OpenOptions::CREATE_EXCL),
        ("TRUNCATE", OpenOptions::TRUNCATE),
    ];
    let mut result = String::new();

    let mut i = 0;
    for &(name, opt) in OPTS {
        if opts.contains(opt) {
            if i != 0 {
                write!(result, "|").ok();
            }
            write!(result, "{name}").ok();
            i += 1;
        }
    }

    if result.is_empty() {
        "<empty>".into()
    } else {
        result
    }
}

impl Arg {
    fn format(
        &self,
        tracer: &mut CommandTracer,
        args: &[usize],
        result: usize,
    ) -> Option<(String, usize)> {
        match self {
            Arg::Dec => Some((format!("{}", args[0]), 1)),
            Arg::Hex => Some((format!("{:#x}", args[0]), 1)),
            Arg::Oct => Some((format!("{:#o}", args[0]), 1)),
            #[cfg(any(rust_analyzer, target_pointer_width = "64"))]
            Arg::U64Dec => {
                Some((format!("{}", args[0]), 1))
            }
            Arg::Bool => {
                let result = if args[0] == 0 {
                    "false".into()
                } else {
                    "true".into()
                };
                Some((result, 1))
            }
            Arg::Fd => Some((format!("{}", args[0]), 1)),
            Arg::OptPid => {
                let result = if args[0] == usize::MAX {
                    "<self>".into()
                } else {
                    format!("#{}", args[0])
                };
                Some((result, 1))
            }
            Arg::OptFd => Some((
                if args[0] == usize::MAX {
                    "<none>".into()
                } else {
                    format!("{}", args[0])
                },
                1,
            )),
            Arg::OutByteSlice => {
                let mut len = args[1];
                if len > MAX_BYTES {
                    len = MAX_BYTES;
                }
                let mut buffer = vec![0; len];
                if tracer.read_memory(args[0], &mut buffer).is_ok() {
                    Some((format_byte_slice(&buffer, args[1], true), 2))
                } else {
                    Some((format!("<{} bytes @ {:#x}>", args[1], args[0]), 2))
                }
            }
            Arg::InByteSlice => {
                let mut len = result;
                if len > MAX_BYTES {
                    len = MAX_BYTES;
                }
                let mut buffer = vec![0; len];
                if tracer.read_memory(args[0], &mut buffer).is_ok() {
                    Some((format_byte_slice(&buffer, result, true), 2))
                } else {
                    Some((
                        format!("<{}/{} bytes @ {:#x}>", result, args[1], args[0]),
                        2,
                    ))
                }
            }
            Arg::OpenOptions => {
                let opts = unsafe { OpenOptions::from_raw(args[0] as u32) };
                Some((format_open_options(opts), 1))
            }
            Arg::ThinPointer => {
                let result = if args[0] == 0 {
                    "<null>".into()
                } else {
                    format!("{:#x}", args[0])
                };
                Some((result, 1))
            }
            Arg::Seek => {
                // TODO this will get encoded as two words in i686
                let seek = SeekFrom::from_syscall_register(args[0]);
                let result = match seek {
                    SeekFrom::Start(pos) => format!("={pos}"),
                    SeekFrom::End(pos) => format!("=END{:+}", -pos),
                    SeekFrom::Current(pos) => format!("{pos:+}"),
                };
                Some((result, 1))
            }
            Arg::Todo => {
                Some(("...".into(), 1))
            }
        }
    }
}

impl Result {
    fn format(&self, result: usize) -> Option<String> {
        Some(match self {
            Self::ResDec | Self::ResFd | Self::ResU32 | Self::ResPid => {
                if result & (1 << 63) != 0 {
                    "<error>".into()
                } else {
                    format!("{result}")
                }
            }
            Self::ResHex => {
                if result & (1 << 63) != 0 {
                    "<error>".into()
                } else {
                    format!("{result:#x}")
                }
            }
            Self::ResUnit => {
                if result == 0 {
                    "<ok>"
                } else {
                    "<error>"
                }.into()
            }
            _ => return None
        })
    }
}

pub fn format_syscall(tracer: &mut CommandTracer, syscall: &SyscallTrace) {
    // Ignore mutex calls for now, they're too noisy
    if syscall.function == SyscallFunction::Mutex as usize {
        return;
    }

    if let Some(Some(formatter)) = FORMATTERS.get(syscall.function) {
        formatter.print(tracer, &syscall.args, syscall.result);
    }
}