k-os/crates/drivers/input/virtio/src/lib.rs

#![no_std]

use hw_virtio::mmio::MmioAccess;
use port_input_device::{InputDevice, InputEvent, Key};

const INPUT_MMIO_VA: usize = 0x25000;
const VQUEUE_VA: usize = 0x26000;
const EVBUF_VA: usize = 0x28000;

const QUEUE_SIZE: usize = 32;
const EVENT_SIZE: usize = 8;

const EV_KEY: u16 = 1;
const KEY_PRESS: u32 = 1;

struct MmioInput;

impl hw_virtio::mmio::MmioAccess for MmioInput {
    fn read(&self, offset: usize) -> u32 {
        unsafe { core::ptr::read_volatile((INPUT_MMIO_VA + offset) as *const u32) }
    }
    fn write(&self, offset: usize, val: u32) {
        unsafe { core::ptr::write_volatile((INPUT_MMIO_VA + offset) as *mut u32, val) }
    }
}

pub struct VirtioInputDevice {
    last_used: u16,
    next_avail: u16,
    evbuf_pa: usize,
}

impl VirtioInputDevice {
    /// Takes over the virtio input device behind `mmio`. `None` if it is
    /// not one or any step fails.
    pub fn try_new(mmio: types::bounded::CapSlot) -> Option<Self> {
        libkos::device::map_mmio(mmio, INPUT_MMIO_VA).ok()?;
        let regs = MmioInput;
        if regs.read(hw_virtio::mmio::REG_MAGIC) != hw_virtio::mmio::MAGIC
            || regs.read(hw_virtio::mmio::REG_DEVICE_ID) != hw_virtio::mmio::DEVICE_ID_INPUT
        {
            return None;
        }
        let vq_pa = libkos::device::dma(2, VQUEUE_VA).ok()?;
        let evbuf_pa = libkos::device::dma(1, EVBUF_VA).ok()?;
        hw_virtio::mmio::init_legacy_device(&regs, vq_pa, QUEUE_SIZE as u32);

        let mut dev = Self { last_used: 0, next_avail: 0, evbuf_pa };
        dev.fill_queue();
        Some(dev)
    }

    fn fill_queue(&mut self) {
        let evbuf_phys = self.evbuf_pa;
        let desc_base = VQUEUE_VA;
        let avail_base = VQUEUE_VA + QUEUE_SIZE * 16;

        for i in 0..QUEUE_SIZE {
            let desc = desc_base + i * 16;
            let buf_pa = evbuf_phys + i * EVENT_SIZE;
            unsafe {
                core::ptr::write_volatile(desc as *mut u64, buf_pa as u64);
                core::ptr::write_volatile((desc + 8) as *mut u32, EVENT_SIZE as u32);
                core::ptr::write_volatile((desc + 12) as *mut u16, 2); // WRITE flag
                core::ptr::write_volatile((desc + 14) as *mut u16, 0); // no next
            }

            unsafe {
                let ring_entry = avail_base + 4 + (i % QUEUE_SIZE) * 2;
                core::ptr::write_volatile(ring_entry as *mut u16, i as u16);
            }
        }

        unsafe {
            core::sync::atomic::fence(core::sync::atomic::Ordering::SeqCst);
            core::ptr::write_volatile((avail_base + 2) as *mut u16, QUEUE_SIZE as u16);
            core::sync::atomic::fence(core::sync::atomic::Ordering::SeqCst);
        }

        let mmio = MmioInput;
        mmio.write(hw_virtio::mmio::REG_QUEUE_NOTIFY, 0);
        self.next_avail = QUEUE_SIZE as u16;
    }

    fn resubmit(&mut self, desc_idx: u16) {
        let avail_base = VQUEUE_VA + QUEUE_SIZE * 16;
        let slot = self.next_avail as usize % QUEUE_SIZE;
        unsafe {
            core::ptr::write_volatile((avail_base + 4 + slot * 2) as *mut u16, desc_idx);
            core::sync::atomic::fence(core::sync::atomic::Ordering::SeqCst);
            self.next_avail = self.next_avail.wrapping_add(1);
            core::ptr::write_volatile((avail_base + 2) as *mut u16, self.next_avail);
            core::sync::atomic::fence(core::sync::atomic::Ordering::SeqCst);
        }
        let mmio = MmioInput;
        mmio.write(hw_virtio::mmio::REG_QUEUE_NOTIFY, 0);
    }

    fn read_event(&self, desc_idx: u16) -> (u16, u16, u32) {
        let buf_va = EVBUF_VA + (desc_idx as usize) * EVENT_SIZE;
        unsafe {
            let ev_type = core::ptr::read_volatile(buf_va as *const u16);
            let ev_code = core::ptr::read_volatile((buf_va + 2) as *const u16);
            let ev_value = core::ptr::read_volatile((buf_va + 4) as *const u32);
            (ev_type, ev_code, ev_value)
        }
    }

    fn used_idx(&self) -> u16 {
        let used_base = VQUEUE_VA + 4096;
        unsafe { core::ptr::read_volatile((used_base + 2) as *const u16) }
    }

    fn used_entry_id(&self, idx: u16) -> u16 {
        let used_base = VQUEUE_VA + 4096;
        let entry = used_base + 4 + (idx as usize % QUEUE_SIZE) * 8;
        unsafe { core::ptr::read_volatile(entry as *const u32) as u16 }
    }
}

fn linux_keycode_to_key(code: u16) -> Option<Key> {
    match code {
        17 => Some(Key::Up),    // KEY_W
        31 => Some(Key::Down),  // KEY_S
        30 => Some(Key::Left),  // KEY_A
        32 => Some(Key::Right), // KEY_D
        103 => Some(Key::Up),   // KEY_UP
        108 => Some(Key::Down), // KEY_DOWN
        105 => Some(Key::Left), // KEY_LEFT
        106 => Some(Key::Right),// KEY_RIGHT
        28 => Some(Key::Enter), // KEY_ENTER
        1 => Some(Key::Escape), // KEY_ESC
        57 => Some(Key::Space), // KEY_SPACE
        14 => Some(Key::Backspace), // KEY_BACKSPACE
        15 => Some(Key::Tab),   // KEY_TAB
        _ => None,
    }
}

impl InputDevice for VirtioInputDevice {
    fn poll(&mut self) -> Option<InputEvent> {
        loop {
            let current_used = self.used_idx();
            if current_used == self.last_used {
                return None;
            }

            let desc_idx = self.used_entry_id(self.last_used);
            self.last_used = self.last_used.wrapping_add(1);

            let (ev_type, ev_code, ev_value) = self.read_event(desc_idx);
            self.resubmit(desc_idx);

            if ev_type == EV_KEY {
                if let Some(key) = linux_keycode_to_key(ev_code) {
                    if ev_value == KEY_PRESS {
                        return Some(InputEvent::KeyDown(key));
                    } else if ev_value == 0 {
                        return Some(InputEvent::KeyUp(key));
                    }
                }
            }
            // Skip SYN events, repeats, and unknown keys — keep draining
        }
    }

    fn has_event(&self) -> bool {
        self.used_idx() != self.last_used
    }

    fn acknowledge(&mut self) {
        let regs = MmioInput;
        let status = regs.read(hw_virtio::mmio::REG_INTERRUPT_STATUS);
        regs.write(hw_virtio::mmio::REG_INTERRUPT_ACK, status);
    }
}