usb hubs, hotplug support, unplug support, saving log files to disk, smt

code/kernel/io_device.c

@@ -1,13 +1,423 @@
-// io_device.c - kernel driver/device dispatch and disk registry
+// io_device.c - I/O manager: driver/device objects, stacks, PnP skeleton, disk namespace
 #include <string.h>
+#include <intrin.h>
 #include <ntlite/ddi/io_driver.h>
 #include <kernel/ke_irp.h>
 #include <kernel/mm_virt.h>
 #include <kernel/ke_debug.h>
-
+#include <kernel/ob_object.h>
 
 #define MAX_PUBLISHED_DEVICES 64
-#define MAX_DRIVER_MATCHERS 32
+#define MAX_DRIVER_MATCHERS   32
+#define MAX_DRIVER_OBJECTS    64
+#define MAX_PDO_OBJECTS       128
+
+static SYS_STRING_VIEW IoSvOrUnnamed(SYS_STRING_VIEW name) {
+    return SysSvIsEmpty(name) ? SYS_SV_LITERAL("<unnamed>") : name;
+}
+
+static const char* IoBusTypeName(IO_BUS_TYPE busType) {
+    switch (busType) {
+        case IoBusPci: return "pci";
+        case IoBusUsb: return "usb";
+        default: return "unknown";
+    }
+}
+
+// -------------------------------------------------------------------------
+// Driver object registry / PDO AddDevice binding
+// -------------------------------------------------------------------------
+
+static DRIVER_OBJECT* gDriverObjects[MAX_DRIVER_OBJECTS];
+static uint32_t gDriverObjectCount;
+static DEVICE_OBJECT* gPhysicalDevices[MAX_PDO_OBJECTS];
+static uint32_t gPhysicalDeviceCount;
+
+static bool IoDriverMatchesPdo(const DRIVER_OBJECT* driver, const DEVICE_OBJECT* pdo) {
+    if (!driver || !driver->IdTable || !pdo || pdo->Type != DeviceObjectPhysical)
+        return false;
+
+    for (size_t d = 0; d < driver->IdTable->Count; d++) {
+        const IO_DEVICE_ID* want = &driver->IdTable->Ids[d];
+        if (want->BusType != IoBusUnknown && want->BusType != pdo->u.Pdo.BusType)
+            continue;
+        for (uint8_t p = 0; p < pdo->u.Pdo.IdCount; p++) {
+            if (SysSvEqual(want->Id, pdo->u.Pdo.Ids[p].Id))
+                return true;
+        }
+    }
+    return false;
+}
+
+static bool IoTryAddDeviceForDriver(DRIVER_OBJECT* driver, DEVICE_OBJECT* pdo) {
+    if (!driver || !pdo || !driver->DriverAddDevice)
+        return false;
+    if (!IoDriverMatchesPdo(driver, pdo))
+        return false;
+
+    KSTATUS st = driver->DriverAddDevice(driver, pdo);
+    if (KSUCCESS(st)) {
+        DbgPrint("io: AddDevice driver=%v pdo=%v bus=%s\n",
+                 IoSvOrUnnamed(driver->Name), IoSvOrUnnamed(pdo->Name), IoBusTypeName(pdo->u.Pdo.BusType));
+
+        KSTATUS start = IoSendPnpIrpSync(pdo, IRP_MN_START_DEVICE);
+        if (!KSUCCESS(start) && start != STATUS_NOT_SUPPORTED) {
+            DbgPrint("io: START_DEVICE failed driver=%v pdo=%v st=0x%x\n",
+                     IoSvOrUnnamed(driver->Name), IoSvOrUnnamed(pdo->Name), start);
+        }
+        return true;
+    } else {
+        DbgPrint("io: AddDevice failed driver=%v pdo=%v st=0x%x\n",
+                 IoSvOrUnnamed(driver->Name), IoSvOrUnnamed(pdo->Name), st);
+        return true;
+    }
+}
+
+__declspec(dllexport) void IoInitializeDriverObject(DRIVER_OBJECT* driver,
+                                                    SYS_STRING_VIEW name) {
+    if (!driver) return;
+    memset(driver, 0, sizeof(*driver));
+    driver->Header.Type = ObjectTypeDriver;
+    driver->Header.RefCount = 1;
+    driver->Header.Destroy = (void*)0;
+    driver->Name = name;
+}
+
+__declspec(dllexport) KSTATUS IoRegisterDriverObject(DRIVER_OBJECT* driver) {
+    if (!driver) return STATUS_INVALID_PARAMETER;
+    if (gDriverObjectCount >= MAX_DRIVER_OBJECTS)
+        return STATUS_NO_MEMORY;
+
+    if (driver->Header.Type != ObjectTypeDriver) {
+        driver->Header.Type = ObjectTypeDriver;
+        driver->Header.RefCount = 1;
+    }
+
+    gDriverObjects[gDriverObjectCount++] = driver;
+    DbgPrint("io: registered driver object %v\n", IoSvOrUnnamed(driver->Name));
+
+    for (uint32_t i = 0; i < gPhysicalDeviceCount; i++)
+        (void)IoTryAddDeviceForDriver(driver, gPhysicalDevices[i]);
+
+    return STATUS_SUCCESS;
+}
+
+// -------------------------------------------------------------------------
+// Device object lifetime
+// -------------------------------------------------------------------------
+
+static void IoDeviceDestroy(OBJECT* object) {
+    DEVICE_OBJECT* dev = (DEVICE_OBJECT*)object;
+
+    for (uint32_t i = 0; i < gPhysicalDeviceCount; i++) {
+        if (gPhysicalDevices[i] == dev) {
+            gPhysicalDevices[i] = gPhysicalDevices[gPhysicalDeviceCount - 1u];
+            gPhysicalDevices[gPhysicalDeviceCount - 1u] = (DEVICE_OBJECT*)0;
+            gPhysicalDeviceCount--;
+            break;
+        }
+    }
+
+    if (dev->LowerDevice)
+        IoDetachDevice(dev);
+
+    if (dev->Parent) {
+        DEVICE_OBJECT* parent = dev->Parent;
+        uintptr_t flags = KeAcquireSpinlock(&parent->Lock);
+        DEVICE_OBJECT** link = &parent->FirstChild;
+        while (*link) {
+            if (*link == dev) {
+                *link = dev->NextSibling;
+                break;
+            }
+            link = &(*link)->NextSibling;
+        }
+        KeReleaseSpinlock(&parent->Lock, flags);
+        dev->Parent = (DEVICE_OBJECT*)0;
+        dev->NextSibling = (DEVICE_OBJECT*)0;
+        IoDereferenceDeviceObject(parent);
+    }
+}
+
+static KSTATUS IoCreateDeviceInternal(DRIVER_OBJECT* driver,
+                                      DEVICE_OBJECT_TYPE type,
+                                      size_t extensionSize,
+                                      DEVICE_OBJECT** outDevice) {
+    if (!driver || !outDevice) return STATUS_INVALID_PARAMETER;
+
+    size_t total = sizeof(DEVICE_OBJECT) + extensionSize;
+    DEVICE_OBJECT* dev = (DEVICE_OBJECT*)ObAllocateObject(ObjectTypeDevice, total);
+    if (!dev) return STATUS_NO_MEMORY;
+
+    dev->Header.Destroy = IoDeviceDestroy;
+    dev->DriverObject = driver;
+    dev->Type = type;
+    dev->PnpState = DevicePnpNotStarted;
+    dev->AlignmentRequirement = 1;
+    dev->Lock = (SPINLOCK)SPINLOCK_INIT;
+
+    // The per-device union is interpreted differently for PDOs and disk FDOs.
+    // Do not write disk defaults into a PDO: u.Disk.DiskIndex aliases the low
+    // byte of u.Pdo.BusType, which prevents ID-table AddDevice matching.
+    if (type != DeviceObjectPhysical) {
+        dev->u.Disk.DiskIndex = 0xFF;
+        dev->u.Disk.SectorSize = 512;
+    }
+
+    if (extensionSize)
+        dev->DeviceExtension = (void*)((uint8_t*)dev + sizeof(DEVICE_OBJECT));
+
+    *outDevice = dev;
+    return STATUS_SUCCESS;
+}
+
+__declspec(dllexport) KSTATUS IoCreateDevice(DRIVER_OBJECT* driver,
+                                             size_t extensionSize,
+                                             DEVICE_OBJECT** outDevice) {
+    return IoCreateDeviceInternal(driver, DeviceObjectFunction, extensionSize, outDevice);
+}
+
+__declspec(dllexport) KSTATUS IoCreatePhysicalDevice(DRIVER_OBJECT* busDriver,
+                                                     DEVICE_OBJECT* parent,
+                                                     const IO_DEVICE_ID* ids,
+                                                     size_t idCount,
+                                                     size_t extensionSize,
+                                                     DEVICE_OBJECT** outPdo) {
+    if (!busDriver || !outPdo) return STATUS_INVALID_PARAMETER;
+    if (idCount > IO_MAX_DEVICE_IDS) return STATUS_INVALID_PARAMETER;
+
+    DEVICE_OBJECT* pdo = (DEVICE_OBJECT*)0;
+    KSTATUS st = IoCreateDeviceInternal(busDriver, DeviceObjectPhysical, extensionSize, &pdo);
+    if (!KSUCCESS(st)) return st;
+
+    if (idCount && ids) {
+        for (size_t i = 0; i < idCount; i++) {
+            if (!SysSvIsValid(ids[i].Id) || ids[i].Id.Length >= IO_MAX_DEVICE_ID_CHARS) {
+                IoDeleteDevice(pdo);
+                return STATUS_INVALID_PARAMETER;
+            }
+
+            char* dst = pdo->u.Pdo.IdStorage[i];
+            if (ids[i].Id.Length)
+                memcpy(dst, ids[i].Id.Data, ids[i].Id.Length);
+            dst[ids[i].Id.Length] = '\0';
+
+            pdo->u.Pdo.Ids[i].BusType = ids[i].BusType;
+            pdo->u.Pdo.Ids[i].Id = SysSvFromParts(dst, ids[i].Id.Length);
+            if (pdo->u.Pdo.BusType == IoBusUnknown)
+                pdo->u.Pdo.BusType = ids[i].BusType;
+        }
+        pdo->u.Pdo.IdCount = (uint8_t)idCount;
+    }
+
+    if (parent) {
+        IoReferenceDeviceObject(parent);
+        pdo->Parent = parent;
+        uintptr_t flags = KeAcquireSpinlock(&parent->Lock);
+        pdo->NextSibling = parent->FirstChild;
+        parent->FirstChild = pdo;
+        KeReleaseSpinlock(&parent->Lock, flags);
+    }
+
+    *outPdo = pdo;
+    return STATUS_SUCCESS;
+}
+
+__declspec(dllexport) KSTATUS IoRegisterPhysicalDevice(DEVICE_OBJECT* pdo) {
+    if (!pdo || pdo->Type != DeviceObjectPhysical)
+        return STATUS_INVALID_PARAMETER;
+
+    for (uint32_t i = 0; i < gPhysicalDeviceCount; i++) {
+        if (gPhysicalDevices[i] == pdo)
+            return STATUS_ALREADY_EXISTS;
+    }
+
+    if (gPhysicalDeviceCount >= MAX_PDO_OBJECTS)
+        return STATUS_NO_MEMORY;
+
+    gPhysicalDevices[gPhysicalDeviceCount++] = pdo;
+
+    bool anyMatched = false;
+    for (uint32_t i = 0; i < gDriverObjectCount; i++) {
+        if (IoTryAddDeviceForDriver(gDriverObjects[i], pdo))
+            anyMatched = true;
+    }
+
+    if (!anyMatched && pdo->u.Pdo.BusType == IoBusUsb) {
+        SYS_STRING_VIEW id = pdo->u.Pdo.IdCount ? pdo->u.Pdo.Ids[0].Id : SYS_SV_LITERAL("<no-id>");
+        DbgPrint("io: no function driver for pdo=%v id=%v bus=%s\n",
+                 IoSvOrUnnamed(pdo->Name), id, IoBusTypeName(pdo->u.Pdo.BusType));
+    }
+    return STATUS_SUCCESS;
+}
+
+
+__declspec(dllexport) KSTATUS IoSetDeviceName(DEVICE_OBJECT* device,
+                                               SYS_STRING_VIEW name) {
+    if (!device || !SysSvIsValid(name))
+        return STATUS_INVALID_PARAMETER;
+    if (name.Length >= IO_MAX_DEVICE_NAME_CHARS)
+        return STATUS_BUFFER_TOO_SMALL;
+
+    if (name.Length)
+        memcpy(device->NameStorage, name.Data, name.Length);
+    device->NameStorage[name.Length] = '\0';
+    device->Name = SysSvFromParts(device->NameStorage, name.Length);
+    return STATUS_SUCCESS;
+}
+
+__declspec(dllexport) void IoReferenceDeviceObject(DEVICE_OBJECT* device) {
+    if (device) ObReferenceObject(&device->Header);
+}
+
+__declspec(dllexport) void IoDereferenceDeviceObject(DEVICE_OBJECT* device) {
+    if (device) ObDereferenceObject(&device->Header);
+}
+
+static void IoRemovePhysicalDeviceRegistryEntry(DEVICE_OBJECT* device) {
+    if (!device || device->Type != DeviceObjectPhysical)
+        return;
+
+    for (uint32_t i = 0; i < gPhysicalDeviceCount; i++) {
+        if (gPhysicalDevices[i] == device) {
+            for (uint32_t j = i + 1; j < gPhysicalDeviceCount; j++)
+                gPhysicalDevices[j - 1] = gPhysicalDevices[j];
+            gPhysicalDeviceCount--;
+            if (gPhysicalDeviceCount < MAX_PDO_OBJECTS)
+                gPhysicalDevices[gPhysicalDeviceCount] = (DEVICE_OBJECT*)0;
+            return;
+        }
+    }
+}
+
+static void IoUnlinkDeviceFromParent(DEVICE_OBJECT* device) {
+    if (!device || !device->Parent)
+        return;
+
+    DEVICE_OBJECT* parent = device->Parent;
+    uintptr_t flags = KeAcquireSpinlock(&parent->Lock);
+    DEVICE_OBJECT** link = &parent->FirstChild;
+    while (*link) {
+        if (*link == device) {
+            *link = device->NextSibling;
+            break;
+        }
+        link = &(*link)->NextSibling;
+    }
+    device->Parent = (DEVICE_OBJECT*)0;
+    device->NextSibling = (DEVICE_OBJECT*)0;
+    KeReleaseSpinlock(&parent->Lock, flags);
+
+    IoDereferenceDeviceObject(parent);
+}
+
+__declspec(dllexport) void IoDeleteDevice(DEVICE_OBJECT* device) {
+    if (!device) return;
+    if (device->DeletePending)
+        return;
+
+    device->DeletePending = true;
+    device->PnpState = DevicePnpDeleted;
+    IoUnlinkDeviceFromParent(device);
+    IoRemovePhysicalDeviceRegistryEntry(device);
+    IoDereferenceDeviceObject(device);
+}
+
+// -------------------------------------------------------------------------
+// Device stack management
+// -------------------------------------------------------------------------
+
+__declspec(dllexport) DEVICE_OBJECT* IoAttachDeviceToDeviceStack(DEVICE_OBJECT* source,
+                                                                 DEVICE_OBJECT* target) {
+    if (!source || !target) return (DEVICE_OBJECT*)0;
+
+    DEVICE_OBJECT* top = target;
+    while (top->AttachedDevice)
+        top = top->AttachedDevice;
+
+    top->AttachedDevice = source;
+    source->LowerDevice = top;
+    IoReferenceDeviceObject(top);
+    return top;
+}
+
+__declspec(dllexport) void IoDetachDevice(DEVICE_OBJECT* source) {
+    if (!source || !source->LowerDevice) return;
+
+    DEVICE_OBJECT* lower = source->LowerDevice;
+    if (lower->AttachedDevice == source)
+        lower->AttachedDevice = (DEVICE_OBJECT*)0;
+    source->LowerDevice = (DEVICE_OBJECT*)0;
+    IoDereferenceDeviceObject(lower);
+}
+
+__declspec(dllexport) DEVICE_OBJECT* IoGetAttachedDeviceReference(DEVICE_OBJECT* device) {
+    if (!device) return (DEVICE_OBJECT*)0;
+    DEVICE_OBJECT* top = device;
+    while (top->AttachedDevice)
+        top = top->AttachedDevice;
+    IoReferenceDeviceObject(top);
+    return top;
+}
+
+// -------------------------------------------------------------------------
+// Device-stack helpers / PnP dispatch
+// -------------------------------------------------------------------------
+
+__declspec(dllexport) uint8_t IoGetDeviceStackSize(DEVICE_OBJECT* device) {
+    uint8_t count = 0;
+    for (DEVICE_OBJECT* cur = device; cur && count < 0xFEu; cur = cur->LowerDevice)
+        count++;
+    return count ? count : 1;
+}
+
+__declspec(dllexport) KSTATUS IoForwardIrpToLower(DEVICE_OBJECT* device, IRP* irp) {
+    if (!device || !irp || !device->LowerDevice)
+        return STATUS_NOT_SUPPORTED;
+
+    IO_STACK_LOCATION* cur = IoGetCurrentIrpStackLocation(irp);
+    IO_STACK_LOCATION* next = IoGetNextIrpStackLocation(irp);
+    if (!cur || !next)
+        return STATUS_INVALID_PARAMETER;
+
+    DEVICE_OBJECT* currentDevice = cur->DeviceObject;
+    memcpy(next, cur, sizeof(*next));
+
+    if (currentDevice == device) {
+        cur->DeviceObject = (DEVICE_OBJECT*)0;
+        _InterlockedDecrement(&device->OutstandingIrps);
+        IoDereferenceDeviceObject(device);
+    }
+
+    irp->CurrentLocation++;
+    return IoCallDriver(device->LowerDevice, irp);
+}
+
+__declspec(dllexport) KSTATUS IoSendPnpIrpSync(DEVICE_OBJECT* device,
+                                                IO_PNP_MINOR_FUNCTION minor) {
+    if (!device) return STATUS_INVALID_PARAMETER;
+
+    DEVICE_OBJECT* top = IoGetAttachedDeviceReference(device);
+    if (!top) return STATUS_INVALID_PARAMETER;
+
+    uint8_t stackCount = IoGetDeviceStackSize(top);
+    IRP* irp = (IRP*)0;
+    KSTATUS st = IoAllocateIrp(stackCount, &irp);
+    if (!KSUCCESS(st)) {
+        IoDereferenceDeviceObject(top);
+        return st;
+    }
+
+    IoBuildPnpIrp(irp, minor);
+    st = IoSubmitIrpSync(top, irp);
+    IoFreeIrp(irp);
+    IoDereferenceDeviceObject(top);
+    return st;
+}
+
+// -------------------------------------------------------------------------
+// Legacy publish/matcher path
+// -------------------------------------------------------------------------
 
 typedef struct IO_DRIVER_MATCHER_ENTRY {
     bool Used;
@@ -23,14 +433,6 @@ static uint32_t gPublishedDeviceCount;
 static IO_DRIVER_MATCHER_ENTRY gDriverMatchers[MAX_DRIVER_MATCHERS];
 static uint32_t gDriverMatcherCount;
 
-static const char* IoBusTypeName(IO_BUS_TYPE busType) {
-    switch (busType) {
-        case IoBusPci: return "pci";
-        case IoBusUsb: return "usb";
-        default: return "unknown";
-    }
-}
-
 static bool IoTryBindDevice(IO_PUBLISHED_DEVICE* device,
                             IO_DRIVER_MATCHER_ENTRY* matcher) {
     if (!device || !matcher || !matcher->Used || device->Bound)
@@ -80,12 +482,10 @@ __declspec(dllexport) KSTATUS IoRegisterDriverMatcher(const char* driverName,
     entry->Bind = bind;
     entry->Context = context;
 
-    DbgPrint("io: registered %s matcher %s\n",
-             IoBusTypeName(busType), driverName);
+    DbgPrint("io: registered %s matcher %s\n", IoBusTypeName(busType), driverName);
 
-    for (uint32_t i = 0; i < gPublishedDeviceCount; i++) {
+    for (uint32_t i = 0; i < gPublishedDeviceCount; i++)
         IoTryBindDevice(&gPublishedDevices[i], entry);
-    }
 
     return STATUS_SUCCESS;
 }
@@ -101,7 +501,6 @@ __declspec(dllexport) KSTATUS IoPublishDevice(const IO_PUBLISHED_DEVICE* device)
     stored->Bound = false;
     stored->BoundDriverName = (const char*)0;
 
-
     for (uint32_t i = 0; i < gDriverMatcherCount; i++) {
         if (IoTryBindDevice(stored, &gDriverMatchers[i]))
             break;
@@ -114,50 +513,44 @@ __declspec(dllexport) uint32_t IoPublishedDeviceCount(void) {
     return gPublishedDeviceCount;
 }
 
-#define MAX_DISK_DEVICES 8
-#define MAX_DISK_ARRIVAL_CALLBACKS 16
+// -------------------------------------------------------------------------
+// Disk namespace
+// -------------------------------------------------------------------------
 
-typedef struct IO_DISK_CALLBACK_ENTRY {
+#define MAX_DISK_DEVICES 8
+#define MAX_DISK_CALLBACKS 16
+
+typedef struct IO_DISK_CHANGE_CALLBACK_ENTRY {
+    bool Used;
+    IO_DISK_CHANGE_CALLBACK Callback;
+    void* Context;
+} IO_DISK_CHANGE_CALLBACK_ENTRY;
+
+typedef struct IO_DISK_ARRIVAL_CALLBACK_ENTRY {
     bool Used;
     IO_DISK_ARRIVAL_CALLBACK Callback;
     void* Context;
-} IO_DISK_CALLBACK_ENTRY;
+} IO_DISK_ARRIVAL_CALLBACK_ENTRY;
 
 static DEVICE_OBJECT* gDiskDevices[MAX_DISK_DEVICES];
-static IO_DISK_CALLBACK_ENTRY gDiskArrivalCallbacks[MAX_DISK_ARRIVAL_CALLBACKS];
+static IO_DISK_CHANGE_CALLBACK_ENTRY gDiskChangeCallbacks[MAX_DISK_CALLBACKS];
+static uint32_t gDiskChangeCallbackCount;
+static IO_DISK_ARRIVAL_CALLBACK_ENTRY gDiskArrivalCallbacks[MAX_DISK_CALLBACKS];
 static uint32_t gDiskArrivalCallbackCount;
 
-static void IoNotifyDiskArrived(uint8_t driveIndex, DEVICE_OBJECT* device) {
-    for (uint32_t i = 0; i < gDiskArrivalCallbackCount; i++) {
-        IO_DISK_CALLBACK_ENTRY* entry = &gDiskArrivalCallbacks[i];
+static void IoNotifyDiskChange(uint8_t driveIndex, DEVICE_OBJECT* device, IO_DEVICE_CHANGE change) {
+    for (uint32_t i = 0; i < gDiskChangeCallbackCount; i++) {
+        IO_DISK_CHANGE_CALLBACK_ENTRY* entry = &gDiskChangeCallbacks[i];
         if (entry->Used && entry->Callback)
-            entry->Callback(driveIndex, device, entry->Context);
+            entry->Callback(driveIndex, device, change, entry->Context);
+    }
+    if (change == IoDeviceArrived) {
+        for (uint32_t i = 0; i < gDiskArrivalCallbackCount; i++) {
+            IO_DISK_ARRIVAL_CALLBACK_ENTRY* entry = &gDiskArrivalCallbacks[i];
+            if (entry->Used && entry->Callback)
+                entry->Callback(driveIndex, device, entry->Context);
+        }
     }
-}
-
-__declspec(dllexport) KSTATUS IoCreateDevice(DRIVER_OBJECT* driver,
-                                             size_t extensionSize,
-                                             DEVICE_OBJECT** outDevice) {
-    if (!driver || !outDevice) return STATUS_INVALID_PARAMETER;
-
-    size_t total = sizeof(DEVICE_OBJECT) + extensionSize;
-    DEVICE_OBJECT* dev = (DEVICE_OBJECT*)MmAllocatePool(total);
-    if (!dev) return STATUS_NO_MEMORY;
-
-    memset(dev, 0, total);
-    dev->DriverObject = driver;
-    if (extensionSize)
-        dev->DeviceExtension = (void*)((uint8_t*)dev + sizeof(DEVICE_OBJECT));
-    dev->AlignmentRequirement = 1;
-    dev->DiskIndex = 0xFF;
-
-    *outDevice = dev;
-    return STATUS_SUCCESS;
-}
-
-__declspec(dllexport) void IoDeleteDevice(DEVICE_OBJECT* device) {
-    if (!device) return;
-    MmFreePool(device);
 }
 
 __declspec(dllexport) KSTATUS IoRegisterDiskDevice(uint8_t driveIndex,
@@ -168,9 +561,57 @@ __declspec(dllexport) KSTATUS IoRegisterDiskDevice(uint8_t driveIndex,
         return STATUS_ALREADY_EXISTS;
 
     device->Flags |= DEVICE_OBJECT_DISK;
-    device->DiskIndex = driveIndex;
+    device->u.Disk.DiskIndex = driveIndex;
+    if (device->u.Disk.SectorSize == 0)
+        device->u.Disk.SectorSize = 512;
+
+    KSTATUS pnpStatus = IoSendPnpIrpSync(device, IRP_MN_START_DEVICE);
+    if (!KSUCCESS(pnpStatus) && pnpStatus != STATUS_NOT_SUPPORTED)
+        return pnpStatus;
+    if (pnpStatus == STATUS_NOT_SUPPORTED)
+        device->PnpState = DevicePnpStarted;
+
+    IoReferenceDeviceObject(device);
     gDiskDevices[driveIndex] = device;
-    IoNotifyDiskArrived(driveIndex, device);
+    IoNotifyDiskChange(driveIndex, device, IoDeviceArrived);
+    return STATUS_SUCCESS;
+}
+
+__declspec(dllexport) KSTATUS IoNotifyDiskSurpriseRemoved(uint8_t driveIndex) {
+    if (driveIndex >= MAX_DISK_DEVICES)
+        return STATUS_INVALID_PARAMETER;
+    DEVICE_OBJECT* dev = gDiskDevices[driveIndex];
+    if (!dev) return STATUS_NOT_FOUND;
+
+    dev->PnpState = DevicePnpSurpriseRemoved;
+    IoNotifyDiskChange(driveIndex, dev, IoDeviceSurpriseRemoved);
+    (void)IoSendPnpIrpSync(dev, IRP_MN_SURPRISE_REMOVAL);
+    return STATUS_SUCCESS;
+}
+
+__declspec(dllexport) KSTATUS IoUnregisterDiskDevice(uint8_t driveIndex) {
+    if (driveIndex >= MAX_DISK_DEVICES)
+        return STATUS_INVALID_PARAMETER;
+    DEVICE_OBJECT* dev = gDiskDevices[driveIndex];
+    if (!dev) return STATUS_NOT_FOUND;
+
+    if (dev->PnpState != DevicePnpSurpriseRemoved) {
+        IoNotifyDiskChange(driveIndex, dev, IoDeviceQueryRemove);
+        KSTATUS queryStatus = IoSendPnpIrpSync(dev, IRP_MN_QUERY_REMOVE_DEVICE);
+        if (!KSUCCESS(queryStatus) && queryStatus != STATUS_NOT_SUPPORTED)
+            return queryStatus;
+
+        dev->PnpState = DevicePnpSurpriseRemoved;
+        IoNotifyDiskChange(driveIndex, dev, IoDeviceSurpriseRemoved);
+        (void)IoSendPnpIrpSync(dev, IRP_MN_SURPRISE_REMOVAL);
+    }
+
+    dev->PnpState = DevicePnpRemovePending;
+    (void)IoSendPnpIrpSync(dev, IRP_MN_REMOVE_DEVICE);
+
+    gDiskDevices[driveIndex] = (DEVICE_OBJECT*)0;
+    IoNotifyDiskChange(driveIndex, dev, IoDeviceRemoved);
+    IoDereferenceDeviceObject(dev);
     return STATUS_SUCCESS;
 }
 
@@ -178,10 +619,10 @@ __declspec(dllexport) KSTATUS IoRegisterDiskArrivalCallback(IO_DISK_ARRIVAL_CALL
                                                             void* context) {
     if (!callback)
         return STATUS_INVALID_PARAMETER;
-    if (gDiskArrivalCallbackCount >= MAX_DISK_ARRIVAL_CALLBACKS)
+    if (gDiskArrivalCallbackCount >= MAX_DISK_CALLBACKS)
         return STATUS_NO_MEMORY;
 
-    IO_DISK_CALLBACK_ENTRY* entry = &gDiskArrivalCallbacks[gDiskArrivalCallbackCount++];
+    IO_DISK_ARRIVAL_CALLBACK_ENTRY* entry = &gDiskArrivalCallbacks[gDiskArrivalCallbackCount++];
     memset(entry, 0, sizeof(*entry));
     entry->Used = true;
     entry->Callback = callback;
@@ -191,7 +632,21 @@ __declspec(dllexport) KSTATUS IoRegisterDiskArrivalCallback(IO_DISK_ARRIVAL_CALL
         if (gDiskDevices[i])
             callback(i, gDiskDevices[i], context);
     }
+    return STATUS_SUCCESS;
+}
 
+__declspec(dllexport) KSTATUS IoRegisterDiskChangeCallback(IO_DISK_CHANGE_CALLBACK callback,
+                                                           void* context) {
+    if (!callback)
+        return STATUS_INVALID_PARAMETER;
+    if (gDiskChangeCallbackCount >= MAX_DISK_CALLBACKS)
+        return STATUS_NO_MEMORY;
+
+    IO_DISK_CHANGE_CALLBACK_ENTRY* entry = &gDiskChangeCallbacks[gDiskChangeCallbackCount++];
+    memset(entry, 0, sizeof(*entry));
+    entry->Used = true;
+    entry->Callback = callback;
+    entry->Context = context;
     return STATUS_SUCCESS;
 }
 
@@ -204,7 +659,6 @@ __declspec(dllexport) KSTATUS IoEnumerateDiskDevices(IO_DISK_ARRIVAL_CALLBACK ca
         if (gDiskDevices[i])
             callback(i, gDiskDevices[i], context);
     }
-
     return STATUS_SUCCESS;
 }
 
@@ -213,28 +667,59 @@ __declspec(dllexport) DEVICE_OBJECT* IoGetDiskDevice(uint8_t driveIndex) {
     return gDiskDevices[driveIndex];
 }
 
+__declspec(dllexport) DEVICE_OBJECT* IoReferenceDiskDevice(uint8_t driveIndex) {
+    DEVICE_OBJECT* dev = IoGetDiskDevice(driveIndex);
+    if (dev) IoReferenceDeviceObject(dev);
+    return dev;
+}
+
 __declspec(dllexport) bool IoDiskDevicePresent(uint8_t driveIndex) {
     return IoGetDiskDevice(driveIndex) != (DEVICE_OBJECT*)0;
 }
 
 __declspec(dllexport) uint64_t IoDiskSectorCount(uint8_t driveIndex) {
     DEVICE_OBJECT* dev = IoGetDiskDevice(driveIndex);
-    return dev ? dev->SectorCount : 0;
+    return dev ? dev->u.Disk.SectorCount : 0;
 }
 
+// -------------------------------------------------------------------------
+// Dispatch
+// -------------------------------------------------------------------------
+
 __declspec(dllexport) KSTATUS IoCallDriver(DEVICE_OBJECT* device, IRP* irp) {
     if (!device || !device->DriverObject || !irp)
         return STATUS_INVALID_PARAMETER;
 
-    IO_MAJOR_FUNCTION major = (irp->Flags & IRP_FLAG_WRITE) ? IRP_MJ_WRITE : IRP_MJ_READ;
+    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
+    if (!ios)
+        return STATUS_INVALID_PARAMETER;
+
+    IO_MAJOR_FUNCTION major = ios->MajorFunction;
     if (major >= IRP_MJ_MAXIMUM_FUNCTION)
         return STATUS_INVALID_PARAMETER;
 
+    if (major != IRP_MJ_PNP &&
+        (device->DeletePending || device->PnpState == DevicePnpSurpriseRemoved ||
+         device->PnpState == DevicePnpRemovePending || device->PnpState == DevicePnpDeleted))
+        return STATUS_OBJECT_DEAD;
+
     PDRIVER_DISPATCH dispatch = device->DriverObject->MajorFunction[major];
     if (!dispatch)
         return STATUS_NOT_SUPPORTED;
 
-    return dispatch(device, irp);
+    ios->DeviceObject = device;
+    IoReferenceDeviceObject(device);
+    _InterlockedIncrement(&device->OutstandingIrps);
+
+    KSTATUS st = dispatch(device, irp);
+    if (st != STATUS_PENDING) {
+        if (ios->DeviceObject == device) {
+            ios->DeviceObject = (DEVICE_OBJECT*)0;
+            _InterlockedDecrement(&device->OutstandingIrps);
+            IoDereferenceDeviceObject(device);
+        }
+    }
+    return st;
 }
 
 __declspec(dllexport) KSTATUS IoReadDiskSectors(uint8_t driveIndex,
@@ -243,17 +728,15 @@ __declspec(dllexport) KSTATUS IoReadDiskSectors(uint8_t driveIndex,
                                                 void* buffer) {
     if (!buffer || count == 0) return STATUS_INVALID_PARAMETER;
 
-    IRP irp;
-    memset(&irp, 0, sizeof(irp));
-    irp.DriveIndex = driveIndex;
-    irp.Flags = 0;
-    irp.Lba = lba;
-    irp.SectorCount = count;
-    irp.Buffer = buffer;
-    irp.Status = STATUS_FAILURE;
-    irp.Information = 0;
+    DEVICE_OBJECT* dev = IoReferenceDiskDevice(driveIndex);
+    if (!dev) return STATUS_NOT_FOUND;
 
-    return IoSubmitIrpSync(&irp);
+    IRP irp;
+    IoBuildReadWriteIrp(&irp, IRP_MJ_READ, lba * 512ull, count * 512u, buffer);
+    KSTATUS st = IoSubmitIrpSync(dev, &irp);
+
+    IoDereferenceDeviceObject(dev);
+    return st;
 }
 
 __declspec(dllexport) KSTATUS IoWriteDiskSectors(uint8_t driveIndex,
@@ -262,15 +745,13 @@ __declspec(dllexport) KSTATUS IoWriteDiskSectors(uint8_t driveIndex,
                                                  const void* buffer) {
     if (!buffer || count == 0) return STATUS_INVALID_PARAMETER;
 
-    IRP irp;
-    memset(&irp, 0, sizeof(irp));
-    irp.DriveIndex = driveIndex;
-    irp.Flags = IRP_FLAG_WRITE;
-    irp.Lba = lba;
-    irp.SectorCount = count;
-    irp.Buffer = (void*)buffer;
-    irp.Status = STATUS_FAILURE;
-    irp.Information = 0;
+    DEVICE_OBJECT* dev = IoReferenceDiskDevice(driveIndex);
+    if (!dev) return STATUS_NOT_FOUND;
 
-    return IoSubmitIrpSync(&irp);
+    IRP irp;
+    IoBuildReadWriteIrp(&irp, IRP_MJ_WRITE, lba * 512ull, count * 512u, (void*)buffer);
+    KSTATUS st = IoSubmitIrpSync(dev, &irp);
+
+    IoDereferenceDeviceObject(dev);
+    return st;
 }