Windows内核驱动--实现修改线程优先级demo

在User下修改优先级比较麻烦,该驱动能够直接用线程ID,和优先级级数两个参数直接修改线程的优先级:spa

 

Client代码:prototype

#include <Windows.h>
#include <stdio.h>
#include "..\PriorityBooster\PriorityBoosterCommon.h"

int Error(const char* message) {
    printf("%s (error=%d)\n", message, GetLastError());
    return 1;
}

int main(int argc, const char* argv[]) {
    if (argc < 3) {
        printf("Usage: Booster <threadid> <priority>\n");
        return 0;
    }

    HANDLE hDevice = CreateFile(L"\\\\.\\PriorityBooster", GENERIC_WRITE, FILE_SHARE_WRITE,
        nullptr, OPEN_EXISTING, 0, nullptr);
    if (hDevice == INVALID_HANDLE_VALUE)
        return Error("Failed to open device");

    ThreadData data;
    data.ThreadId = atoi(argv[1]);
    data.Priority = atoi(argv[2]);

    DWORD returned;
    BOOL success = DeviceIoControl(hDevice, IOCTL_PRIORITY_BOOSTER_SET_PRIORITY, &data, sizeof(data), nullptr, 0, &returned, nullptr);
    if (success)
        printf("Priority change succeeded!\n");
    else
        Error("Priority change failed!");

    CloseHandle(hDevice);
}

 

kernel代码:线程

//PriorityBooster.cpp
#include <ntifs.h>
#include <ntddk.h>
#include "PriorityBoosterCommon.h"

// prototypes

void PriorityBoosterUnload(_In_ PDRIVER_OBJECT DriverObject);
NTSTATUS PriorityBoosterCreateClose(_In_ PDEVICE_OBJECT DeviceObject, _In_ PIRP Irp);
NTSTATUS PriorityBoosterDeviceControl(_In_ PDEVICE_OBJECT DeviceObject, _In_ PIRP Irp);

// DriverEntry

extern "C" NTSTATUS
DriverEntry(_In_ PDRIVER_OBJECT DriverObject, _In_ PUNICODE_STRING RegistryPath) {
    UNREFERENCED_PARAMETER(RegistryPath);

    KdPrint(("PriorityBooster DriverEntry started\n"));

    DriverObject->DriverUnload = PriorityBoosterUnload;

    DriverObject->MajorFunction[IRP_MJ_CREATE] = PriorityBoosterCreateClose;
    DriverObject->MajorFunction[IRP_MJ_CLOSE] = PriorityBoosterCreateClose;
    DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = PriorityBoosterDeviceControl;

    UNICODE_STRING devName = RTL_CONSTANT_STRING(L"\\Device\\PriorityBooster");
    //RtlInitUnicodeString(&devName, L"\\Device\\ThreadBoost");
    PDEVICE_OBJECT DeviceObject;
    NTSTATUS status = IoCreateDevice(DriverObject, 0, &devName, FILE_DEVICE_UNKNOWN, 0, FALSE, &DeviceObject);
    if (!NT_SUCCESS(status)) {
        KdPrint(("Failed to create device (0x%08X)\n", status));
        return status;
    }

    UNICODE_STRING symLink = RTL_CONSTANT_STRING(L"\\??\\PriorityBooster");
    status = IoCreateSymbolicLink(&symLink, &devName);
    if (!NT_SUCCESS(status)) {
        KdPrint(("Failed to create symbolic link (0x%08X)\n", status));
        IoDeleteDevice(DeviceObject);
        return status;
    }

    KdPrint(("PriorityBooster DriverEntry completed successfully\n"));

    return STATUS_SUCCESS;
}

void PriorityBoosterUnload(_In_ PDRIVER_OBJECT DriverObject) {
    UNICODE_STRING symLink = RTL_CONSTANT_STRING(L"\\??\\PriorityBooster");
    // delete symbolic link
    IoDeleteSymbolicLink(&symLink);

    // delete device object
    IoDeleteDevice(DriverObject->DeviceObject);

    KdPrint(("PriorityBooster unloaded\n"));
}

_Use_decl_annotations_
NTSTATUS PriorityBoosterCreateClose(PDEVICE_OBJECT DeviceObject, PIRP Irp) {
    UNREFERENCED_PARAMETER(DeviceObject);

    Irp->IoStatus.Status = STATUS_SUCCESS;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    return STATUS_SUCCESS;
}



_Use_decl_annotations_
NTSTATUS PriorityBoosterDeviceControl(PDEVICE_OBJECT, PIRP Irp) {
    // get our IO_STACK_LOCATION
    auto stack = IoGetCurrentIrpStackLocation(Irp);
    auto status = STATUS_SUCCESS;

    switch (stack->Parameters.DeviceIoControl.IoControlCode) {
        case IOCTL_PRIORITY_BOOSTER_SET_PRIORITY:
        {
            // do the work
            if (stack->Parameters.DeviceIoControl.InputBufferLength < sizeof(ThreadData)) {
                status = STATUS_BUFFER_TOO_SMALL;
                break;
            }

            auto data = (ThreadData*)stack->Parameters.DeviceIoControl.Type3InputBuffer;
            if (data == nullptr) {
                status = STATUS_INVALID_PARAMETER;
                break;
            }

            if (data->Priority < 1 || data->Priority > 31) {
                status = STATUS_INVALID_PARAMETER;
                break;
            }

            PETHREAD Thread;
            status = PsLookupThreadByThreadId(ULongToHandle(data->ThreadId), &Thread);
            if (!NT_SUCCESS(status))
                break;

            KeSetPriorityThread((PKTHREAD)Thread, data->Priority);
            ObDereferenceObject(Thread);
            KdPrint(("Thread Priority change for %d to %d succeeded!\n",
                data->ThreadId, data->Priority));
            break;
        }

        default:
            status = STATUS_INVALID_DEVICE_REQUEST;
            break;
    }

    Irp->IoStatus.Status = status;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    return status;
}
//PriorityBoosterCommon.h
#pragma once

struct ThreadData {
    ULONG ThreadId;
    int Priority;
};

#define PRIORITY_BOOSTER_DEVICE 0x8000

#define IOCTL_PRIORITY_BOOSTER_SET_PRIORITY CTL_CODE(PRIORITY_BOOSTER_DEVICE, \
    0x800, METHOD_NEITHER, FILE_ANY_ACCESS)

采用VS2019创建如图相似工程就能够了:code