drivercontrol.cpp

CPU降温程序

#include "DriverControl.h"
#include <shlwapi.h>
#include <winioctl.h>
#include "..\commondef.h"

#pragma comment( lib, "shlwapi.lib" )

#ifdef __cplusplus
extern "C"{
#endif

typedef LONG NTSTATUS;
#define STATUS_SUCCESS				((NTSTATUS)0x00000000L)
#define STATUS_IMAGE_ALREADY_LOADED	((NTSTATUS)0xC000010EL)

typedef struct _UNICODE_STRING {
    USHORT Length;
    USHORT MaximumLength;
    PWSTR  Buffer;
} UNICODE_STRING, *PUNICODE_STRING;

NTSYSAPI NTSTATUS NTAPI ZwLoadDriver( IN PUNICODE_STRING DriverServiceName );
NTSYSAPI NTSTATUS NTAPI ZwUnloadDriver( IN PUNICODE_STRING DriverServiceName );

NTSYSAPI VOID NTAPI RtlInitUnicodeString( PUNICODE_STRING DestinationString, PCWSTR SourceString );

#ifdef __cplusplus
}//end extern "C"
#endif

const WCHAR wszSvcName[] = L"CoolCPUsvc";
const WCHAR wszRegKeyRoot[] = L"\\registry\\machine\\SYSTEM\\CurrentControlSet\\Services\\";
const TCHAR szRegSvcKeyRoot[] = TEXT("SYSTEM\\CurrentControlSet\\Services");

BOOL EnablePrivilege( LPCTSTR lpszPrivilege )
{
	TOKEN_PRIVILEGES tp;
	HANDLE			 hToken;

	if( !OpenProcessToken( GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken ) )
		{
		return FALSE;
		}//end if

	if( !LookupPrivilegeValue( NULL, lpszPrivilege, &tp.Privileges[0].Luid ) )
		{
		return FALSE;
		}//end if

	tp.PrivilegeCount = 1;
	tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

	if( !AdjustTokenPrivileges( hToken, FALSE, &tp, 0, NULL, 0 ) )
		{
		return FALSE;
		}//end if

	CloseHandle( hToken );

	return TRUE;
}//end EnablePrivilege

void GetDriverFullName( LPWSTR lpszDriverPath, LPCWSTR lpszDriverFileName )
{
	int i;

	wcscpy( lpszDriverPath, L"\\??\\" );
	lpszDriverPath += 4;
	GetModuleFileName( NULL, lpszDriverPath, MAX_PATH - 4 );
	for( i = wcslen( lpszDriverPath ) - 1; lpszDriverPath[i] != '\\'; --i );//end for
	++i;
	wcscpy( lpszDriverPath + i, lpszDriverFileName );
}//end GetDriverFullName

BOOL LoadDriver( LPCWSTR lpszDriverFile )
{
	WCHAR wszTemp[MAX_PATH];
	UNICODE_STRING usDevSvcName;
	HKEY hParentKey;
	HKEY hKey;
	DWORD dwType;
	NTSTATUS status;
	BOOL bRes = FALSE;

	if( RegOpenKeyEx( HKEY_LOCAL_MACHINE,
					  szRegSvcKeyRoot,
					  0,
					  KEY_WRITE,
					  &hParentKey ) == ERROR_SUCCESS )
		{
		if( RegCreateKeyExW( hParentKey, wszSvcName, 0, NULL, 0, KEY_WRITE, NULL, &hKey, NULL ) == ERROR_SUCCESS )
			{
			dwType = SERVICE_KERNEL_DRIVER;
			if( RegSetValueExW( hKey,
								L"ImagePath",
								0,
								REG_SZ,
								(LPBYTE)lpszDriverFile,
								sizeof(WCHAR) * ( wcslen( lpszDriverFile ) + 1 ) ) == ERROR_SUCCESS
				&& RegSetValueEx( hKey,
								  TEXT("Type"),
								  0,
								  REG_DWORD,
								  (LPBYTE)&dwType,
								  sizeof( DWORD ) ) == ERROR_SUCCESS )
				{
				wcscpy( wszTemp, wszRegKeyRoot );
				wcscat( wszTemp, wszSvcName );
				RtlInitUnicodeString( &usDevSvcName, wszTemp );

				status = ZwLoadDriver( &usDevSvcName );
				if( status == STATUS_SUCCESS || status == STATUS_IMAGE_ALREADY_LOADED )
					{
					bRes = TRUE;
					}//end if
				}//end if
			RegCloseKey( hKey );
			}//end if
		RegCloseKey( hParentKey );
		}//end if

	return bRes;
}//end LoadDriver

BOOL UnloadDriver(void)
{
	WCHAR wszTemp[MAX_PATH];
	UNICODE_STRING usDevSvcName;
	HKEY hKey;
	BOOL bRes = FALSE;

	wcscpy( wszTemp, wszRegKeyRoot );
	wcscat( wszTemp, wszSvcName );
	RtlInitUnicodeString( &usDevSvcName, wszTemp );

	if( ZwUnloadDriver( &usDevSvcName ) == STATUS_SUCCESS )
		{
		if( RegOpenKeyEx( HKEY_LOCAL_MACHINE,
						  szRegSvcKeyRoot,
						  0,
						  DELETE,
						  &hKey ) == ERROR_SUCCESS )
			{
			if( SHDeleteKeyW( hKey, wszSvcName ) == ERROR_SUCCESS )
				{
				bRes = TRUE;
				}//end if
			RegCloseKey( hKey );
			}//end if
		}//end if

	return bRes;
}//end UnloadDriver

HANDLE OpenCoolDevice(void)
{
	WCHAR szSymbolicLink[128] = { L"\\\\.\\" };

	wcscpy( szSymbolicLink + 4, wszSymbolicLinkName );
	return CreateFileW( szSymbolicLink, GENERIC_READ | GENERIC_WRITE, 0, NULL, OPEN_EXISTING, 0, NULL );
}//end OpenCoolDevice

void ChangeCoolPriority( HANDLE hDevice, BOOL bLowPrior )
{
	DWORD dwRet;
	DeviceIoControl( hDevice, (bLowPrior) ? IOCTL_LOWPRIOR : IOCTL_HIGHPRIOR, NULL, 0, NULL, 0, &dwRet, NULL );
}//end ChangeCoolPriority

void GrabCPU( HANDLE hDevice, BOOL bGrab )
{
	DWORD dwRet;
	DeviceIoControl( hDevice, (bGrab) ? IOCTL_GRABCPU : IOCTL_UNGRABCPU, NULL, 0, NULL, 0, &dwRet, NULL );
}//end ChangeCoolPriority