Windows 内核编程【驱动程序：从头到尾】

起因

本章主要因为Windows的API SetThreadPriority功能太弱，只能对线程设置几个级别，于是调用驱动的APIKeSetPriorityThread来进行设置线程优先级别，于是客户端与驱动的代码如下

代码实现

客户端

#include <stdio.h>
#include <Windows.h>

#define CTL_CODE(DeviceType,Function,Method,Access)(\
((DeviceType) << 16) | ((Access)<<14) | ((Function) << 2) | Method)
#define PRIORITY_BOOSTER_DEVICE 0x8000
#define IOCTL_PRIORITY_BOOSTER_SET_PRIORITY CTL_CODE(PRIORITY_BOOSTER_DEVICE,0x800,METHOD_NEITHER,FILE_ANY_ACCESS)
typedef struct _ThreadData {
	ULONG ThreadId;		// 需要修改的线程的ID
	int Priority;		// 需要修改的优先级
}ThreadData, * PThreadData;

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>\r\n");
		return 0;
	}

	// 打开设备对象的句柄,该函数调用会触发IRP_MJ_CREATE分发例程,如果驱动未加载,则此处会--文件未找到
	HANDLE hDevice = CreateFileW(L"\\\\.\\PriorityVooster", GENERIC_ALL, FILE_SHARE_WRITE, nullptr, OPEN_EXISTING, 0, nullptr);
	// 构造缓冲区结构
	ThreadData data = { 0 };
	data.ThreadId = atoi(argv[1]);
	data.Priority = atoi(argv[2]);

	// 这里函数触发的是IRP_MJ_DEVICE_CONTROL分发例程
	DWORD ret = 0;
	BOOL success = DeviceIoControl(hDevice, IOCTL_PRIORITY_BOOSTER_SET_PRIORITY, &data, sizeof(data), nullptr, 0, &ret, nullptr);
	if (success) {
		printf("Priority change succedded!\n");
	}
	else {
		Error("Priority change failed!");
	}

	// 这里函数触发的是IRP_MJ_CLOSE分发例程
	CloseHandle(hDevice);

	return 0;
}

驱动

#include <ntifs.h>
#include <ntddk.h>
#include "PBCommon.h"
#define DEVICE_NAME L"\\Device\\PriorityBooster"
#define SYMBLE_LINK_NAME L"\\??\\PriorityVooster"

VOID DriverUnload(PDRIVER_OBJECT pDriverObject) {
	// 还原驱动所做的操作
	// 删除符号链接
	UNICODE_STRING symLink = RTL_CONSTANT_STRING(SYMBLE_LINK_NAME);
	IoDeleteSymbolicLink(&symLink);
	IoDeleteDevice(pDriverObject->DeviceObject);
	KdPrint(("Bye World!\n"));
}

NTSTATUS
PriorityBoosterCreateClose(
	_In_ struct _DEVICE_OBJECT* DeviceObject,
	_Inout_ struct _IRP* Irp
) {
	Irp->IoStatus.Status = STATUS_SUCCESS;	// 指用什么状态完成该请求
	Irp->IoStatus.Information = 0;		// 不同请求有不同含义,这里请求需要设置为0
	IoCompleteRequest(Irp, IO_NO_INCREMENT);	// 完成IRP响应,会把IRP传回给创建者(IO管理器),然后管理器通知客户程序操作完成情况
	return STATUS_SUCCESS;
}

NTSTATUS
PriorityBoosterDeviceControl(
	_In_ struct _DEVICE_OBJECT* DeviceObject,
	_Inout_ struct _IRP* Irp
) {
	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 = (PThreadData)stack->Parameters.DeviceIoControl.Type3InputBuffer;
		if (data == nullptr) {
			status = STATUS_INVALID_PARAMETER;
			break;
		}
		// 参数合法性检查
		if (data->Priority < 1 || data->Priority>31) {
			status = STATUS_INVALID_PARAMETER;
			break;
		}
		// 将线程id转换成句柄,然后通过API获取线程结构,该API会增加线程结构的引用计数,需要手动减少
		PETHREAD pEthread;
		status = PsLookupThreadByThreadId(ULongToHandle(data->ThreadId), &pEthread);
		if (!NT_SUCCESS(status)) {
			break;
		}
		// 修改优先级
		KeSetPriorityThread(pEthread, data->Priority);
		
		// 手动减少引用计数
		ObDereferenceObject(pEthread);

		break;
	}
	default: {
		status = STATUS_INVALID_DEVICE_REQUEST;
		break;
	}
	}
	Irp->IoStatus.Status = status;
	Irp->IoStatus.Information = 0;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return status;
}

EXTERN_C NTSTATUS DriverEntry(PDRIVER_OBJECT pDriverObject, PUNICODE_STRING pReg) {
	// 指定Unload例程
	pDriverObject->DriverUnload = DriverUnload;
	// 操作设备对象至少需要CREATE和CLOSE操作分发例程
	pDriverObject->MajorFunction[IRP_MJ_CREATE] = PriorityBoosterCreateClose;
	pDriverObject->MajorFunction[IRP_MJ_CLOSE] = PriorityBoosterCreateClose;
	// 通过设备对象使得信息能从R3传入驱动,需要使用DeviceIoControl,该API通知的是DEVICE_CONTROL
	// 通过这种方法进行通信需要控制码和输入缓冲区
	pDriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = PriorityBoosterDeviceControl;

	// 创建设备对象,经典的驱动程序只需要一个设备对象,通过符号链接指向它即可
	UNICODE_STRING devName = RTL_CONSTANT_STRING(DEVICE_NAME);
	PDEVICE_OBJECT pDeviceObject = NULL;
	NTSTATUS ntStatus = IoCreateDevice(pDriverObject, 0, &devName, FILE_DEVICE_UNKNOWN, 0, FALSE, &pDeviceObject);

	// 提供符号链接
	UNICODE_STRING symLink = RTL_CONSTANT_STRING(SYMBLE_LINK_NAME);
	ntStatus = IoCreateSymbolicLink(&symLink, &devName);
	if (!NT_SUCCESS(ntStatus)) {
		// 如果失败了,需要回退之前的操作,因为DriverEntry中返回了非成功的状态,DriverUnload是不会被调用执行的
		KdPrint(("Failed to create symbolic link (0x%08x)\n", ntStatus));
		IoDeleteDevice(pDeviceObject);
		return ntStatus;
	}


	KdPrint(("Hello World!\n"));
	return STATUS_SUCCESS;
}

驱动头文件

#pragma once
#include <ntifs.h>
// 参数:
// DeviceType: 设备类型,第三方驱动从0x8000开始
// Function: 用来指明操作类型,对于第三方驱动从0x800开始
// Method: 指明用户提供的输入输出缓冲区是如何传入驱动的
// Access: 指明是到达驱动程序还是来自驱动程序
#define CTL_CODE(DeviceType,Function,Method,Access)(\
((DeviceType) << 16) | ((Access)<<14) | ((Function) << 2) | Method)
#define PRIORITY_BOOSTER_DEVICE 0x8000
#define IOCTL_PRIORITY_BOOSTER_SET_PRIORITY CTL_CODE(PRIORITY_BOOSTER_DEVICE,0x800,METHOD_NEITHER,FILE_ANY_ACCESS)
typedef struct _ThreadData {
	ULONG ThreadId;		// 需要修改的线程的ID
	int Priority;		// 需要修改的优先级
}ThreadData, * PThreadData;

注意事项

所有代码使用的是C++，所以如果使用C文件的话例如nullptr等会出现未识别的标识符

参考资料

https://www.kn0sky.com/?p=105
《Windows 内核编程》作者: [美] 帕维尔·约西福维奇（Pavel Yosifovich）