dev(device): implement VirtIO

See dev log for detail

13 files changed, 781 insertions, 231 deletions

diff --git a/src/drivers/plic.cc b/src/drivers/plic.cc
index a13537c..a759597 100644
--- a/src/drivers/plic.cc
+++ b/src/drivers/plic.cc
@@ -1,60 +1,59 @@
 #include "common/types.h"
 #include "kernel/riscv.h"
 #include "drivers/plic.h"
+#include "drivers/uart.h"
 
-// QEMU Virt PLIC address constants
+// Base Constants
 constexpr uint64 PLIC_BASE = 0x0c000000L;
 constexpr uint64 PLIC_PRIORITY = PLIC_BASE + 0x0;
 constexpr uint64 PLIC_PENDING = PLIC_BASE + 0x1000;
 
-// S-Mode interrupt enable register: each Hart occupies 0x100 bytes
-// Context 1 is Hart 0 S-Mode, Context 3 is Hart 1 S-Mode...
-inline uint64 PLIC_SENABLE(int hart)
-{
-    return PLIC_BASE + 0x2080 + static_cast<uint64>(hart) * 0x100;
-}
+// Helper Macros to calculate addresses
+#define PLIC_SENABLE_ADDR(hart) (PLIC_BASE + 0x2080 + (hart) * 0x100)
+#define PLIC_STHRESHOLD_ADDR(hart) (PLIC_BASE + 0x201000 + (hart) * 0x2000)
+#define PLIC_SCLAIM_ADDR(hart) (PLIC_BASE + 0x201004 + (hart) * 0x2000)
 
-// S-Mode priority threshold and claim register: 0x2000 bytes per Hart
-inline uint64 PLIC_STHRESHOLD(int hart)
+// Helper for Volatile Access
+static inline volatile uint32 *REG32(uint64 addr)
 {
-    return PLIC_BASE + 0x201000 + static_cast<uint64>(hart) * 0x2000;
+    return reinterpret_cast<volatile uint32 *>(addr);
 }
 
-inline uint64 PLIC_SCLAIM(int hart)
-{
-    return PLIC_BASE + 0x201004 + static_cast<uint64>(hart) * 0x2000;
-}
-
-namespace PILC
+namespace PLIC
 {
     void init()
     {
-        *reinterpret_cast<uint32 *>(PLIC_PRIORITY + 10 * 4) = 1; // UART
-        *reinterpret_cast<uint32 *>(PLIC_PRIORITY + 1 * 4) = 1; // VirtIO 0
+        // Set Priority for UART (IRQ 10) and VirtIO (IRQ 1)
+        // Must use volatile write
+        *REG32(PLIC_PRIORITY + 10 * 4) = 1;
+        for (int i = 1; i <= 8; i++)
+        {
+            *REG32(PLIC_PRIORITY + i * 4) = 1;
+        }
     }
 
     void inithart()
     {
-        int hart = static_cast<int>(r_tp()); // 获取当前 hartid
+        int hart = static_cast<int>(r_tp());
+
+        // Enable UART(10) and VirtIO(1) for this Hart (S-Mode)
+        *REG32(PLIC_SENABLE_ADDR(hart)) = (1 << 10) | 0x1FE;
 
-        // Enable UART and VirtIO interrupts for the current Hart's S-Mode
-        // Set bit 10 and bit 1 in the register bitmap
-        *reinterpret_cast<uint32 *>(PLIC_SENABLE(hart)) = (1 << 10) | (1 << 1);
+        // Set Priority Threshold to 0 (Allow all)
+        *REG32(PLIC_STHRESHOLD_ADDR(hart)) = 0;
 
-        // Set the current Hart's S-Mode priority threshold to 0
-        // Only interrupts with a priority higher than this threshold will be delivered
-        *reinterpret_cast<uint32 *>(PLIC_STHRESHOLD(hart)) = 0;
+        Drivers::uart_puts("[PLIC] Hart Init Done.\n");
     }
 
     int claim()
     {
         int hart = static_cast<int>(r_tp());
-        return *reinterpret_cast<uint32 *>(PLIC_SCLAIM(hart));
+        return *REG32(PLIC_SCLAIM_ADDR(hart));
     }
 
     void complete(int irq)
     {
         int hart = static_cast<int>(r_tp());
-        *reinterpret_cast<uint32 *>(PLIC_SCLAIM(hart)) = static_cast<uint32>(irq);
+        *REG32(PLIC_SCLAIM_ADDR(hart)) = static_cast<uint32>(irq);
     }
-} // namespace PILC
+}
\ No newline at end of file
diff --git a/src/drivers/uart.cc b/src/drivers/uart.cc
index fe051f8..9475d49 100644
--- a/src/drivers/uart.cc
+++ b/src/drivers/uart.cc
@@ -1,102 +1,206 @@
 #include "drivers/uart.h"
-#include "common/types.h"
-#include "kernel/config.h"
-
-// #define UART0_BASE          0x10000000L
-
-#define UART_RBR            0x00 // Receiver Buffer Register (Read)
-#define UART_THR            0x00 // Transmitter Holding Register (Write)
-#define UART_IER            0x01 // Interrupt Enable Register
-#define UART_FCR            0x02 // FIFO Control Register
-#define UART_LCR            0x03 // Line Control Register
-#define UART_LSR            0x05 // Line Status Register
-
-#define UART_LSR_RX_READY   0x01
-#define UART_LSR_TX_EMPTY   0x20
+#include "kernel/riscv.h"
+#include "kernel/spinlock.h"
+#include "kernel/proc.h"
+#include "kernel/mm.h"
+#include "lib/string.h"
 
 namespace Drivers
 {
-
-    inline volatile uint8_t &reg(uint32_t offset)
+    // Circular buffer structure
+    struct
     {
-        return *reinterpret_cast<volatile uint8_t *>(g_uart_base + offset);
-    }
+        Spinlock lock;
+        char buf[UART_RX_BUF_SIZE];
+        uint32 r; // read index
+        uint32 w; // write index
+    } uart_rx;
+
+    Spinlock uart_tx_lock;
 
     void uart_init()
     {
-        if (g_uart_base == 0)
-            return;
+        uart_rx.lock.init("uart_rx");
+        uart_tx_lock.init("uart_tx");
+        uart_rx.r = 0;
+        uart_rx.w = 0;
 
-        reg(UART_IER) = 0x00; // Disable interrupt
-        reg(UART_LCR) = 0x80; // Set baud rate (LATCH bit)
-        reg(0x00) = 0x03;     // LSB
-        reg(0x01) = 0x00;     // MSB
-        reg(UART_LCR) = 0x03; // 8 bits, no parity, one stop bit
-        reg(UART_FCR) = 0x07; // Enable FIFO, clear TX/RX queues
-    }
+        // Disable interrupts and configure
+        *(volatile uint8 *)(UART0 + IER) = 0x00;
 
-    void uart_putc(char c)
-    {
-        if (g_uart_base == 0)
-            return;
+        // Enter baud rate setting mode (LCR_DLAB) = 0x80;
+        *(volatile uint8 *)(UART0 + LCR) = 0x80;
+        *(volatile uint8 *)(UART0 + 0) = 0x03; // DLL: 38.4K baud
+        *(volatile uint8 *)(UART0 + 1) = 0x00; // DLM
 
-        if (c == '\n')
-        {
-            uart_putc('\r');
-        }
+        // Exit baud rate setting, set 8-bit word length, no parity
+        *(volatile uint8 *)(UART0 + LCR) = 0x03;
 
-        // Waiting for the send buffer to be empty
-        while ((reg(UART_LSR) & UART_LSR_TX_EMPTY) == 0)
-            ;
+        // Reset and enable FIFO
+        *(volatile uint8 *)(UART0 + FCR) = 0x07;
 
-        reg(UART_THR) = c;
+        // Enable Receive Data Available Interrupt
+        *(volatile uint8 *)(UART0 + IER) = 0x01;
     }
 
+    void uart_putc(char c)
+    {
+        uart_tx_lock.acquire();
+        while ((*(volatile uint8 *)(UART0 + LSR) & LSR_TX_IDLE) == 0)
+            ;
+        *(volatile uint8 *)(UART0 + THR) = c;
+        uart_tx_lock.release();
+    }
 
     void uart_puts(const char *s)
     {
         while (*s)
-        {
             uart_putc(*s++);
-        }
     }
 
-    void uart_put_int(int value)
+    void uart_put_int(int val)
     {
-        if (value == 0)
+        char buf[16];
+        int i = 0;
+        if (val == 0)
         {
-            uart_putc('0');
+            uart_puts("0");
             return;
         }
-
-        if (value < 0)
+        if (val < 0)
         {
-            uart_putc('-');
-            value = -value;
+            uart_puts("-");
+            val = -val;
         }
-
-        char buffer[16];
-        int i = 0;
-        while (value > 0)
+        while (val > 0)
         {
-            buffer[i++] = (value % 10) + '0';
-            value /= 10;
+            buf[i++] = (val % 10) + '0';
+            val /= 10;
         }
-
         while (i > 0)
         {
-            uart_putc(buffer[--i]);
+            uart_putc(buf[--i]);
         }
     }
 
-    void print_hex(uint64 value)
+    void print_hex(uint64 val)
     {
+        char hex[] = "0123456789ABCDEF";
         uart_puts("0x");
-        char hex[] = "0123456789abcdef";
-        for (int i = 15; i >= 0; i--)
+        for (int i = 60; i >= 0; i -= 4)
         {
-            int nibble = (value >> (i * 4)) & 0xF;
-            uart_putc(hex[nibble]);
+            uart_putc(hex[(val >> i) & 0xF]);
         }
     }
+
+    void uart_intr()
+    {
+        uart_rx.lock.acquire();
+
+        while (1)
+        {
+            if ((*(volatile uint8 *)(UART0 + LSR) & LSR_RX_READY) == 0)
+                break;
+
+            char c = *(volatile uint8 *)(UART0 + RHR);
+
+            uart_rx.buf[uart_rx.w % UART_RX_BUF_SIZE] = c;
+            uart_rx.w++;
+
+            ProcManager::wakeup(&uart_rx);
+        }
+
+        uart_rx.lock.release();
+    }
+
+    // Modern console read: Line buffering + Editing Keys + Echo
+    int console_read(uint64 target, int len)
+    {
+        struct Proc *p = myproc();
+        int target_idx = 0;
+
+        uart_rx.lock.acquire();
+
+        while (target_idx < len)
+        {
+            while (uart_rx.r == uart_rx.w)
+            {
+                if (p->xstate != 0)
+                {
+                    uart_rx.lock.release();
+                    return -1;
+                }
+                ProcManager::sleep(&uart_rx, &uart_rx.lock);
+            }
+
+            char c = uart_rx.buf[uart_rx.r % UART_RX_BUF_SIZE];
+            uart_rx.r++;
+
+            uart_rx.lock.release();
+
+            // 1. Handle Backspace (127 or \b)
+            if (c == 127 || c == '\b')
+            {
+                if (target_idx > 0)
+                {
+                    // Erase character from screen: Backspace, Space, Backspace
+                    uart_putc('\b');
+                    uart_putc(' ');
+                    uart_putc('\b');
+                    target_idx--;
+                }
+                uart_rx.lock.acquire();
+                continue;
+            }
+
+            // 2. Handle Ctrl+U (Kill Line)
+            if (c == 0x15)
+            {
+                while (target_idx > 0)
+                {
+                    uart_putc('\b');
+                    uart_putc(' ');
+                    uart_putc('\b');
+                    target_idx--;
+                }
+                uart_rx.lock.acquire();
+                continue;
+            }
+
+            // 3. Handle Ctrl+L (Clear Screen)
+            if (c == 0x0C)
+            {
+                uart_puts(ANSI_CLEAR);
+                uart_puts("$ ");
+                uart_rx.lock.acquire();
+                continue;
+            }
+
+            // 4. Handle Return (CR -> LF)
+            if (c == '\r')
+            {
+                c = '\n';
+            }
+
+            // Echo back to user
+            uart_putc(c);
+
+            if (VM::copyout(p->pagetable, target + target_idx, &c, 1) < 0)
+            {
+                return -1;
+            }
+
+            target_idx++;
+
+            if (c == '\n')
+            {
+                return target_idx;
+            }
+
+            uart_rx.lock.acquire();
+        }
+
+        uart_rx.lock.release();
+        return target_idx;
+    }
 }
\ No newline at end of file
diff --git a/src/drivers/virtio_dist.cc b/src/drivers/virtio_dist.cc
new file mode 100644
index 0000000..99d100f
--- /dev/null
+++ b/src/drivers/virtio_dist.cc
@@ -0,0 +1,338 @@
+#include "drivers/virtio.h"
+#include "kernel/riscv.h"
+#include "kernel/pmm.h"
+#include "kernel/spinlock.h"
+#include "lib/string.h"
+#include "drivers/uart.h"
+#include "kernel/proc.h"
+
+static uint64 virtio_base = 0x10001000;
+
+// VirtIO Block Device Config Structure
+struct virtio_blk_config
+{
+    uint64 capacity;
+    uint32 size_max;
+    uint32 seg_max;
+    struct virtio_blk_geometry
+    {
+        uint16 cylinders;
+        uint8 heads;
+        uint8 sectors;
+    } geometry;
+    uint32 blk_size;
+    uint8 physical_block_exp;
+    uint8 alignment_offset;
+    uint16 min_io_size;
+    uint32 opt_io_size;
+} __attribute__((packed));
+
+static inline volatile uint32 *REG(uint32 offset)
+{
+    return reinterpret_cast<volatile uint32 *>(virtio_base + offset);
+}
+
+namespace VirtIO
+{
+    struct Disk
+    {
+        struct VRingDesc *desc;
+        struct VRingAvail *avail;
+        struct VRingUsed *used;
+
+        int free_head;
+        int num_free;
+        int free_chain[NUM_DESCS];
+
+        uint16 used_idx;
+
+        struct TraceInfo
+        {
+            int busy;
+            int status;
+        } info[NUM_DESCS];
+
+        struct BlockRequest reqs[NUM_DESCS];
+        uint8 req_status[NUM_DESCS];
+
+        Spinlock lock;
+    } disk;
+
+    void init()
+    {
+        disk.lock.init("virtio_disk");
+
+        // 1. Probing for VirtIO-MMIO device
+        bool found = false;
+        // Probe common VirtIO addresses in QEMU virt machine
+        for (uint64 addr = 0x10001000; addr < 0x10009000; addr += 0x1000)
+        {
+            volatile uint32 *magic = (volatile uint32 *)(addr + VIRTIO_MMIO_MAGIC_VALUE);
+            volatile uint32 *dev_id = (volatile uint32 *)(addr + VIRTIO_MMIO_DEVICE_ID);
+
+            if (*magic == 0x74726976 && *dev_id == 2) // ID 2 = Block Device
+            {
+                virtio_base = addr;
+                found = true;
+                Drivers::uart_puts(ANSI_GREEN "[VirtIO] Found Block Device @ ");
+                Drivers::print_hex(addr);
+                Drivers::uart_puts(ANSI_RESET "\n");
+                break;
+            }
+        }
+
+        if (!found)
+        {
+            Drivers::uart_puts(ANSI_RED "[VirtIO] Fatal - No Block Device found.\n" ANSI_RESET);
+            return;
+        }
+
+        // 2. Check Device Version (We want Modern v2)
+        if (*REG(VIRTIO_MMIO_VERSION) != 2)
+        {
+            Drivers::uart_puts(ANSI_YELLOW "[VirtIO] Warning: Legacy device detected (expected v2).\n" ANSI_RESET);
+        }
+
+        // 3. Reset Device
+        *REG(VIRTIO_MMIO_STATUS) = 0;
+
+        // 4. Set ACKNOWLEDGE & DRIVER status
+        uint32 status = *REG(VIRTIO_MMIO_STATUS);
+        status |= VIRTIO_CONFIG_S_ACKNOWLEDGE | VIRTIO_CONFIG_S_DRIVER;
+        *REG(VIRTIO_MMIO_STATUS) = status;
+
+        // 5. Feature Negotiation
+        *REG(VIRTIO_MMIO_DEVICE_FEATURES_SEL) = 0;
+        uint32 device_features_lo = *REG(VIRTIO_MMIO_DEVICE_FEATURES);
+        (void)device_features_lo;
+
+        *REG(VIRTIO_MMIO_DEVICE_FEATURES_SEL) = 1;
+        uint32 device_features_hi = *REG(VIRTIO_MMIO_DEVICE_FEATURES);
+        (void)device_features_hi;
+
+        *REG(VIRTIO_MMIO_DRIVER_FEATURES_SEL) = 0;
+        *REG(VIRTIO_MMIO_DRIVER_FEATURES) = 0;
+
+        *REG(VIRTIO_MMIO_DRIVER_FEATURES_SEL) = 1;
+        // VIRTIO_F_VERSION_1 is usually bit 32.
+        // If not defined in header, we hardcode logic here: (1ULL << 0) for high bits
+        *REG(VIRTIO_MMIO_DRIVER_FEATURES) = 1;
+
+        status |= VIRTIO_CONFIG_S_FEATURES_OK;
+        *REG(VIRTIO_MMIO_STATUS) = status;
+
+        if (!(*REG(VIRTIO_MMIO_STATUS) & VIRTIO_CONFIG_S_FEATURES_OK))
+        {
+            Drivers::uart_puts("VirtIO: Feature negotiation failed.\n");
+            return;
+        }
+
+        // 6. Queue Configuration
+        *REG(VIRTIO_MMIO_QUEUE_SEL) = 0;
+        if (*REG(VIRTIO_MMIO_QUEUE_READY))
+        {
+            Drivers::uart_puts("VirtIO: Queue 0 not ready.\n");
+        }
+
+        uint32 max_queue = *REG(VIRTIO_MMIO_QUEUE_NUM_MAX);
+        if (max_queue < NUM_DESCS)
+        {
+            Drivers::uart_puts("VirtIO: Queue size too small.\n");
+            return;
+        }
+
+        *REG(VIRTIO_MMIO_QUEUE_NUM) = NUM_DESCS;
+
+        // Allocate Memory for VirtQueues
+        void *ptr = PMM::alloc_pages(1); // 4KB is enough for small rings
+        memset(ptr, 0, PGSIZE);
+
+        uint64 base_addr = (uint64)ptr;
+        disk.desc = reinterpret_cast<struct VRingDesc *>(base_addr);
+        disk.avail = reinterpret_cast<struct VRingAvail *>(base_addr + NUM_DESCS * 16);
+        disk.used = reinterpret_cast<struct VRingUsed *>(base_addr + PGSIZE / 2);
+
+        *REG(VIRTIO_MMIO_QUEUE_DESC_LOW) = (uint32)base_addr;
+        *REG(VIRTIO_MMIO_QUEUE_DESC_HIGH) = (uint32)(base_addr >> 32);
+
+        uint64 avail_addr = (uint64)disk.avail;
+        *REG(VIRTIO_MMIO_QUEUE_DRIVER_LOW) = (uint32)avail_addr;
+        *REG(VIRTIO_MMIO_QUEUE_DRIVER_HIGH) = (uint32)(avail_addr >> 32);
+
+        uint64 used_addr = (uint64)disk.used;
+        *REG(VIRTIO_MMIO_QUEUE_DEVICE_LOW) = (uint32)used_addr;
+        *REG(VIRTIO_MMIO_QUEUE_DEVICE_HIGH) = (uint32)(used_addr >> 32);
+
+        *REG(VIRTIO_MMIO_QUEUE_READY) = 1;
+
+        // 7. Initialize Internal Structures
+        disk.free_head = 0;
+        disk.num_free = NUM_DESCS;
+        disk.used_idx = 0;
+        for (int i = 0; i < NUM_DESCS; i++)
+        {
+            disk.free_chain[i] = i + 1;
+            disk.info[i].busy = 0;
+        }
+
+        // 8. Finalize Setup
+        status |= VIRTIO_CONFIG_S_DRIVER_OK;
+        *REG(VIRTIO_MMIO_STATUS) = status;
+
+        // === [Feature] Read Capacity ===
+        volatile struct virtio_blk_config *config = (volatile struct virtio_blk_config *)(virtio_base + VIRTIO_MMIO_CONFIG);
+        uint64 capacity_sectors = config->capacity;
+        uint64 size_mb = (capacity_sectors * 512) / (1024 * 1024);
+
+        Drivers::uart_puts(ANSI_GREEN "[VirtIO] Initialized. Capacity: ");
+        Drivers::uart_put_int(size_mb);
+        Drivers::uart_puts(" MB\n" ANSI_RESET);
+    }
+
+    static int alloc_desc()
+    {
+        if (disk.num_free == 0)
+            return -1;
+        int id = disk.free_head;
+        disk.free_head = disk.free_chain[id];
+        disk.num_free--;
+        return id;
+    }
+
+    static void free_desc(int id)
+    {
+        disk.free_chain[id] = disk.free_head;
+        disk.free_head = id;
+        disk.num_free++;
+        ProcManager::wakeup(&disk.free_head);
+    }
+
+    static void free_chain_soft(int head_idx)
+    {
+        int curr = head_idx;
+        while (1)
+        {
+            int flag = disk.desc[curr].flags;
+            int next = disk.desc[curr].next;
+            free_desc(curr);
+            if (flag & VRING_DESC_F_NEXT)
+                curr = next;
+            else
+                break;
+        }
+    }
+
+    void rw(int write, uint64 blockno, char *buf)
+    {
+        uint64 sector = blockno * (1024 / 512); // Assume 1KB blocks
+
+        disk.lock.acquire();
+
+        // Wait for enough descriptors
+        while (disk.num_free < 3)
+        {
+            ProcManager::sleep(&disk.free_head, &disk.lock);
+        }
+
+        int idx[3];
+        idx[0] = alloc_desc();
+        idx[1] = alloc_desc();
+        idx[2] = alloc_desc();
+
+        // 1. Request Header
+        struct BlockRequest *req = &disk.reqs[idx[0]];
+        req->type = write ? VIRTIO_BLK_T_OUT : VIRTIO_BLK_T_IN;
+        req->reserved = 0;
+        req->sector = sector;
+
+        disk.desc[idx[0]].addr = (uint64)req;
+        disk.desc[idx[0]].len = sizeof(struct BlockRequest);
+        disk.desc[idx[0]].flags = VRING_DESC_F_NEXT;
+        disk.desc[idx[0]].next = idx[1];
+
+        // 2. Data Buffer
+        disk.desc[idx[1]].addr = (uint64)buf;
+        disk.desc[idx[1]].len = 1024; // Fixed 1KB for this test
+        disk.desc[idx[1]].flags = VRING_DESC_F_NEXT;
+        if (!write)
+            disk.desc[idx[1]].flags |= VRING_DESC_F_WRITE;
+        disk.desc[idx[1]].next = idx[2];
+
+        // 3. Status Byte
+        disk.desc[idx[2]].addr = (uint64)&disk.req_status[idx[0]];
+        disk.desc[idx[2]].len = 1;
+        disk.desc[idx[2]].flags = VRING_DESC_F_WRITE;
+        disk.desc[idx[2]].next = 0;
+
+        disk.info[idx[0]].busy = 1;
+
+        // Submit to Avail Ring
+        disk.avail->ring[disk.avail->idx % NUM_DESCS] = idx[0];
+        __sync_synchronize();
+        disk.avail->idx++;
+        __sync_synchronize();
+
+        // Notify Device
+        *REG(VIRTIO_MMIO_QUEUE_NOTIFY) = 0;
+
+        // Wait for Completion
+        while (disk.info[idx[0]].busy)
+        {
+            ProcManager::sleep(&disk.info[idx[0]], &disk.lock);
+        }
+
+        free_chain_soft(idx[0]);
+        disk.lock.release();
+    }
+
+    void intr()
+    {
+        disk.lock.acquire();
+
+        uint32 status = *REG(VIRTIO_MMIO_INTERRUPT_STATUS);
+        *REG(VIRTIO_MMIO_INTERRUPT_ACK) = status & 0x3;
+
+        volatile struct VRingUsed *used_ring = (volatile struct VRingUsed *)disk.used;
+
+        while (disk.used_idx != used_ring->idx)
+        {
+            int id = used_ring->ring[disk.used_idx % NUM_DESCS].id;
+
+            if (disk.info[id].busy)
+            {
+                disk.info[id].busy = 0;
+                ProcManager::wakeup(&disk.info[id]);
+            }
+            disk.used_idx++;
+        }
+
+        disk.lock.release();
+    }
+
+    void test_rw()
+    {
+        Drivers::uart_puts(ANSI_BLUE "[VirtIO] Disk Test Start...\n" ANSI_RESET);
+
+        char *buf = (char *)PMM::alloc_page();
+        if (!buf)
+            return;
+
+        const char *msg = "Lume OS Disk Driver Test Pattern";
+        strcpy(buf, msg);
+
+        rw(1, 0, buf); // Write
+        memset(buf, 0, PGSIZE);
+        rw(0, 0, buf); // Read
+
+        if (strcmp(buf, msg) == 0)
+        {
+            Drivers::uart_puts(ANSI_GREEN "[VirtIO] Test PASSED (R/W Verified).\n" ANSI_RESET);
+        }
+        else
+        {
+            Drivers::uart_puts(ANSI_RED "[VirtIO] Test FAILED.\n" ANSI_RESET);
+        }
+
+        PMM::free_page(buf);
+    }
+}
\ No newline at end of file
diff --git a/src/kernel/main.cc b/src/kernel/main.cc
index 081ab3e..6e3bdde 100644
--- a/src/kernel/main.cc
+++ b/src/kernel/main.cc
@@ -1,6 +1,8 @@
 #include "common/types.h"
 #include "kernel/riscv.h"
 #include "drivers/uart.h"
+#include "drivers/plic.h"
+#include "drivers/virtio.h"
 #include "kernel/pmm.h"
 #include "kernel/mm.h"
 #include "kernel/proc.h"
@@ -8,7 +10,6 @@
 #include "kernel/timer.h"
 #include "kernel/cpu.h"
 
-// 引用外部 FDT 初始化函数
 extern void fdt_init(uint64 dtb);
 extern uint64 g_dtb_addr;
 
@@ -30,6 +31,9 @@ extern "C" void kernel_main(uint64 hartid, uint64 dtb)
         VM::kvminithart();        // Enable MMU
         Trap::init();             // Trap Management
         Trap::inithart();
+        PLIC::init();
+        PLIC::inithart();
+        VirtIO::init();
         Timer::init();
         ProcManager::init();      // Process Management
         ProcManager::user_init(); // Initialize first user process
@@ -45,6 +49,7 @@ extern "C" void kernel_main(uint64 hartid, uint64 dtb)
 
         VM::kvminithart();
         Trap::inithart();
+        PLIC::inithart();
         Timer::init();
         Drivers::uart_puts("[Boot] Hart started.\n");
     }
diff --git a/src/kernel/mm/pmm.cc b/src/kernel/mm/pmm.cc
index 09bd65b..3474a49 100644
--- a/src/kernel/mm/pmm.cc
+++ b/src/kernel/mm/pmm.cc
@@ -9,73 +9,6 @@
 extern "C" char kernel_end[];
 extern uint64 g_dtb_addr;
 
-// namespace PMM
-// {
-//     struct Run
-//     {
-//         struct Run *next;
-//     };
-
-//     struct
-//     {
-//         struct Run *freelist;
-//     } kmem;
-
-//     inline uint64_t PGROUNDUP(uint64_t a)
-//     {
-//         return (a + PGSIZE - 1) & ~(PGSIZE - 1);
-//     }
-
-//     void init()
-//     {
-//         uint64_t start = PGROUNDUP((uint64_t)kernel_end);
-
-//         Drivers::uart_puts("PMM: Initializing from");
-
-//         for (uint64_t p = start; p + PGSIZE <= PHYSTOP; p += PGSIZE)
-//         {
-//             free_page((void *)p);
-//         }
-//     }
-
-//     void free_page(void* pa)
-//     {
-//         struct Run *r;
-//         if (((uint64_t)pa % PGSIZE) != 0 || (uint64_t)pa < (uint64_t)kernel_end || (uint64_t)pa >= PHYSTOP)
-//         {
-//             Drivers::uart_puts("PMM: Error - Invalid free address\n");
-//             return;
-//         }
-//         r = (struct Run*)pa;
-
-//         r->next = kmem.freelist;
-//         kmem.freelist = r;
-        
-//     }
-
-//     void* alloc_page()
-//     {
-//         struct Run* r;
-
-//         r = kmem.freelist;
-//         if (r)
-//         {
-//             kmem.freelist = r->next;
-
-//             uint64_t *p = (uint64_t *)r;
-//             for (int i = 0; i < PGSIZE / 8; i++)
-//             {
-//                 p[i] = 0;
-//             }
-//         }
-//         else
-//         {
-//             Drivers::uart_puts("PMM: Out of memory!\n");
-//         }
-//         return (void *)r;
-//     }
-// }
-
 namespace PMM
 {
     // Global metadata array pointer
diff --git a/src/kernel/mm/slab.cc b/src/kernel/mm/slab.cc
index 370074a..702dc5c 100644
--- a/src/kernel/mm/slab.cc
+++ b/src/kernel/mm/slab.cc
@@ -19,6 +19,7 @@ struct SlabHeader
     uint16 used_count;       // Number of allocated objects
     // Followed immediately by object data...
     // The free list of objects is embedded in the first 2 bytes of the object itself (stores the index of the next object)
+    uint32 padding;
 };
 
 struct KmemCache
@@ -82,6 +83,7 @@ namespace Slab
         slab->next = nullptr;
         slab->used_count = 0;
         slab->free_idx = 0;
+        slab->padding = 0;
 
         // Initialize the free list of objects
         // The starting position of each object stores the index of the next object
@@ -186,7 +188,7 @@ namespace Slab
         if (!page)
             return;
 
-        if (!(page->flags & PG_slab))
+        if (!page || !(page->flags & PG_slab))
         {
             PMM::free_pages(ptr, page->order);
             return;
@@ -204,6 +206,13 @@ namespace Slab
 
         struct SlabHeader *slab = (struct SlabHeader *)PGROUNDDOWN((uint64)ptr);
 
+        if ((uint64)ptr < (uint64)(slab + 1) || (uint64)ptr >= (uint64)slab + PGSIZE)
+        {
+            Drivers::uart_puts("Slab: kfree invalid pointer\n");
+            c->lock.release();
+            return;
+        }
+
         char *base = (char *)(slab + 1);
         int obj_idx = ((uint64)ptr - (uint64)base) / c->obj_size;
 
diff --git a/src/kernel/mm/vm.cc b/src/kernel/mm/vm.cc
index c3c5f01..7057d88 100644
--- a/src/kernel/mm/vm.cc
+++ b/src/kernel/mm/vm.cc
@@ -360,11 +360,10 @@ namespace VM
     {
         kernel_pagetable = uvmcreate();
 
-        kvmmap(0x10000000, 0x10000000, PGSIZE, PTE_R | PTE_W);
+        kvmmap(0x10000000, 0x10000000, 0x10000, PTE_R | PTE_W);
         kvmmap(0x0c000000, 0x0c000000, 0x400000, PTE_R | PTE_W);
         kvmmap(0x02000000, 0x02000000, 0x10000, PTE_R | PTE_W); // CLINT
         kvmmap(KERNBASE, KERNBASE, PHYSTOP - KERNBASE, PTE_R | PTE_W | PTE_X);
-        // kvmmap(TRAMPOLINE, (uint64)trampoline, PGSIZE, PTE_R | PTE_X);
     }
 
     void kvminithart()
diff --git a/src/kernel/proc/proc.cc b/src/kernel/proc/proc.cc
index e99b83d..c8b6d1d 100644
--- a/src/kernel/proc/proc.cc
+++ b/src/kernel/proc/proc.cc
@@ -471,7 +471,6 @@ namespace ProcManager
 
     int wait(uint64 addr)
     {
-        (void)addr;
 
         struct Proc *np;
         int havekids, pid;
@@ -499,6 +498,11 @@ namespace ProcManager
                         }
 
                         pid = np->pid;
+                        if (addr != 0 && VM::copyout(p->pagetable, addr, (char *)&np->xstate, sizeof(np->xstate)) < 0)
+                        {
+                            proc_mem_lock.release();
+                            return -1;
+                        }
                         freeproc(np);
                         proc_mem_lock.release();
                         return pid;
@@ -509,8 +513,33 @@ namespace ProcManager
 
             if (!havekids || p->state == ZOMBIE)
                 return -1;
-            sleep(p, nullptr);
+            sleep(p, &proc_mem_lock);
+        }
+    }
+
+    int growproc(int n)
+    {
+        uint64 sz;
+        struct Proc *p = myproc();
+
+        sz = p->sz;
+
+        if (n > 0)
+        {
+            // Call VM::uvmalloc to allocate memory
+            // Parameters: page table, old size, new size, permissions (PTE_W)
+            if ((sz = VM::uvmalloc(p->pagetable, sz, sz + n, PTE_W)) == 0)
+            {
+                return -1;
+            }
         }
+        else if (n < 0)
+        {
+            sz = VM::uvmdealloc(p->pagetable, sz, sz + n);
+        }
+
+        p->sz = sz;
+        return 0;
     }
 
     void create_kernel_thread(void (*func)(), const char *name)
diff --git a/src/kernel/syscall.cc b/src/kernel/syscall.cc
index 1dfec73..837ac2c 100644
--- a/src/kernel/syscall.cc
+++ b/src/kernel/syscall.cc
@@ -3,6 +3,7 @@
 #include "kernel/proc.h"
 #include "kernel/riscv.h"
 #include "drivers/uart.h"
+#include "drivers/virtio.h"
 #include "kernel/timer.h"
 #include "kernel/mm.h"
 #include "kernel/trap.h"
@@ -124,6 +125,42 @@ static uint64_t sys_write()
     return static_cast<uint64_t>(-1);
 }
 
+static uint64_t sys_read()
+{
+    int fd;
+    int n;
+    uint64_t p;
+
+    // get param: fd, buf, len
+    if (argint(0, &fd) < 0 || argint(2, &n) < 0)
+        return -1;
+    p = argraw(1);
+
+    if (fd == 0) // stdin
+    {
+        return Drivers::console_read(p, n);
+    }
+    return -1;
+}
+
+static uint64_t sys_sbrk()
+{
+    int n;
+    if (argint(0, &n) < 0)
+        return -1;
+
+    uint64_t addr = myproc()->sz;
+    if (ProcManager::growproc(n) < 0)
+        return -1;
+
+    return addr;
+}
+
+static uint64_t sys_disk_test()
+{
+    VirtIO::test_rw();
+    return 0;
+}
 
 void syscall()
 {
@@ -131,15 +168,19 @@ void syscall()
     int num = p->tf->a7;
     uint64_t ret = static_cast<uint64_t>(-1);
 
-    Drivers::uart_puts("Syscall: ");
-    Drivers::print_hex(num);
-    Drivers::uart_puts("\n");
+    /* Just For Debug */
+    // Drivers::uart_puts("Syscall: ");
+    // Drivers::print_hex(num);
+    // Drivers::uart_puts("\n");
 
     switch (num)
     {
     case SYS_write:
         ret = sys_write();
         break;
+    case SYS_read:
+        ret = sys_read();
+        break;
     case SYS_putc:
         ret = sys_putc();
         break;
@@ -155,6 +196,12 @@ void syscall()
     case SYS_getpid:
         ret = sys_getpid();
         break;
+    case SYS_sbrk:
+        ret = sys_sbrk();
+        break;
+    case SYS_disk_test:
+        ret = sys_disk_test();
+        break;
     default:
         Drivers::uart_puts("Unknown Syscall ID: ");
         Drivers::print_hex(num);
diff --git a/src/kernel/trap/handler.cc b/src/kernel/trap/handler.cc
index d7f78de..c1626cc 100644
--- a/src/kernel/trap/handler.cc
+++ b/src/kernel/trap/handler.cc
@@ -6,6 +6,7 @@
 #include "kernel/mm.h"
 #include "kernel/timer.h"
 #include "drivers/plic.h"
+#include "drivers/virtio.h"
 #include "lib/string.h"
 
 
@@ -59,20 +60,24 @@ namespace Trap
         if (is_interrupt && code == 9)
         {
             // Supervisor External Interrupt (IRQ 9)
-            int irq = PILC::claim();
+            int irq = PLIC::claim();
 
             if (irq == 10)
             {
-                // uart_intr();
+                Drivers::uart_intr();
+                // Debug
+                // Drivers::uart_puts("[Interrupt] UART Input detected!\n");
             }
             else if (irq >= 1 && irq <= 8)
             {
-                // virtio_disk_intr();
+                VirtIO::intr();
+                // Debug
+                // Drivers::uart_puts("[Interrupt] VirtIO disk finished!\n");
             }
 
             if (irq > 0)
             {
-                PILC::complete(irq);
+                PLIC::complete(irq);
             }
 
             return 1;
diff --git a/src/user/entry.S b/src/user/entry.S
new file mode 100644
index 0000000..110ebbf
--- /dev/null
+++ b/src/user/entry.S
@@ -0,0 +1,12 @@
+.section .text
+.global _start
+
+_start:
+    li a0, 0
+    li a1, 0
+
+    call main
+
+    call exit
+
+1:  j 1b
\ No newline at end of file
diff --git a/src/user/init.cc b/src/user/init.cc
index 85252be..1646f6b 100644
--- a/src/user/init.cc
+++ b/src/user/init.cc
@@ -13,106 +13,175 @@ void print(const char *s)
     write(1, s, strlen(s));
 }
 
-void print_int(int val)
+void print_int(int n)
 {
     char buf[16];
     int i = 0;
-    if (val == 0)
+    int neg = 0;
+    if (n == 0)
     {
-        print("0");
+        write(1, "0", 1);
         return;
     }
-    if (val < 0)
+    if (n < 0)
     {
-        print("-");
-        val = -val;
+        neg = 1;
+        n = -n;
     }
-    while (val > 0)
+    while (n > 0)
     {
-        buf[i++] = (val % 10) + '0';
-        val /= 10;
+        buf[i++] = (n % 10) + '0';
+        n /= 10;
     }
+    if (neg)
+        buf[i++] = '-';
     while (i > 0)
-    {
-        char c = buf[--i];
-        write(1, &c, 1);
-    }
+        write(1, &buf[--i], 1);
 }
 
-void pass() { print(" -> [PASS]\n"); }
-void fail()
+void print_hex(uint64 n)
 {
-    print(" -> [FAIL]\n");
-    exit(1);
+    char buf[16];
+    char digits[] = "0123456789ABCDEF";
+    write(1, "0x", 2);
+    if (n == 0)
+    {
+        write(1, "0", 1);
+        return;
+    }
+    int i = 0;
+    while (n > 0)
+    {
+        buf[i++] = digits[n % 16];
+        n /= 16;
+    }
+    while (i > 0)
+        write(1, &buf[--i], 1);
 }
 
-int main()
+void test_syscalls()
 {
-    print("\n=== Lume OS Full Validation ===\n");
-
-    print("[TEST 1] Write\n");
-    print("  System call write is working.");
-    pass();
-
-    print("[TEST 2] GetPID\n");
-    int parent_pid = getpid();
-    print("  Parent PID: ");
-    print_int(parent_pid);
-    if (parent_pid > 0)
-        pass();
+    print("\n=== [Test] Starting System Call Check ===\n");
+
+    print("[1/4] Checking getpid()... ");
+    int my_pid = getpid();
+    if (my_pid > 0)
+    {
+        print("PASS (PID=");
+        print_int(my_pid);
+        print(")\n");
+    }
     else
-        fail();
+    {
+        print("FAIL (Invalid PID)\n");
+    }
+
+    print("[2/4] Checking sbrk() heap allocation... ");
+    char *p = sbrk(4096);
+    if (p == (char *)-1)
+    {
+        print("FAIL (sbrk returned error)\n");
+    }
+    else
+    {
+        p[0] = 'A';
+        p[4095] = 'Z';
+        if (p[0] == 'A' && p[4095] == 'Z')
+        {
+            print("PASS (Allocated at ");
+            print_hex((uint64)p);
+            print(")\n");
+        }
+        else
+        {
+            print("FAIL (Memory read/write check failed)\n");
+        }
+    }
 
-    print("[TEST 3] Fork, Wait & COW\n");
-    int val = 100;
+    print("[3/4] Checking fork() and wait()...\n");
     int pid = fork();
 
     if (pid < 0)
     {
-        print("  Fork Failed\n");
-        fail();
+        print("      FAIL: fork error\n");
     }
-
-    if (pid == 0)
+    else if (pid == 0)
     {
-        print("  [Child] PID: ");
-        print_int(getpid());
-        print("\n");
-
-        if (val != 100)
+        int child_pid = getpid();
+        if (child_pid == my_pid)
         {
-            print("  [Child] Data corruption!\n");
+            print("      FAIL: Child has same PID as parent\n");
             exit(1);
         }
-
-        val = 200;
-        print("  [Child] Modified val to 200. Exiting.\n");
+        print("      [Child] I am alive. PID=");
+        print_int(child_pid);
+        print("\n");
+        print("      [Child] Exiting now.\n");
         exit(0);
     }
     else
     {
-        print("  [Parent] Waiting for child...\n");
-        int child_pid = wait(0);
-
-        print("  [Parent] Child ");
-        print_int(child_pid);
-        print(" exited.\n");
+        print("      [Parent] Forked child PID=");
+        print_int(pid);
+        print("\n");
+        print("      [Parent] Waiting for child...\n");
 
-        print("  [Parent] Check val (Expect 100): ");
-        print_int(val);
+        int waited_pid = wait(0);
 
-        if (val == 100)
-            pass();
+        if (waited_pid == pid)
+        {
+            print("      [Parent] Wait returned correct PID.\n");
+            print("PASS (Fork/Wait cycle)\n");
+        }
         else
         {
-            print("  [Parent] Memory Isolation Failed! Val=");
-            print_int(val);
-            fail();
+            print("FAIL (Wait returned ");
+            print_int(waited_pid);
+            print(")\n");
         }
     }
 
-    print("\n=== All Tests Passed ===\n");
+    print("[4/4] Checking write()... PASS (Self-evident)\n");
+
+    print("=== [Test] All Tests Finished ===\n");
+}
+
+
+int main()
+{
+    print("\n[Lume Shell] Ready.\n");
+    print("Commands:\n");
+    print("  s: Run System Call Tests (fork, wait, mm)\n");
+    print("  t: Run Disk Driver Test (VirtIO)\n");
+
     while (1)
-        ;
+    {
+        char buf[16];
+        write(1, "$ ", 2);
+
+        int n = read(0, buf, sizeof(buf));
+        if (n <= 0)
+            continue;
+
+        if (buf[0] == 's')
+        {
+            test_syscalls();
+        }
+        else if (buf[0] == 't')
+        {
+            if (fork() == 0)
+            {
+                disk_test();
+                exit(0);
+            }
+            wait(0);
+        }
+        else if (buf[0] != '\n')
+        {
+            print("Unknown command.\n");
+        }
+    }
+
+    exit(0);
     return 0;
 }
\ No newline at end of file
diff --git a/src/user/usys.S b/src/user/usys.S
index e805d82..65937cc 100644
--- a/src/user/usys.S
+++ b/src/user/usys.S
@@ -28,4 +28,5 @@ SYSCALL(getpid)
 SYSCALL(sbrk)
 SYSCALL(sleep)
 SYSCALL(uptime)
-SYSCALL(putc)
\ No newline at end of file
+SYSCALL(putc)
+SYSCALL(disk_test)
\ No newline at end of file