test.c

来自「灰狐驱动学习笔记系列 参考 windows驱动开发详解和 楚狂人windows驱」· C语言 代码 · 共 452 行

/*

  Test.c
  
  Author: Adly
  Last Updated: 2008-11-16
	
*/

#include <ntddk.h>

#include "Test.h"

//
// A structure representing the instance information associated with
// a particular device
//

typedef struct _DEVICE_EXTENSION
{
	ULONG  StateVariable;
	
} DEVICE_EXTENSION, *PDEVICE_EXTENSION;


//
// Function Declare
//

NTSTATUS
DriverEntry(
	IN PDRIVER_OBJECT		DriverObject,
	IN PUNICODE_STRING		RegistryPath
);

NTSTATUS
TestDispatchCreate(
	IN PDEVICE_OBJECT		DeviceObject,
	IN PIRP					Irp
);

NTSTATUS
TestDispatchClose(
	IN PDEVICE_OBJECT		DeviceObject,
	IN PIRP					Irp
);

NTSTATUS
TestDispatchDeviceControl(
	IN PDEVICE_OBJECT		DeviceObject,
	IN PIRP					Irp
);

VOID
TestUnload(
	IN PDRIVER_OBJECT		DriverObject
);

//////////////////////////////////////////////////////////////////////////
//
// 自定义函数声明
//
BOOLEAN
MyCopyFile(
           IN PUNICODE_STRING    ustrDestFile,
           IN PUNICODE_STRING    ustrSrcFile
            );

//////////////////////////////////////////////////////////////////////////

#ifdef ALLOC_PRAGMA
#pragma alloc_text(INIT, DriverEntry)
#pragma alloc_text(PAGE, TestDispatchCreate)
#pragma alloc_text(PAGE, TestDispatchClose)
#pragma alloc_text(PAGE, TestDispatchDeviceControl)
#pragma alloc_text(PAGE, TestUnload)
#endif // ALLOC_PRAGMA



NTSTATUS
DriverEntry(
	IN PDRIVER_OBJECT		DriverObject,
	IN PUNICODE_STRING		RegistryPath
)
{
	NTSTATUS			Status = STATUS_SUCCESS;    
	UNICODE_STRING		ntDeviceName;
	UNICODE_STRING		dosDeviceName;
	PDEVICE_EXTENSION	deviceExtension;
	PDEVICE_OBJECT		deviceObject = NULL;
	
	// 添加变量声明
	UNICODE_STRING		ustrSrcFile;
	UNICODE_STRING		ustrDestFile;

	
	KdPrint(("[Test] DriverEntry: %wZ\n", RegistryPath));
		
	RtlInitUnicodeString(&ntDeviceName, TEST_DEVICE_NAME_W);
	
	
	Status = IoCreateDevice(
		DriverObject,
		sizeof(DEVICE_EXTENSION),		// DeviceExtensionSize
		&ntDeviceName,					// DeviceName
		FILE_DEVICE_TEST,	// DeviceType
		0,								// DeviceCharacteristics
		TRUE,							// Exclusive
		&deviceObject					// [OUT]
		);
	
	if(!NT_SUCCESS(Status))
	{
		KdPrint(("[Test] IoCreateDevice Error Code = 0x%X\n", Status));
		
		return Status;
	}
	
	deviceExtension = (PDEVICE_EXTENSION)deviceObject->DeviceExtension;
	
	//
	// Set up synchronization objects, state info,, etc.
	//
	
	//
	// Create a symbolic link that Win32 apps can specify to gain access
	// to this driver/device
	//
	
	RtlInitUnicodeString(&dosDeviceName, TEST_DOS_DEVICE_NAME_W);
	
	Status = IoCreateSymbolicLink(&dosDeviceName, &ntDeviceName);
	
	if(!NT_SUCCESS(Status))
	{
		KdPrint(("[Test] IoCreateSymbolicLink Error Code = 0x%X\n", Status));
		
		//
		// Delete Object
		//

		IoDeleteDevice(deviceObject);
		
		return Status;
	}
	
	//
	// Create dispatch points for device control, create, close.
	//
	
	DriverObject->MajorFunction[IRP_MJ_CREATE]			= TestDispatchCreate;
	DriverObject->MajorFunction[IRP_MJ_CLOSE]			= TestDispatchClose;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL]	= TestDispatchDeviceControl;
	DriverObject->DriverUnload							= TestUnload;

	//////////////////////////////////////////////////////////////////////////
	//
	// 自定义函数测试
	//
	RtlInitUnicodeString(&ustrSrcFile, L"\\??\\C:\\windows\\notepad.exe");
    RtlInitUnicodeString(&ustrDestFile, L"\\??\\C:\\notepad.exe");
	
    if(MyCopyFile(&ustrDestFile, &ustrSrcFile))
    {
		KdPrint(("[Test] CopyFile Success!"));
    }
    else
    {  
		KdPrint(("[Test] CopyFile Error!"));
    }
	

	return Status;
}

NTSTATUS
TestDispatchCreate(
	IN PDEVICE_OBJECT		DeviceObject,
	IN PIRP					Irp
)
{
	NTSTATUS Status = STATUS_SUCCESS;
	
	Irp->IoStatus.Information = 0;
	
	KdPrint(("[Test] IRP_MJ_CREATE\n"));
	
	Irp->IoStatus.Status = Status;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	
	return Status;
}

NTSTATUS
TestDispatchClose(
	IN PDEVICE_OBJECT		DeviceObject,
	IN PIRP					Irp
)
{
	NTSTATUS Status = STATUS_SUCCESS;
	
	Irp->IoStatus.Information = 0;
	
	KdPrint(("[Test] IRP_MJ_CLOSE\n"));
	
	Irp->IoStatus.Status = Status;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	
	return Status;
}

NTSTATUS
TestDispatchDeviceControl(
	IN PDEVICE_OBJECT		DeviceObject,
	IN PIRP					Irp
)
{
	NTSTATUS			Status = STATUS_SUCCESS;
	PIO_STACK_LOCATION	irpStack;
	PDEVICE_EXTENSION	deviceExtension;
	PVOID				ioBuf;
	ULONG				inBufLength, outBufLength;
	ULONG				ioControlCode;
	
	irpStack = IoGetCurrentIrpStackLocation(Irp);
	deviceExtension = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;
	
	Irp->IoStatus.Information = 0;
	
	//
	// Get the pointer to the input/output buffer and it's length
	//
	
	ioBuf = Irp->AssociatedIrp.SystemBuffer;
	inBufLength = irpStack->Parameters.DeviceIoControl.InputBufferLength;
	outBufLength = irpStack->Parameters.DeviceIoControl.OutputBufferLength;
	ioControlCode = irpStack->Parameters.DeviceIoControl.IoControlCode;
	// Irp->UserBuffer;		// If METHOD_NEITHER, This is Output Buffer
	
	switch (ioControlCode)
	{
	case IOCTL_TEST_TEST0:
		{
			//
			// Sample IO Control
			//
			KdPrint(("[Test] TEST0 IOCTL: 0x%X", ioControlCode));
			
			break;
		}

	case IOCTL_TEST_TEST1:
		{
			//
			// Sample IO Control
			//
			KdPrint(("[Test] TEST1 IOCTL: 0x%X", ioControlCode));
			
			break;
		}
		
	default:
		{
			Status = STATUS_INVALID_PARAMETER;
			
			KdPrint(("[Test] Unknown IOCTL: 0x%X (%04X,%04X)",
				ioControlCode, DEVICE_TYPE_FROM_CTL_CODE(ioControlCode),
				IoGetFunctionCodeFromCtlCode(ioControlCode)));
			
			break;
		}
	}
	
	
	Irp->IoStatus.Status = Status;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	
	
	return Status;
}

VOID
TestUnload(
	IN PDRIVER_OBJECT		DriverObject
)
{
	UNICODE_STRING dosDeviceName;
	
	//
	// Free any resources
	//
	
	//
	// Delete the symbolic link
	//
	
	RtlInitUnicodeString(&dosDeviceName, TEST_DOS_DEVICE_NAME_W);
	
	IoDeleteSymbolicLink(&dosDeviceName);
	
	//
	// Delete the device object
	//
	
	IoDeleteDevice(DriverObject->DeviceObject);
	
	KdPrint(("[Test] Unloaded"));
}


//////////////////////////////////////////////////////////////////////////
//
// 自定义函数实现
//

BOOLEAN
MyCopyFile(
           IN PUNICODE_STRING    ustrDestFile,
           IN PUNICODE_STRING    ustrSrcFile
            )
{
    HANDLE	hSrcFile, hDestFile;
    PVOID	buffer = NULL;
    ULONG	length = 0;
    LARGE_INTEGER	offset = {0};
    IO_STATUS_BLOCK	Io_Status_Block = {0};
    OBJECT_ATTRIBUTES	obj_attrib;
    NTSTATUS	status;
    BOOLEAN		bRet = FALSE;

    do 
    {
        // 打开源文件
        InitializeObjectAttributes(&obj_attrib,
								   ustrSrcFile,
								   OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE,
								   NULL,
								   NULL);
        status = ZwCreateFile(&hSrcFile,
							  GENERIC_READ,
							  &obj_attrib,
							  &Io_Status_Block,
							  NULL,
							  FILE_ATTRIBUTE_NORMAL,
							  FILE_SHARE_READ,
							  FILE_OPEN,
							  FILE_NON_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT,
							  NULL,
							  0);
        if (!NT_SUCCESS(status))
        {
            bRet = FALSE;
            goto END;
        }

        // 打开目标文件
        InitializeObjectAttributes(&obj_attrib, 
								   ustrDestFile,
								   OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE,
								   NULL,
								   NULL);
        status = ZwCreateFile(&hDestFile,
							  GENERIC_WRITE,
							  &obj_attrib,
							  &Io_Status_Block,
							  NULL,
							  FILE_ATTRIBUTE_NORMAL,
							  FILE_SHARE_READ,
							  FILE_OPEN_IF,
							  FILE_NON_DIRECTORY_FILE | FILE_SYNCHRONOUS_IO_NONALERT,
							  NULL,
							  0);
        if (!NT_SUCCESS(status))
        {
            bRet = FALSE;
            goto END;
        }

        // 为buffer分配4KB空间
        buffer = ExAllocatePool(NonPagedPool, 1024 * 4);
        if (buffer == NULL)
        {
            bRet = FALSE;
            goto END;
        }

        // 复制文件
        while (1)
        {
            length = 4 * 1024;
            // 读取源文件
            status = ZwReadFile(hSrcFile,
                            NULL,
                            NULL,
                            NULL,
                            &Io_Status_Block,
                            buffer,
                            length,
                            &offset,
                            NULL);
            if (!NT_SUCCESS(status))
            {
                // 如果状态为STATUS_END_OF_FILE，说明文件已经读取到末尾
                if (status == STATUS_END_OF_FILE)
                {
                    bRet = TRUE;
                    goto END;
                }
            }

            // 获得实际读取的长度
            length = Io_Status_Block.Information;

            // 写入到目标文件
            status = ZwWriteFile(hDestFile,
                                NULL,
                                NULL,
                                NULL,
                                &Io_Status_Block,
                                buffer,
                                length,
                                &offset,
                                NULL);
            if (!NT_SUCCESS(status))
            {
                bRet = FALSE;
                goto END;
            }
            
            // 移动文件指针
            offset.QuadPart += length;
        }

    } while (0);

END:
    if (hSrcFile)
    {
        ZwClose(hSrcFile);
    }
    if (hDestFile)
    {
        ZwClose(hDestFile);
    }
    if (buffer = NULL)
    {
        ExFreePool(buffer);
    }

    return bRet;
}