coolcore.c

CPU降温程序

#define _INC_STRING
#include <ntddk.h>
#include "..\commondef.h"

//////////////////////////////////////////////////////////////////////////////////////////////
NTHALAPI VOID NTAPI HalProcessorIdle(VOID);
NTKERNELAPI NTSTATUS NTAPI ZwYieldExecution(VOID);
NTKERNELAPI KAFFINITY NTAPI KeSetAffinityThread( PRKTHREAD prkThread, KAFFINITY kAffinity );
NTKERNELAPI LPWSTR __cdecl wcscpy( LPWSTR lpDest, LPCWSTR lpSrc );
#pragma function( wcscpy )
//////////////////////////////////////////////////////////////////////////////////////////////

const WCHAR wszDeviceName[] = L"\\Device\\CTQY_CoolCPU";

const LONG LOW_BASEPRIOR = -16;

//////////////////////////////////////////////////////////////////////////////////////////////
LONG g_lNormal_BasePrior;

volatile void (__stdcall *g_pfnCoolProc)(void);
volatile BOOLEAN g_bRun;

ULONG g_uThreadNum;
PRKTHREAD g_prkThread[32];

//////////////////////////////////////////////////////////////////////////////////////////////
#define NormalCoolProc HalProcessorIdle

void __stdcall IdleCoolProc(void)
{
	HalProcessorIdle();
	ZwYieldExecution();
}//end IdleCoolProc

void CoolThreadProc( PVOID pParam )
{
	while( g_bRun )
		{
		g_pfnCoolProc();
		}//end while

	PsTerminateSystemThread( STATUS_SUCCESS );
}//end CoolThreadProc

NTSTATUS DispatchCreate( IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp )
{
	// CreateFile was called, to get device handle
	// We are in user process context here
	HANDLE hThread;
	KAFFINITY kAffinity;
	KAFFINITY bit;

	g_pfnCoolProc = IdleCoolProc;
	g_bRun = TRUE;

	kAffinity = KeQueryActiveProcessors();
	for( g_uThreadNum = 0; kAffinity != 0; kAffinity ^= bit )
		{
		bit = kAffinity & ( -(*(LONG*)&kAffinity) ); // get rightest bit
		if( PsCreateSystemThread( &hThread,
								  THREAD_ALL_ACCESS,
								  NULL,
								  NULL,
								  NULL,
								  CoolThreadProc,
								  NULL ) == STATUS_SUCCESS )
			{
			ObReferenceObjectByHandle( hThread,
									   THREAD_ALL_ACCESS,
									   NULL,
									   KernelMode,
									   &g_prkThread[g_uThreadNum],
									   NULL );
			ZwClose( hThread );
			KeSetAffinityThread( g_prkThread[g_uThreadNum], bit );
			g_lNormal_BasePrior = KeSetBasePriorityThread( g_prkThread[g_uThreadNum], LOW_BASEPRIOR );
			++g_uThreadNum;
			}//end if
		}//end for

	pIrp->IoStatus.Status = STATUS_SUCCESS;
	pIrp->IoStatus.Information = 0;
	IoCompleteRequest( pIrp, IO_NO_INCREMENT );
	return STATUS_SUCCESS;
}//end DispatchCreate

NTSTATUS DispatchDeviceControl( IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp )
{
	// DeviceIoControl was called
	// We are in user process context here
	NTSTATUS status = STATUS_SUCCESS;
	ULONG uBytesReturned = 0;
	ULONG i;

	switch( IoGetCurrentIrpStackLocation( pIrp )->Parameters.DeviceIoControl.IoControlCode )
		{
		case IOCTL_HIGHPRIOR:
			{
			for( i = 0; i < g_uThreadNum; ++i )
				{
				KeSetBasePriorityThread( g_prkThread[i], g_lNormal_BasePrior );
				}//end for
			}break;//end IOCTL_HIGHPRIOR

		case IOCTL_LOWPRIOR:
			{
			for( i = 0; i < g_uThreadNum; ++i )
				{
				KeSetBasePriorityThread( g_prkThread[i], LOW_BASEPRIOR );
				}//end for
			}break;//end IOCTL_LOWPRIOR

		case IOCTL_GRABCPU:
			{
			g_pfnCoolProc = HalProcessorIdle;
			}break;//end IOCTL_GRABCPU

		case IOCTL_UNGRABCPU:
			{
			g_pfnCoolProc = IdleCoolProc;
			}break;//end IOCTL_UNGRABCPU

		default:
			{
			status = STATUS_INVALID_DEVICE_REQUEST;
			}
		}//end switch

	pIrp->IoStatus.Status = status;
	pIrp->IoStatus.Information = uBytesReturned;
	IoCompleteRequest( pIrp, IO_NO_INCREMENT );
	return STATUS_SUCCESS;
}//end DispatchDeviceControl

NTSTATUS DispatchClose( IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp )
{
	// CloseHandle was called, to close device handle
	// We are in user process context here
	ULONG i;

	g_bRun = FALSE;
	KeWaitForMultipleObjects( g_uThreadNum, g_prkThread, WaitAll, UserRequest, KernelMode, FALSE, NULL, NULL );
	for( i = 0; i < g_uThreadNum; ++i )
		{
		ObDereferenceObject( g_prkThread[i] );
		}//end for

	pIrp->IoStatus.Status = STATUS_SUCCESS;
	pIrp->IoStatus.Information = 0;
	IoCompleteRequest( pIrp, IO_NO_INCREMENT );
	return STATUS_SUCCESS;
}//end DispatchClose

VOID DriverUnload( IN PDRIVER_OBJECT pDriverObject )
{
	// ControlService(,,SERVICE_CONTROL_STOP) (ZwUnloadDriver) was called
	// We are in System process (pid = 4) context here
	UNICODE_STRING usSymbolicLinkName;

	RtlInitUnicodeString( &usSymbolicLinkName, wszSymbolicLinkName );
	IoDeleteSymbolicLink( &usSymbolicLinkName );
	IoDeleteDevice( pDriverObject->DeviceObject );
}//end DriverUnload

#pragma code_seg( "INIT" )
NTSTATUS DriverEntry( IN PDRIVER_OBJECT pDriverObject, IN PUNICODE_STRING pusRegistryPath )
{
	// StartService(ZwLoadDriver) was called
	// We are in System process (pid = 4) context here

	WCHAR szSymbolicLink[128];
	UNICODE_STRING usDeviceName;
	UNICODE_STRING usSymbolicLinkName;
	PDEVICE_OBJECT pDeviceObject;
	NTSTATUS result;
	NTSTATUS status = STATUS_DEVICE_CONFIGURATION_ERROR;

	RtlInitUnicodeString( &usDeviceName, wszDeviceName );
	if( IoCreateDevice( pDriverObject,
						0,
						&usDeviceName,
						FILE_DEVICE_UNKNOWN,
						0,
						FALSE,
						&pDeviceObject ) == STATUS_SUCCESS )
		{
		wcscpy( szSymbolicLink, L"\\??\\" );
		wcscpy( szSymbolicLink + 4, wszSymbolicLinkName );
		RtlInitUnicodeString( &usSymbolicLinkName, szSymbolicLink );

csl:	result = IoCreateSymbolicLink( &usSymbolicLinkName, &usDeviceName );
		switch( result )
			{
			case STATUS_SUCCESS:
				{
				pDriverObject->MajorFunction[IRP_MJ_CREATE] = DispatchCreate;
				pDriverObject->MajorFunction[IRP_MJ_CLOSE] = DispatchClose;
				pDriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DispatchDeviceControl;
				pDriverObject->DriverUnload = DriverUnload;

				status = STATUS_SUCCESS;
				}break;//end STATUS_SUCCESS

			case STATUS_OBJECT_NAME_COLLISION:
				{
				IoDeleteSymbolicLink( &usSymbolicLinkName );
				goto csl;
				}break;//end STATUS_OBJECT_NAME_COLLISION

			default:
				{
				IoDeleteDevice( pDeviceObject );
				}
			}//end switch
		}//end if

	return status;
}//end DriverEntry
#pragma code_seg()