main.cpp

显示Device Treee

//********************************************************************
//	日期:	2005/02/23 - 23:2:2005   8:39
//	名前:	tiamo
//	描述:	main
//*********************************************************************

#include "stdafx.h"

// driver entry
NTSTATUS DriverEntry(PDRIVER_OBJECT pDriverObject,PUNICODE_STRING pRegPath)
{
	pDriverObject->MajorFunction[IRP_MJ_CREATE] = DispatchCreate;
	pDriverObject->MajorFunction[IRP_MJ_CLOSE] = DispatchClose;
	pDriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DispathDeviceIoControl;
	pDriverObject->DriverUnload = DriverUnload;

	NTSTATUS status = STATUS_SUCCESS;

	PDEVICE_OBJECT pDeviceObject = 0;

	UNICODE_STRING strDeviceName;
	UNICODE_STRING strSymbolicName;

	RtlInitUnicodeString(&strDeviceName,L"\\Device\\DeviceTreeSys");

	RtlInitUnicodeString(&strSymbolicName,L"\\??\\DeviceTreeSys");

	__try
	{
		status = IoCreateDevice(pDriverObject,sizeof(DEVICE_EXTENSION),&strDeviceName,
								FILE_DEVICE_UNKNOWN,0,0,&pDeviceObject);

		if(!NT_SUCCESS(status))
			ExRaiseStatus(status);

		PDEVICE_EXTENSION pExt;
		pExt = static_cast<PDEVICE_EXTENSION>(pDeviceObject->DeviceExtension);

		pExt->m_pSelf = pDeviceObject;

		status = IoCreateSymbolicLink(&strSymbolicName,&strDeviceName);

		if(!NT_SUCCESS(status))
			ExRaiseStatus(status);
	}
	__except(EXCEPTION_EXECUTE_HANDLER)
	{
		if(pDeviceObject)
			IoDeleteDevice(pDeviceObject);
	}

	devDebugPrint("DeviceTreeSys:DriverEntry exit %#x\n",status);

	return status;
}

// driver unload
VOID DriverUnload(PDRIVER_OBJECT pDriverObject)
{
	if(pDriverObject->DeviceObject)
	{
		UNICODE_STRING strSymbolicName;

		RtlInitUnicodeString(&strSymbolicName,L"\\??\\DeviceTreeSys");

		IoDeleteSymbolicLink(&strSymbolicName);

		IoDeleteDevice(pDriverObject->DeviceObject);
	}

	devDebugPrint("DeviceTreeSys:DriverUnload exit\n");
}

// iocontrol
NTSTATUS DispathDeviceIoControl(PDEVICE_OBJECT pDeviceObject,PIRP pIrp)
{
	NTSTATUS status = STATUS_SUCCESS;

	PIO_STACK_LOCATION pIrpSp = IoGetCurrentIrpStackLocation(pIrp);

	ULONG ulInputLen = pIrpSp->Parameters.DeviceIoControl.InputBufferLength;
	ULONG ulOutputLen = pIrpSp->Parameters.DeviceIoControl.OutputBufferLength;

	__try
	{
		switch(pIrpSp->Parameters.DeviceIoControl.IoControlCode)
		{
			// get object pointer from name
		case DEVICE_TREE_GET_POINTER_FROM_NAME:
			{
				PGET_OBJ_BY_NAME pParam = static_cast<PGET_OBJ_BY_NAME>(pIrp->AssociatedIrp.SystemBuffer);

				if(!pParam)
				{
					status = STATUS_INVALID_PARAMETER;
					break;
				}

				ULONG i = 0;
				while(i < ulInputLen - sizeof(GET_OBJ_BY_NAME))
				{
					if(!pParam->m_strName[i])
						break;

					i ++;
				}

				if(i == ulInputLen - sizeof(GET_OBJ_BY_NAME))
				{
					status = STATUS_INVALID_PARAMETER;
					break;
				}

				UNICODE_STRING strName;
				RtlInitUnicodeString(&strName,pParam->m_strName);

				HANDLE hObjectHandle;

				OBJECT_ATTRIBUTES objectAttribute;
				InitializeObjectAttributes(&objectAttribute,&strName,
					OBJ_CASE_INSENSITIVE,pParam->m_hRootHandle,0);

				status = ObOpenObjectByName(&objectAttribute,0,KernelMode,0,
					GENERIC_READ,0,&hObjectHandle);

				if(!NT_SUCCESS(status))
					break;

				PVOID pObject = 0;
				status = ObReferenceObjectByHandle(hObjectHandle,GENERIC_READ,0,KernelMode,
					&pObject,0);

				ZwClose(hObjectHandle);

				if(!NT_SUCCESS(status))
					break;

				RtlCopyMemory(pIrp->AssociatedIrp.SystemBuffer,&pObject,sizeof(PVOID));

				ObDereferenceObject(pObject);

				ulOutputLen = 4;
			}
			break;

			// query bus rel
		case DEVICE_TREE_QUERY_BUS_REL:
			{
				PVOID pRelBuffer = pIrp->AssociatedIrp.SystemBuffer;
				if( !pRelBuffer || 
					ulInputLen < sizeof(PVOID))
				{
					status = STATUS_INVALID_PARAMETER;
				}
				else
				{
					// first dword will be device object address
					PDEVICE_OBJECT pDeviceInStack = *static_cast<PDEVICE_OBJECT*>(pRelBuffer);
					if(pDeviceInStack->Type != IO_TYPE_DEVICE)
					{
						status = STATUS_INVALID_PARAMETER;
						break;
					}

					status = PnpQueryDeviceBusRel(pDeviceInStack,pRelBuffer,ulOutputLen);
				}
			}
			break;

			// query device id
		case DEVICE_TREE_QUERY_DEVICE_ID:
			{
				PVOID pIDBuffer = pIrp->AssociatedIrp.SystemBuffer;
				if( !pIDBuffer || 
					ulInputLen < sizeof(PVOID))
				{
					status = STATUS_INVALID_PARAMETER;
				}
				else
				{
					// first dword will be device object address
					PDEVICE_OBJECT pDeviceInStack = *static_cast<PDEVICE_OBJECT*>(pIDBuffer);

					status = PnpQueryDeviceID(pDeviceInStack,pIDBuffer,ulOutputLen);
				}
			}
			break;

			// query hardware id
		case DEVICE_TREE_QUERY_HARDWARE_ID:
			{
				PVOID pIDBuffer = pIrp->AssociatedIrp.SystemBuffer;
				if( !pIDBuffer || 
					ulInputLen < sizeof(PVOID))
				{
					status = STATUS_INVALID_PARAMETER;
				}
				else
				{
					// first dword will be device object address
					PDEVICE_OBJECT pDeviceInStack = *static_cast<PDEVICE_OBJECT*>(pIDBuffer);

					status = PnpQueryHardwareID(pDeviceInStack,pIDBuffer,ulOutputLen);
				}
			}
			break;

			// query compatible id
		case DEVICE_TREE_QUERY_COMPATIBLE_ID:
			{
				PVOID pIDBuffer = pIrp->AssociatedIrp.SystemBuffer;
				if( !pIDBuffer || 
					ulInputLen < sizeof(PVOID))
				{
					status = STATUS_INVALID_PARAMETER;
				}
				else
				{
					// first dword will be device object address
					PDEVICE_OBJECT pDeviceInStack = *static_cast<PDEVICE_OBJECT*>(pIDBuffer);

					status = PnpQueryCompatibleID(pDeviceInStack,pIDBuffer,ulOutputLen);
				}
			}
			break;

			// read memory
		case DEVICE_TREE_READ_KERNEL_MEMORY:
			{
				if(ulInputLen < sizeof(ULONG))
					status = STATUS_INVALID_PARAMETER;
				else
				{
					PULONG pAddr = *static_cast<PULONG*>(pIrp->AssociatedIrp.SystemBuffer);
					RtlCopyMemory(pIrp->AssociatedIrp.SystemBuffer,pAddr,ulOutputLen);
				}
			}
			break;

		case DEVICE_TREE_QUERY_OBJECT_NAME:
			{
				PVOID pBuffer = pIrp->AssociatedIrp.SystemBuffer;
				if( !pBuffer || 
					ulInputLen < sizeof(PVOID))
				{
					status = STATUS_INVALID_PARAMETER;
				}
				else
				{
					// first dword will be device object address
					PVOID pObject = *static_cast<PVOID*>(pBuffer);

					WCHAR temp[260];
					
					status = ObQueryNameString(pObject,temp,sizeof(temp),&ulOutputLen);

					if(!NT_SUCCESS(status))
						break;

					PUNICODE_STRING pString = reinterpret_cast<PUNICODE_STRING>(temp);

					RtlCopyMemory(pBuffer,pString->Buffer,pString->MaximumLength);
					ulOutputLen = pString->MaximumLength;
				}
			}
			break;

		default:
			status = STATUS_INVALID_PARAMETER;
			break;
		}
	}
	__except(EXCEPTION_EXECUTE_HANDLER)
	{
		status = STATUS_UNSUCCESSFUL;
	}
	
	pIrp->IoStatus.Information = ulOutputLen;
	pIrp->IoStatus.Status = status;
	IoCompleteRequest(pIrp,IO_NO_INCREMENT);
	return status;
}

// create
NTSTATUS DispatchCreate(PDEVICE_OBJECT pDeviceObject,PIRP pIrp)
{
	pIrp->IoStatus.Status = STATUS_SUCCESS;
	pIrp->IoStatus.Information = 0;

	IoCompleteRequest(pIrp,IO_NO_INCREMENT);

	return STATUS_SUCCESS;
}

// close
NTSTATUS DispatchClose(PDEVICE_OBJECT pDeviceObject,PIRP pIrp)
{
	pIrp->IoStatus.Status = STATUS_SUCCESS;
	pIrp->IoStatus.Information = 0;

	IoCompleteRequest(pIrp,IO_NO_INCREMENT);

	return STATUS_SUCCESS;
}