一個小小工程師的心情抒發天地

Driver-----


#include 
#include 
#include "PriorityBoosterCommon.h"

NTSTATUS BoosterCreateClose(PDEVICE_OBJECT DeviceObject, PIRP Irp);
NTSTATUS BoosterWrite(PDEVICE_OBJECT DeviceObject, PIRP Irp);
void BoosterUnload(PDRIVER_OBJECT DriverObject);


extern "C" NTSTATUS
DriverEntry(_In_ PDRIVER_OBJECT DriverObject, _In_ PUNICODE_STRING RegistryPath)
{
	UNREFERENCED_PARAMETER(RegistryPath);
	KdPrint(("Boster: DriverEntry \n"));


	DriverObject->DriverUnload = BoosterUnload;

	DriverObject->MajorFunction[IRP_MJ_CREATE] = BoosterCreateClose;
	DriverObject->MajorFunction[IRP_MJ_CLOSE] = BoosterCreateClose;
	DriverObject->MajorFunction[IRP_MJ_WRITE] = BoosterWrite;


	UNICODE_STRING devName = RTL_CONSTANT_STRING(L"\\Device\\Booster");

	PDEVICE_OBJECT DeviceObject;
	NTSTATUS status = IoCreateDevice
	(
		DriverObject,			// our driver object
		0,						// no need for extra bytes
		&devName,				// the device name
		FILE_DEVICE_UNKNOWN,	// device type
		0,						// characteristics flags
		FALSE,					// not exclusive
		&DeviceObject			// the resulting pointer
	);
	if (!NT_SUCCESS(status)) {
		KdPrint(("Failed to create device object (0x%08X)\n", status));
		return status;
	}

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

 
void BoosterUnload(PDRIVER_OBJECT DriverObject)
{
	KdPrint(("Boster: Driver unload\n"));

	UNICODE_STRING symLink = RTL_CONSTANT_STRING(L"\\??\\Booster");
	// delete symbolic link
	IoDeleteSymbolicLink(&symLink);

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


NTSTATUS BoosterCreateClose(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;
}

NTSTATUS BoosterWrite(PDEVICE_OBJECT, PIRP Irp) {
	auto status = STATUS_SUCCESS;
	ULONG_PTR information = 0;


	auto irpSp = IoGetCurrentIrpStackLocation(Irp);
	do {
		if (irpSp->Parameters.Write.Length < sizeof(ThreadData)) {
			status = STATUS_BUFFER_TOO_SMALL;
			break;
		}

		auto data = static_cast(Irp->UserBuffer);
		if (data == nullptr || data->Priority < 1 || data->Priority > 31) {
			status = STATUS_INVALID_PARAMETER;
			break;
		}

		PETHREAD thread;
		status = PsLookupThreadByThreadId(ULongToHandle(data->ThreadId), &thread);
		if (!NT_SUCCESS(status))
		{
			break;
		}
		auto oldPriority = KeSetPriorityThread(thread, data->Priority);
		KdPrint(("Priority change ofr thread %u from %d to %d succeeded!\n",
			data->ThreadId, oldPriority, data->Priority));

		ObDereferenceObject(thread);
		information = sizeof(ThreadData);
	} while (false);

	Irp->IoStatus.Status = status;
	Irp->IoStatus.Information = information;

	IoCompleteRequest(Irp, IO_NO_INCREMENT);

	return status;

}

AP-----

// PriorityBoosterApp.cpp : This file contains the 'main' function. Program execution begins and ends there.
//

#include 
#include 
#include 
#include 

#include "..\PriorityBooster\PriorityBoosterCommon.h"

#pragma comment(lib, "ntdll")
using namespace std;

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


int main()
{

	int tid = 0;
	int priority = 0;

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

	ThreadData data;
	data.ThreadId = tid;
	data.Priority = priority;

	DWORD returned;
	BOOL success = WriteFile(hDevice,
		&data, sizeof(data),			// buffer and length
		&returned, nullptr);

	if (!success)
		return Error("Priority change failed");

	printf("Priority change succeeded!\n");
	CloseHandle(hDevice);

	return 0;
}