Merge pull request #8 from hubblo-org/multi-sockets

Driver: reorder code and add socket selection
hubblo-org · Oct 19, 2023 · 27ef035 · 27ef035
2 parents 659a54b + 91620ca
commit 27ef035
Show file tree

Hide file tree

Showing 4 changed files with 83 additions and 39 deletions.
diff --git a/ScaphandreDrv/Driver.c b/ScaphandreDrv/Driver.c
@@ -6,16 +6,13 @@
 #pragma alloc_text (INIT, DriverEntry)
 #endif
 
-NTSTATUS
-DriverEntry(
-    _In_ PDRIVER_OBJECT  DriverObject,
-    _In_ PUNICODE_STRING RegistryPath
-    )
+NTSTATUS DriverEntry(PDRIVER_OBJECT  DriverObject,
+                     PUNICODE_STRING RegistryPath)
 {
-    NTSTATUS status;
     PDEVICE_OBJECT device_object;
     UNICODE_STRING device_name;
     UNICODE_STRING sym_name;
+    NTSTATUS status;
 
     DbgPrint("Registry path address: %p\n", RegistryPath);
 
@@ -55,18 +52,23 @@ NTSTATUS DispatchCreate(PDEVICE_OBJECT device, PIRP irp)
     memset(manufacturer, 0, sizeof(manufacturer));
     __cpuid(cpu_regs, 0);
     memcpy(manufacturer, &cpu_regs[1], sizeof(unsigned __int32));
-    memcpy(manufacturer + sizeof(unsigned __int32), &cpu_regs[3], sizeof(unsigned __int32));
-    memcpy(manufacturer + 2 * sizeof(unsigned __int32), &cpu_regs[2], sizeof(unsigned __int32));
+    memcpy(manufacturer + sizeof(unsigned __int32), &cpu_regs[3],
+           sizeof(unsigned __int32));
+    memcpy(manufacturer + 2 * sizeof(unsigned __int32), &cpu_regs[2],
+           sizeof(unsigned __int32));
 
-    if (strncmp(manufacturer, "GenuineIntel", sizeof(manufacturer) - 1) == 0)
+    if (!strncmp(manufacturer, "GenuineIntel", sizeof(manufacturer) - 1))
         machine_type = E_MACHINE_INTEL;
-    else if (strncmp(manufacturer, "AMDisbetter!", sizeof(manufacturer) - 1) == 0)
+    else if (!strncmp(manufacturer, "AMDisbetter!", sizeof(manufacturer) - 1))
         machine_type = E_MACHINE_AMD;
-    else if (strncmp(manufacturer, "AuthenticAMD", sizeof(manufacturer) - 1) == 0)
+    else if (!strncmp(manufacturer, "AuthenticAMD", sizeof(manufacturer) - 1))
         machine_type = E_MACHINE_AMD;
     else
         machine_type = E_MACHINE_UNK;
 
+    /* Get the number of virtual processors */
+    max_processors = KeQueryMaximumProcessorCountEx(ALL_PROCESSOR_GROUPS);
+
     irp->IoStatus.Status = STATUS_SUCCESS;
     irp->IoStatus.Information = STATUS_SUCCESS;
     IofCompleteRequest(irp, IO_NO_INCREMENT);
@@ -96,44 +98,75 @@ NTSTATUS DispatchCleanup(PDEVICE_OBJECT device, PIRP irp)
 
 NTSTATUS DispatchDeviceControl(PDEVICE_OBJECT device, PIRP irp)
 {
-    NTSTATUS ntStatus;
-    UINT32 msrRegister;
-    ULONG inputBufferLength;
-    ULONG outputBufferLength;
+    GROUP_AFFINITY affinity, old;
+    PIO_STACK_LOCATION stackLoc;
+    PROCESSOR_NUMBER pnumber;
     ULONGLONG msrResult;
-    PIO_STACK_LOCATION stackLocation;
+    NTSTATUS ntStatus;
+    struct data data;
+    size_t inLength;
 
-    stackLocation = irp->Tail.Overlay.CurrentStackLocation;
-    inputBufferLength = stackLocation->Parameters.DeviceIoControl.InputBufferLength;
-    outputBufferLength = stackLocation->Parameters.DeviceIoControl.OutputBufferLength;
+    stackLoc = irp->Tail.Overlay.CurrentStackLocation;
+    inLength = stackLoc->Parameters.DeviceIoControl.InputBufferLength;
 
     DbgPrint("Received event for driver %s... \n", device->DriverObject->DriverName);
 
-    /* METHOD_BUFFERED */
-    if (inputBufferLength == sizeof(ULONGLONG))
-    {
-        /* MSR register codes provided by userland must not exceed 8 bytes */
-        memcpy(&msrRegister, irp->AssociatedIrp.SystemBuffer, sizeof(ULONGLONG));
-        if (validate_msr_lookup(msrRegister) != 0)
-        {
-            DbgPrint("Requested MSR register (%08x) access is not allowed!\n", msrRegister);
-            ntStatus = STATUS_INVALID_DEVICE_REQUEST;
-        }
-        else
-        {
-            /* Call readmsr instruction */
-            msrResult = __readmsr(msrRegister);
-            memcpy(irp->AssociatedIrp.SystemBuffer, &msrResult, sizeof(ULONGLONG));
-            ntStatus = STATUS_SUCCESS;
-            irp->IoStatus.Information = sizeof(ULONGLONG);
-        }
+    if (inLength != sizeof(data)) {
+        DbgPrint("Bad input length provided. Expected %zu bytes, got %zu.\n",
+                 sizeof(data), inLength);
+        ntStatus = STATUS_INVALID_DEVICE_REQUEST;
+        goto error;
     }
-    else
+
+    /* Convert input data into structure */
+    memcpy(&data, irp->AssociatedIrp.SystemBuffer, sizeof(data));
+    if (validate_msr_lookup(data.msrRegister) != 0)
     {
-        DbgPrint("Bad input length provided. Expected %u bytes, got %u.\n", sizeof(ULONGLONG), inputBufferLength);
+        DbgPrint("Requested MSR register (%04x) access is not allowed!\n",
+                 data.msrRegister);
         ntStatus = STATUS_INVALID_DEVICE_REQUEST;
+        goto error;
     }
 
+    /* Check the requested CPU index is not above the maximum CPU available
+     * Otherwise, cap to max_processors - 1
+     */
+    data.cpuIndex = (data.cpuIndex >= max_processors) ? (max_processors - 1) : data.cpuIndex;
+
+    /* Run code on the specified socket
+     * MSDN says: 'Systems with fewer than 64 logical processors always have a single group, Group 0.'
+     * https://learn.microsoft.com/en-us/windows/win32/procthread/processor-groups
+     */
+    if ((ntStatus = KeGetProcessorNumberFromIndex(data.cpuIndex, &pnumber))
+        != STATUS_SUCCESS) {
+        DbgPrint("Failed to get processor info!\n");
+        goto error;
+    }
+
+    try {
+        /* Set affinity */
+        memset(&affinity, 0, sizeof(GROUP_AFFINITY));
+        affinity.Group = pnumber.Group;
+        affinity.Mask = data.cpuIndex;
+        KeSetSystemGroupAffinityThread(&affinity, &old);
+
+        /* Call readmsr instruction */
+        msrResult = __readmsr(data.msrRegister);
+
+        /* Restore affinity */
+        KeRevertToUserGroupAffinityThread(&old);
+    }
+    except (EXCEPTION_EXECUTE_HANDLER) {
+        ntStatus = STATUS_IO_DEVICE_ERROR;
+        goto error;
+    }
+
+    /* Save result */
+    memcpy(irp->AssociatedIrp.SystemBuffer, &msrResult, sizeof(data));
+    irp->IoStatus.Information = sizeof(data);
+    ntStatus = STATUS_SUCCESS;
+
+error:
     irp->IoStatus.Status = ntStatus;
     IofCompleteRequest(irp, IO_NO_INCREMENT);
 

diff --git a/ScaphandreDrv/Driver.h b/ScaphandreDrv/Driver.h
@@ -29,6 +29,11 @@ DRIVER_INITIALIZE DriverEntry;
 
 EXTERN_C_END
 
+struct data {
+    UINT32 msrRegister;
+    UINT32 cpuIndex;
+};
+
 void DriverUnload(PDRIVER_OBJECT driver);
 NTSTATUS DispatchCreate(PDEVICE_OBJECT device, PIRP irp);
 NTSTATUS DispatchClose(PDEVICE_OBJECT device, PIRP irp);

diff --git a/ScaphandreDrv/msr.c b/ScaphandreDrv/msr.c
@@ -25,6 +25,8 @@ int validate_msr_lookup(unsigned __int64 msrRegister)
                 case MSR_PP1_POWER_LIMIT:
                 case MSR_PP1_ENERGY_STATUS:
                 case MSR_PP1_POLICY:
+                case MSR_PLATFORM_ENERGY_STATUS:
+                case MSR_PLATFORM_POWER_LIMIT:
                     err = 0;
                     break;
 

diff --git a/ScaphandreDrv/msr.h b/ScaphandreDrv/msr.h
@@ -28,6 +28,9 @@
 #define MSR_AMD_CORE_ENERGY_STATUS  0xc001029a
 #define MSR_AMD_PKG_ENERGY_STATUS   0xc001029b
 
+#define MSR_PLATFORM_ENERGY_STATUS 0x0000064d
+#define MSR_PLATFORM_POWER_LIMIT 0x0000065c
+
 
 typedef enum {
     E_MACHINE_INTEL,
@@ -37,6 +40,7 @@ typedef enum {
 
 
 static e_machine_type machine_type;
+static unsigned long max_processors;
 
 int validate_msr_lookup(unsigned __int64 msrRegister);