yggdrasil/userspace/tools/rdb/src/debugger.rs

use std::{
    cell::RefCell,
    collections::HashMap,
    fmt::Write,
    fs::File,
    io::{BufReader, Read, Seek, SeekFrom},
    os::{
        fd::{AsRawFd, RawFd},
        yggdrasil::{io::poll::PollChannel, process::CommandExt},
    },
    path::Path,
    process::{Child, Command},
    sync::Arc,
};

use elf::{ElfStream, endian::AnyEndian, segment::ProgramHeader, symbol::Symbol};
use libterm::{Color, Term, TermKey};
use rangemap::RangeMap;
use yggdrasil_rt::{debug::DebugFrame, process::ProcessId};

use crate::InstructionFormatter;
use crate::state::State;
use crate::{Error, Target, comm::Comm};

pub struct SymbolResolver {
    image: Arc<ImageInfo>,
    ip_offset: usize,
    ip: u64,
}

pub struct ImageInfo {
    symbol_table: HashMap<String, Symbol>,
    segment_headers: RangeMap<usize, ProgramHeader>,
    segments: RefCell<RangeMap<usize, Vec<u8>>>,
    functions: RangeMap<u64, String>,
}

pub struct Debugger<T: Target> {
    comm: Comm,
    term: Term,
    file: BufReader<File>,

    image: Arc<ImageInfo>,

    command: Option<String>,
    status: Option<String>,

    term_fd: RawFd,
    comm_fd: RawFd,
    poll: PollChannel,

    child: Child,
    child_exited: bool,

    state: Option<State<T>>,
}

impl<T: Target> Debugger<T> {
    pub fn from_command<P: AsRef<Path>>(image: P, mut command: Command) -> Result<Self, Error> {
        let image = image.as_ref();

        let image_info = read_image(image)?;

        let file = BufReader::new(File::open(image)?);
        let comm = Comm::open("rdb-1")?;
        let term = Term::open().unwrap();
        let mut poll = PollChannel::new()?;

        let comm_fd = comm.as_raw_fd();
        let term_fd = term.input_fd();

        poll.add(comm_fd)?;
        poll.add(term_fd)?;

        unsafe {
            command.attach_debugger(comm_fd);
        }

        let child = command.spawn()?;

        Ok(Self {
            file,
            comm,
            term,

            image: image_info.into(),

            command: None,
            status: None,

            term_fd,
            comm_fd,
            poll,

            child,
            child_exited: false,

            state: None,
        })
    }

    fn handle_frame(&mut self, frame: DebugFrame) -> Result<(), Error> {
        self.status = None;

        match frame {
            DebugFrame::Startup {
                image_base,
                ip_offset,
                ip,
            } => {
                let pid = unsafe { ProcessId::from_raw(self.child.id()) };
                let mut state = State::new(image_base, ip_offset, pid);
                if ip >= ip_offset {
                    state.update_ip(ip - ip_offset);
                }
                self.status = Some(format!("Attached to #{} @ {:#x}", pid, state.current_ip));
                self.state = Some(state);
                Ok(())
            }
            DebugFrame::Step { frame } => {
                let state = self.state.as_mut().unwrap();
                state.update(&frame, true)?;
                Ok(())
            }
            DebugFrame::HitBreakpoint { frame } => {
                let state = self.state.as_mut().unwrap();
                state.update(&frame, true)?;
                self.status = Some(format!("Hit breakpoint @ {:#x}", state.current_ip));
                Ok(())
            }
            DebugFrame::Exited => {
                self.child_exited = true;
                Ok(())
            }
        }
    }

    fn handle_key(&mut self, key: TermKey) -> Result<(), Error> {
        if let Some(command) = self.command.as_mut() {
            match key {
                TermKey::Char('\x7F') => {
                    if !command.is_empty() {
                        command.pop();
                    }
                    Ok(())
                }
                TermKey::Char('\n') => {
                    let command = command.clone();
                    self.command = None;
                    self.run_command(&command)?;
                    Ok(())
                }
                TermKey::Char(ch) if ch.is_ascii_graphic() || ch == ' ' => {
                    command.push(ch);
                    Ok(())
                }
                _ => Ok(()),
            }
        } else {
            match key {
                TermKey::Char(':') => {
                    self.command = Some(String::new());
                    Ok(())
                }
                TermKey::Char('q') => {
                    // TODO send exit to the child
                    // self.child.kill()?;
                    todo!();
                }
                TermKey::Char('s') => {
                    // Send resume to the debugee
                    if let Some(state) = self.state.as_mut() {
                        state.resume(true)?;
                    }
                    Ok(())
                }
                TermKey::Char('c') => {
                    if let Some(state) = self.state.as_mut() {
                        state.resume(false)?;
                    }
                    Ok(())
                }
                _ => Ok(()),
            }
        }
    }

    fn run_command(&mut self, command: &str) -> Result<(), Error> {
        let words: Vec<&str> = command.split(' ').collect();

        if words.is_empty() {
            return Ok(());
        }
        let Some(state) = self.state.as_mut() else {
            return Ok(());
        };

        match words[0] {
            "break" | "b" if words.len() == 2 => {
                let address = Self::parse_location(&self.image, state.ip_offset, words[1])?;
                if let Err(error) = state.set_breakpoint(address) {
                    self.status = Some(format!(
                        "Couldn't set breakpoint @ {:#x}: {:?}",
                        address, error
                    ));
                }
                Ok(())
            }
            "read" | "r" if words.len() == 3 => {
                let ty = match words[1] {
                    "b" => 1,
                    "w" => 2,
                    "d" => 4,
                    "q" => 8,
                    _ => return Ok(()),
                };
                let address = convert_address(words[2])?;

                let mut buf = [0; 8];
                if state.read_memory(address, &mut buf[..ty]).is_ok() {
                    let value = match ty {
                        1 => buf[0] as u64,
                        2 => u16::from_ne_bytes([buf[0], buf[1]]) as u64,
                        4 => u32::from_ne_bytes([buf[0], buf[1], buf[2], buf[3]]) as u64,
                        8 => u64::from_ne_bytes(buf),
                        _ => unreachable!(),
                    };

                    self.status = Some(format!("*{:#x} -> {:#x}", address, value));
                } else {
                    self.status = Some(format!("Could not read memory at {:#x}", address));
                }
                Ok(())
            }
            _ => Ok(()),
        }
    }

    fn disassembly(&mut self, amount: usize) -> Result<Vec<(usize, T::Instruction)>, Error> {
        let Some(state) = self.state.as_ref() else {
            return Ok(vec![]);
        };

        let mut segments = self.image.segments.borrow_mut();

        // Find segment
        let (range, segment) = match segments.get_key_value(&state.current_ip) {
            Some(seg) => seg,
            None if let Some(header) = self.image.segment_headers.get(&state.current_ip) => {
                let start = header.p_vaddr as usize;
                let end = (header.p_vaddr + header.p_memsz) as usize;
                let mut buffer = vec![0; header.p_memsz as usize];

                self.file.seek(SeekFrom::Start(header.p_offset))?;
                self.file.read_exact(&mut buffer)?;

                segments.insert(start..end, buffer);

                segments.get_key_value(&state.current_ip).unwrap()
            }
            // Outside of any segments
            None => return Ok(vec![]),
        };

        let offset_within_segment = state.current_ip - range.start;
        let upper_limit = std::cmp::min(segment.len(), offset_within_segment + amount * 8);

        T::disassemble(
            &segment[offset_within_segment..upper_limit],
            state.current_ip + state.ip_offset,
            amount,
        )
    }

    fn print_registers(&mut self, width: usize) -> Result<usize, Error> {
        const REG_WIDTH: usize = 32;

        let Some(state) = self.state.as_ref() else {
            return Ok(0);
        };

        let columns = core::cmp::max((width - 2) / REG_WIDTH, 1);
        let mut gpregs = vec![];
        T::register_list(&state.last_frame, &mut gpregs);
        let rows = 1 + gpregs.len().div_ceil(columns);

        self.term.set_cursor_position(0, 0)?;
        write!(self.term, "+ Registers ").ok();
        for _ in 13..width {
            write!(self.term, "-").ok();
        }
        write!(self.term, "+").ok();

        for (i, (name, reg)) in gpregs.into_iter().enumerate() {
            let row = i / columns;
            let col = i % columns;

            self.term
                .set_cursor_position(row + 1, col * REG_WIDTH + 1)?;
            write!(self.term, " {:<4} = {: >#18x}", name, reg).ok();
        }

        self.term.set_cursor_position(rows, 0)?;
        write!(
            self.term,
            "  flags = {}",
            T::flags_register_as_display(&state.last_frame)
        )
        .ok();

        // 1 extra row for RFLAGS
        let rows = rows + 1;

        for i in 1..rows {
            self.term.set_cursor_position(i, 0)?;
            write!(self.term, "|").ok();
            self.term.set_cursor_position(i, width - 1)?;
            write!(self.term, "|").ok();
        }

        self.term.set_cursor_position(rows, 0)?;
        write!(self.term, "+").ok();
        for _ in 2..width {
            write!(self.term, "-").ok();
        }
        write!(self.term, "+").ok();

        Ok(rows + 1)
    }

    fn redraw(&mut self) -> Result<(), Error> {
        let resolver = self.symbol_resolver();
        let (width, height) = self.term.size()?;
        self.term.clear(libterm::Clear::All)?;

        // Show register block
        let regs_rows = self.print_registers(width)?;

        if let Some((symbol, offset)) = resolver
            .as_ref()
            .and_then(SymbolResolver::resolve_current_function)
        {
            self.term.set_cursor_position(regs_rows, 1).ok();
            write!(self.term, "<{}> + {}", symbol, offset).ok();
        }

        let disassembly = self.disassembly(height - regs_rows - 2)?;

        if !disassembly.is_empty() {
            let mut formatter = T::new_instruction_formatter(resolver.unwrap());
            let mut buffer = String::new();

            // Show disassembly block
            for (i, (ip, insn)) in disassembly.into_iter().enumerate() {
                let is_current = i == 0;
                self.term.set_cursor_position(i + regs_rows + 1, 0)?;
                buffer.clear();
                formatter.format_instruction(&insn, &mut buffer);
                if is_current {
                    write!(self.term, ">").ok();

                    self.term.set_bright(true)?;
                    self.term.set_foreground(Color::Black)?;
                    self.term.set_background(Color::Green)?;

                    write!(self.term, "{:016x}  {}", ip, buffer).ok();
                } else {
                    write!(self.term, " ").ok();

                    self.term.set_foreground(Color::Blue)?;
                    write!(self.term, "{:016x}  ", ip).ok();
                    self.term.set_foreground(Color::Yellow)?;
                    write!(self.term, "{}", buffer).ok();
                }
                self.term.reset_style()?;
            }
        }

        if self.state.is_none() {
            self.term.set_foreground(Color::Yellow)?;
            write!(self.term, " Waiting for inferior process").ok();
        }

        // TODO deconflict status and command
        self.term.set_cursor_position(height - 1, 0)?;
        if let Some(command) = self.command.as_ref() {
            write!(self.term, ":{}", command).ok();
        } else if let Some(status) = self.status.as_ref() {
            write!(self.term, ">> {}", status).ok();
        }
        self.term.reset_style()?;
        self.term.flush()?;

        Ok(())
    }

    pub fn run(mut self) -> Result<(), Error> {
        while !self.child_exited {
            self.redraw()?;

            let (fd, result) = self.poll.wait(None, true)?.unwrap();
            result?;

            match fd {
                _ if fd == self.comm_fd => {
                    let frame = self.comm.recv()?;
                    self.handle_frame(frame)?;
                }
                _ if fd == self.term_fd => {
                    let key = self.term.read_key()?;
                    self.handle_key(key)?;
                }
                _ => unreachable!(),
            }
        }

        let status = self.child.wait()?;
        println!("Program exited with status {:?}", status);
        Ok(())
    }

    pub fn symbol_resolver(&self) -> Option<SymbolResolver> {
        let state = self.state.as_ref()?;
        Some(SymbolResolver {
            image: self.image.clone(),
            ip_offset: state.ip_offset,
            ip: state.current_ip as _,
        })
    }

    fn parse_location(image: &ImageInfo, ip_offset: usize, string: &str) -> Result<u64, Error> {
        // TODO validate that the breakpoint is within .text segment (or maybe delegate this
        //      validation to the kernel?)
        if let Some(sym) = image.symbol_table.get(string) {
            Ok(sym.st_value + ip_offset as u64)
        } else {
            convert_address(string)
        }
    }
}

impl SymbolResolver {
    fn resolve_symbol_inner(image: &ImageInfo, ip: u64) -> Option<(&str, usize)> {
        if let Some((range, function)) = image.functions.get_key_value(&ip) {
            assert!(range.start <= ip);
            return Some((function.as_str(), (ip - range.start).try_into().unwrap()));
        }

        None
    }

    pub fn resolve_symbol(&self, ip: u64) -> Option<(&str, usize)> {
        Self::resolve_symbol_inner(&self.image, ip)
    }

    pub fn resolve_current_function(&self) -> Option<(&str, usize)> {
        Self::resolve_symbol_inner(&self.image, self.ip)
    }

    pub fn to_image_address(&self, address: u64) -> Option<u64> {
        address.checked_sub(self.ip_offset as _)
    }
}

#[allow(unused)]
fn read_image<P: AsRef<Path>>(path: P) -> Result<ImageInfo, Error> {
    let file = BufReader::new(File::open(path)?);
    let mut elf = ElfStream::<AnyEndian, _>::open_stream(file).unwrap();

    let mut ranges = RangeMap::new();
    let mut symbols = HashMap::new();
    let mut functions = RangeMap::new();

    let (symtab, strtab) = elf.symbol_table().unwrap().unwrap();

    #[cfg(any(not(target_arch = "aarch64"), rust_analyzer))]
    for sym in symtab {
        let raw_name = strtab.get(sym.st_name as _).unwrap();
        let demangled_name = rustc_demangle::demangle(raw_name).to_string();

        if sym.st_symtype() == elf::abi::STT_FUNC && sym.st_size != 0 {
            functions.insert(
                sym.st_value..sym.st_value + sym.st_size,
                demangled_name.clone(),
            );
        }

        symbols.insert(demangled_name, sym);
    }

    for seg in elf.segments() {
        if seg.p_type != elf::abi::PT_LOAD {
            continue;
        }

        let start = seg.p_vaddr as usize;
        let end = (seg.p_vaddr + seg.p_memsz) as usize;

        ranges.insert(start..end, *seg);
    }

    Ok(ImageInfo {
        functions,
        symbol_table: symbols,
        segment_headers: ranges,
        segments: RangeMap::new().into(),
    })
}

fn convert_address(s: &str) -> Result<u64, Error> {
    if let Some(v) = s.strip_prefix("0x") {
        return u64::from_str_radix(v, 16).map_err(|_| Error::InvalidAddress(s.into()));
    }
    todo!()
}