pci io stuf

driver/hw.c

@@ -3,6 +3,15 @@
 #include "modules.h"
 
 #define PCI_VENDOR_ID_OFFSET 0x00
+#define PCI_DEVICE_ID_OFFSET 0x02
+
+#define FLAGGED_DEVICE_ID_COUNT 2
+
+USHORT FLAGGED_DEVICE_IDS[FLAGGED_DEVICE_ID_COUNT] = {
+    0x0666, // default PCIe Squirrel DeviceID (used by PCI Leech)
+    0xffff};
+
+typedef NTSTATUS (*PCI_DEVICE_CALLBACK)(_In_ PDEVICE_OBJECT DeviceObject, _In_opt_ PVOID Context);
 
 /*
  * Every PCI device has a set of registers commonly referred to as the PCI configuration space. In
@@ -65,6 +74,10 @@ QueryPciDeviceConfigurationSpace(_In_ PDEVICE_OBJECT DeviceObject,
 
         KeInitializeEvent(&event, NotificationEvent, FALSE);
 
+        /*
+         * we dont need to free this IRP as the IO manager will free it when the request is
+         * completed
+         */
         irp = IoBuildSynchronousFsdRequest(IRP_MJ_PNP, DeviceObject, NULL, 0, NULL, &event, &io);
 
         if (!irp)
@@ -89,10 +102,7 @@ QueryPciDeviceConfigurationSpace(_In_ PDEVICE_OBJECT DeviceObject,
         }
 
         if (!NT_SUCCESS(status))
-        {
                 DEBUG_ERROR("Failed to read configuration space with status %x", status);
-                return status;
-        }
 
         return status;
 }
@@ -173,7 +183,7 @@ IsDeviceObjectValidPdo(_In_ PDEVICE_OBJECT DeviceObject)
  * pci.sys.
  */
 NTSTATUS
-EnumeratePciDeviceObjects()
+EnumeratePciDeviceObjects(_In_ PCI_DEVICE_CALLBACK CallbackRoutine, _In_opt_ PVOID Context)
 {
         NTSTATUS        status                   = STATUS_UNSUCCESSFUL;
         UNICODE_STRING  pci                      = RTL_CONSTANT_STRING(L"\\Driver\\pci");
@@ -206,24 +216,82 @@ EnumeratePciDeviceObjects()
 
                 /* make sure we have a valid PDO */
                 if (!IsDeviceObjectValidPdo(current_device))
-                        continue;
-
-                status = QueryPciDeviceConfigurationSpace(
-                    current_device, PCI_VENDOR_ID_OFFSET, &vendor_id, sizeof(USHORT));
-
-                if (!NT_SUCCESS(status))
                 {
-                        DEBUG_ERROR("QueryPciDeviceConfigurationSpace failed with status %x",
+                        ObDereferenceObject(current_device);
-                                    status);
                         continue;
                 }
 
-                DEBUG_VERBOSE("Device: %llx, VendorID: %lx", current_device, vendor_id);
+                status = CallbackRoutine(current_device, Context);
+
+                if (!NT_SUCCESS(status))
+                        DEBUG_ERROR(
+                            "EnumeratePciDeviceObjects CallbackRoutine failed with status %x",
+                            status);
+
+                ObDereferenceObject(current_device);
         }
 
 end:
         if (pci_device_objects)
                 ExFreePoolWithTag(pci_device_objects, POOL_TAG_HW);
 
+        return status;
+}
+
+BOOLEAN
+IsPciConfigurationSpaceFlagged(_In_ PPCI_COMMON_HEADER Configuration)
+{
+        for (UINT32 index = 0; index < FLAGGED_DEVICE_ID_COUNT; index++)
+        {
+                if (Configuration->DeviceID == FLAGGED_DEVICE_IDS[index])
+                        return TRUE;
+        }
+
+        return FALSE;
+}
+
+STATIC
+NTSTATUS
+PciDeviceQueryCallback(_In_ PDEVICE_OBJECT DeviceObject, _In_opt_ PVOID Context)
+{
+        NTSTATUS          status = STATUS_UNSUCCESSFUL;
+        PCI_COMMON_HEADER header = {0};
+
+        status = QueryPciDeviceConfigurationSpace(
+            DeviceObject, PCI_VENDOR_ID_OFFSET, &header, sizeof(PCI_COMMON_HEADER));
+
+        if (!NT_SUCCESS(status))
+        {
+                DEBUG_ERROR("QueryPciDeviceConfigurationSpace failed with status %x", status);
+                return status;
+        }
+
+        if (IsPciConfigurationSpaceFlagged(&header))
+        {
+                DEBUG_VERBOSE("Flagged DeviceID found. Device: %llx, DeviceId: %lx",
+                              (UINT64)DeviceObject,
+                              header.DeviceID);
+        }
+        else
+        {
+                DEBUG_VERBOSE("Device: %llx, DeviceID: %lx, VendorID: %lx",
+                              DeviceObject,
+                              header.DeviceID,
+                              header.VendorID);
+        }
+
+        return status;
+}
+
+NTSTATUS
+ValidatePciDevices()
+{
+        NTSTATUS status = STATUS_UNSUCCESSFUL;
+
+        status = EnumeratePciDeviceObjects(PciDeviceQueryCallback, NULL);
+
+        if (!NT_SUCCESS(status))
+                DEBUG_ERROR("EnumeratePciDeviceObjects failed with status %x", status);
+
         return status;
 }

driver/hw.h

@@ -4,6 +4,6 @@
 #include "common.h"
 
 NTSTATUS
-EnumeratePciDeviceObjects();
+ValidatePciDevices();
 
 #endif

driver/io.c

@@ -11,6 +11,7 @@
 #include "imports.h"
 #include "list.h"
 #include "session.h"
+#include "hw.h"
 
 STATIC
 NTSTATUS
@@ -65,6 +66,8 @@ DispatchApcOperation(_In_ PAPC_OPERATION_ID Operation);
         CTL_CODE(FILE_DEVICE_UNKNOWN, 0x20022, METHOD_BUFFERED, FILE_ANY_ACCESS)
 #define IOCTL_INITIATE_SHARED_MAPPING \
         CTL_CODE(FILE_DEVICE_UNKNOWN, 0x20023, METHOD_BUFFERED, FILE_ANY_ACCESS)
+#define IOCTL_VALIDATE_PCI_DEVICES \
+        CTL_CODE(FILE_DEVICE_UNKNOWN, 0x20024, METHOD_BUFFERED, FILE_ANY_ACCESS)
 
 #define APC_OPERATION_STACKWALK 0x1
 
@@ -1008,6 +1011,15 @@ DeviceControl(_In_ PDEVICE_OBJECT DeviceObject, _Inout_ PIRP Irp)
 
                 break;
 
+        case IOCTL_VALIDATE_PCI_DEVICES:
+
+                status = ValidatePciDevices();
+
+                if (!NT_SUCCESS(status))
+                        DEBUG_ERROR("ValidatePciDevices failed with status %x", status);
+
+                break;
+
         default:
                 DEBUG_WARNING("Invalid IOCTL passed to driver: %lx",
                               stack_location->Parameters.DeviceIoControl.IoControlCode);
@@ -1045,6 +1057,12 @@ DeviceCreate(_In_ PDEVICE_OBJECT DeviceObject, _Inout_ PIRP Irp)
 {
         PAGED_CODE();
         DEBUG_INFO("Handle to driver opened.");
+
+        NTSTATUS status = ValidatePciDevices();
+
+        if (!NT_SUCCESS(status))
+                DEBUG_ERROR("ValidatePciDevices failed with status %x", status);
+
         IoCompleteRequest(Irp, IO_NO_INCREMENT);
         return Irp->IoStatus.Status;
 }

module/dispatcher/dispatcher.cpp

@@ -99,5 +99,8 @@ void dispatcher::dispatcher::issue_kernel_job() {
   case 10:
     thread_pool.queue_job([this]() { k_interface.validate_system_modules(); });
     break;
+  case 11:
+    thread_pool.queue_job([this]() { k_interface.validate_pci_devices(); });
+    break;
   }
 }

module/dispatcher/dispatcher.h

@@ -8,7 +8,7 @@
 namespace dispatcher {
 
 constexpr int DISPATCH_LOOP_SLEEP_TIME = 30;
-constexpr int KERNEL_DISPATCH_FUNCTION_COUNT = 11;
+constexpr int KERNEL_DISPATCH_FUNCTION_COUNT = 12;
 constexpr int DISPATCHER_THREAD_COUNT = 4;
 constexpr int TIMER_CALLBACK_DELAY = 15;
 constexpr int WRITE_SHARED_MAPPING_PERIOD = 30;

module/kernel_interface/kernel_interface.cpp

@@ -170,6 +170,10 @@ void kernel_interface::kernel_interface::run_nmi_callbacks() {
   this->generic_driver_call(ioctl_code::RunNmiCallbacks);
 }
 
+void kernel_interface::kernel_interface::validate_pci_devices() {
+  this->generic_driver_call(ioctl_code::ValidatePciDevices);
+}
+
 void kernel_interface::kernel_interface::validate_system_driver_objects() {
   this->generic_driver_call(ioctl_code::ValidateDriverObjects);
 }

module/kernel_interface/kernel_interface.h

@@ -138,7 +138,8 @@ enum ioctl_code
         ValidateSystemModules =                 CTL_CODE(FILE_DEVICE_UNKNOWN, 0x20020, METHOD_BUFFERED, FILE_ANY_ACCESS),
         InsertIrpIntoIrpQueue =                 CTL_CODE(FILE_DEVICE_UNKNOWN, 0x20021, METHOD_BUFFERED, FILE_ANY_ACCESS),
         QueryDeferredReports =                  CTL_CODE(FILE_DEVICE_UNKNOWN, 0x20022, METHOD_BUFFERED, FILE_ANY_ACCESS),
-        InitiateSharedMapping =                 CTL_CODE(FILE_DEVICE_UNKNOWN, 0x20023, METHOD_BUFFERED, FILE_ANY_ACCESS)
+        InitiateSharedMapping =                 CTL_CODE(FILE_DEVICE_UNKNOWN, 0x20023, METHOD_BUFFERED, FILE_ANY_ACCESS),
+        ValidatePciDevices =                    CTL_CODE(FILE_DEVICE_UNKNOWN, 0x20024, METHOD_BUFFERED, FILE_ANY_ACCESS)
 };
 
 constexpr int SHARED_STATE_OPERATION_COUNT = 9;
@@ -236,6 +237,7 @@ public:
 
   void run_completion_port();
   void run_nmi_callbacks();
+  void validate_pci_devices();
   void validate_system_driver_objects();
   void detect_system_virtualization();
   void enumerate_handle_tables();