NT驱动使用设备扩展的例子

前言

设备扩展，可以消掉操作每个设备时需要的全局变量.

试验记录

模拟向文件中写入2次字符串，第二次写时调整文件指针的位置.然后调整文件指针，从文件头开始读字符串.
设备IO采用缓冲区方式.
设备扩展是DEVICE_OBJECT.DeviceExtension, 不是DEVICE_OBJECT.DeviceObjectExtension, 这个有可能会看错.

R3

// hw.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include <windows.h>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <WinIoCtl.h>

#define SYM_LINK_NAME_R3 L"\\\\?\\TestVf"
#define LEN_BUFFER_ON_DEVICE_EXT 4096

#define CTL_BASE 0x1000
#define MY_CTL_CODE(x) CTL_CODE(FILE_DEVICE_UNKNOWN, CTL_BASE + (x), METHOD_BUFFERED, FILE_ANY_ACCESS)

#define MY_IOCTL_GET_FILE_SIZE  MY_CTL_CODE(1)
#define MY_IOCTL_SET_FILE_POS   MY_CTL_CODE(2)

int main(int argc, char* argv[])
{
    HANDLE hFile = INVALID_HANDLE_VALUE;
    char szBuf[LEN_BUFFER_ON_DEVICE_EXT] = {'\0'};
    size_t nLenWrite = 0;
    DWORD dwCnt = 0;
    ULONG ulFileSize = 0;
    ULONG ulFilePos = 0;
    DWORD dwRc = 0;
    printf("CreateFile...\r\n");
    hFile = CreateFileW(
                SYM_LINK_NAME_R3,
                GENERIC_READ | GENERIC_WRITE,
                0,
                NULL,
                OPEN_EXISTING,
                FILE_ATTRIBUTE_NORMAL,
                NULL);

    if (INVALID_HANDLE_VALUE != hFile) {
        if (::DeviceIoControl(hFile, MY_IOCTL_GET_FILE_SIZE,
                              &ulFileSize, sizeof(ulFileSize),
                              &ulFileSize, sizeof(ulFileSize), &dwRc, NULL)) {
            printf("ulFileSize = 0x%X\r\n", ulFileSize);
        }

        strcpy(szBuf, "hello_");

        if (WriteFile(hFile, szBuf, strlen(szBuf), &dwCnt, NULL)) {
            nLenWrite += dwCnt;
        }

        ZeroMemory(szBuf, sizeof(szBuf));
        ReadFile(hFile, szBuf, nLenWrite, &dwCnt, NULL);
        printf("dwCnt = %d, ReadFrom Drv is : %s\r\n", dwCnt, szBuf);
        strcpy(szBuf, "world");

        ulFilePos = 6;
        if (::DeviceIoControl(hFile, MY_IOCTL_SET_FILE_POS,
            &ulFilePos, sizeof(ulFilePos),
            &ulFilePos, sizeof(ulFilePos), &dwRc, NULL)) {
            printf("ulFilePos = 0x%X\r\n", ulFilePos);
        }

        if (WriteFile(hFile, szBuf, strlen(szBuf), &dwCnt, NULL)) {
            nLenWrite += dwCnt;
        }

        ulFilePos = 0;
        if (::DeviceIoControl(hFile, MY_IOCTL_SET_FILE_POS,
            &ulFilePos, sizeof(ulFilePos),
            &ulFilePos, sizeof(ulFilePos), &dwRc, NULL)) {
            printf("ulFilePos = 0x%X\r\n", ulFilePos);
        }

        ZeroMemory(szBuf, sizeof(szBuf));
        ReadFile(hFile, szBuf, nLenWrite, &dwCnt, NULL);
        printf("dwCnt = %d, ReadFrom Drv is : %s\r\n", dwCnt, szBuf);
        CloseHandle(hFile);
        hFile = NULL;
    } else {
        printf("failed : CreateFile errSn = %d\r\n", GetLastError());
    }

    system("pause");
    return 0;
}

R0

// @file hw.c
// @brief 设备扩展的例子

#include <ntddk.h>

#define DEVICE_NAME L"\\device\\TestVf"
#define SYM_LINK_NAME_R0 L"\\??\\TestVf"
#define SYM_LINK_NAME_R3 L"\\\\?\\TestVf"
#define POOL_TAG_DEVICE_EXT 'ext1'
#define LEN_BUFFER_ON_DEVICE_EXT 4096

#define CTL_BASE 0x1000
#define MY_CTL_CODE(x) CTL_CODE(FILE_DEVICE_UNKNOWN, CTL_BASE + (x), METHOD_BUFFERED, FILE_ANY_ACCESS)

#define MY_IOCTL_GET_FILE_SIZE  MY_CTL_CODE(1)
#define MY_IOCTL_SET_FILE_POS   MY_CTL_CODE(2)

typedef struct _DEVICE_EXTENSION {
    PDEVICE_OBJECT pDevice;
    UNICODE_STRING strDeviceName;
    UNICODE_STRING strSymLinkName;

    PUCHAR pBuffer;
    ULONG uLenBuffer;
    ULONG ulFilePos;
    KEVENT EventBuffer;
} DEVICE_EXTENSION, *PDEVICE_EXTENSION;

NTSTATUS DriverEntry(__in struct _DRIVER_OBJECT*  DriverObject, __in PUNICODE_STRING  RegistryPath);
#pragma alloc_text( "INIT", DriverEntry)

VOID OnUnload(__in struct _DRIVER_OBJECT* DriverObject);
#pragma alloc_text( "PAGE", OnUnload)

NTSTATUS DispatchEmpty(__in struct _DEVICE_OBJECT*  DeviceObject,
                       __in struct _IRP*  Irp);
#pragma alloc_text( "PAGE", DispatchEmpty)

NTSTATUS OnCreate(__in struct _DEVICE_OBJECT*  DeviceObject,
                  __in struct _IRP*  Irp);
#pragma alloc_text( "PAGE", OnCreate)

NTSTATUS OnRead(__in struct _DEVICE_OBJECT*  DeviceObject,
                __in struct _IRP*  Irp);
#pragma alloc_text( "PAGE", OnRead)

NTSTATUS OnWrite(__in struct _DEVICE_OBJECT*  DeviceObject,
                 __in struct _IRP*  Irp);
#pragma alloc_text( "PAGE", OnWrite)

NTSTATUS OnControl(__in struct _DEVICE_OBJECT*  DeviceObject,
                   __in struct _IRP*  Irp);
#pragma alloc_text( "PAGE", OnControl)

NTSTATUS OnClose(__in struct _DEVICE_OBJECT*  DeviceObject,
                 __in struct _IRP*  Irp);
#pragma alloc_text( "PAGE", OnClose)

NTSTATUS fnCreateDevice(IN PDRIVER_OBJECT pDriverObject);
#pragma alloc_text( "PAGE", fnCreateDevice)

NTSTATUS DriverEntry(__in struct _DRIVER_OBJECT*  DriverObject, __in PUNICODE_STRING  RegistryPath)
{
    DriverObject->DriverUnload = OnUnload;
    DriverObject->MajorFunction[IRP_MJ_CREATE] = OnCreate;
    DriverObject->MajorFunction[IRP_MJ_CLOSE] = OnClose;
    DriverObject->MajorFunction[IRP_MJ_READ] = OnRead;
    DriverObject->MajorFunction[IRP_MJ_WRITE] = OnWrite;
    DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = OnControl;
    DriverObject->MajorFunction[IRP_MJ_CREATE_NAMED_PIPE] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_QUERY_INFORMATION] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_SET_INFORMATION] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_QUERY_EA] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_SET_EA] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_FLUSH_BUFFERS] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_QUERY_VOLUME_INFORMATION] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_SET_VOLUME_INFORMATION] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_DIRECTORY_CONTROL] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_FILE_SYSTEM_CONTROL] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_INTERNAL_DEVICE_CONTROL] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_SHUTDOWN] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_LOCK_CONTROL] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_CLEANUP] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_CREATE_MAILSLOT] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_QUERY_SECURITY] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_SET_SECURITY] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_POWER] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_SYSTEM_CONTROL] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_DEVICE_CHANGE] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_QUERY_QUOTA] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_SET_QUOTA] = DispatchEmpty;
    DriverObject->MajorFunction[IRP_MJ_PNP] = DispatchEmpty;
    return fnCreateDevice(DriverObject);
}

NTSTATUS fnCreateDevice(IN PDRIVER_OBJECT pDriverObject)
{
    NTSTATUS status;
    PDEVICE_OBJECT pDevObj;
    PDEVICE_EXTENSION pDevExt;
    UNICODE_STRING strSymLinkName;
    UNICODE_STRING strDeviceName;

    // DbgBreakPoint();
    RtlInitUnicodeString(&strDeviceName, DEVICE_NAME);
    status = IoCreateDevice(pDriverObject,
                            sizeof(DEVICE_EXTENSION),
                            &strDeviceName,
                            FILE_DEVICE_UNKNOWN,
                            0, TRUE,
                            &pDevObj);

    if (!NT_SUCCESS(status)) {
        KdPrint(("IoCreateDevice failed"));
        return status;
    }

    pDevObj->Flags = DO_BUFFERED_IO;
    pDevExt = (PDEVICE_EXTENSION)pDevObj->DeviceExtension;
    pDevExt->pDevice = pDevObj;
    RtlInitUnicodeString(&pDevExt->strDeviceName, DEVICE_NAME);

    pDevExt->pBuffer = NULL;
    pDevExt->uLenBuffer = 0;
    pDevExt->ulFilePos = 0;

    RtlInitUnicodeString(&strSymLinkName, SYM_LINK_NAME_R0);
    RtlInitUnicodeString(&pDevExt->strSymLinkName, SYM_LINK_NAME_R0);

    status = IoCreateSymbolicLink(&strSymLinkName, &strDeviceName);

    if (!NT_SUCCESS(status)) {
        KdPrint(("IoCreateSymbolicLink failed"));
        IoDeleteDevice(pDevObj);
        return status;
    }

    KdPrint(("fnCreateDevice ok."));
    return STATUS_SUCCESS;
}

VOID OnUnload(__in struct _DRIVER_OBJECT* DriverObject)
{
    PDEVICE_EXTENSION pDevExt;
    PDEVICE_OBJECT  pNextObj = DriverObject->DeviceObject;
    KdPrint((">> fnDrvUnLoad"));
    // DbgBreakPoint();

    while (pNextObj != NULL) {
        pDevExt = (PDEVICE_EXTENSION)pNextObj->DeviceExtension;

        if (NULL != pDevExt) {
            KdPrint(("IoDeleteSymbolicLink %wZ", pDevExt->strSymLinkName));
            IoDeleteSymbolicLink(&pDevExt->strSymLinkName);
            IoDeleteDevice(pDevExt->pDevice);
        }

        pNextObj = pNextObj->NextDevice;
    }
}

NTSTATUS DispatchEmpty(__in struct _DEVICE_OBJECT*  DeviceObject,
                       __in struct _IRP*  Irp)
{
    return STATUS_SUCCESS;
}

NTSTATUS OnCreate(__in struct _DEVICE_OBJECT*  DeviceObject,
                  __in struct _IRP*  Irp)
{
    PDEVICE_EXTENSION pDevExt = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;
    KdPrint((">> OnCreate\n"));
    // DbgBreakPoint();

    pDevExt->pBuffer = (PUCHAR)ExAllocatePoolWithTag(PagedPool, LEN_BUFFER_ON_DEVICE_EXT, POOL_TAG_DEVICE_EXT);
    pDevExt->uLenBuffer = LEN_BUFFER_ON_DEVICE_EXT;
    KeInitializeEvent(&pDevExt->EventBuffer, SynchronizationEvent, TRUE);
    return STATUS_SUCCESS;
}

NTSTATUS OnRead(__in struct _DEVICE_OBJECT*  DeviceObject,
                __in struct _IRP*  Irp)
{
    PDEVICE_EXTENSION pDevExt = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;
    NTSTATUS status = STATUS_SUCCESS;
    PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);
    ULONG ulLen = stack->Parameters.Read.Length;
    ULONG ulOffset = (ULONG)stack->Parameters.Read.ByteOffset.QuadPart;

    // DbgBreakPoint();
    if ((ulOffset + ulLen + pDevExt->ulFilePos) > pDevExt->uLenBuffer) {
        status = STATUS_FILE_INVALID;
        ulLen = 0;
    } else {
        KeWaitForSingleObject(&pDevExt->EventBuffer, Executive, KernelMode, FALSE , NULL);
        memcpy(Irp->AssociatedIrp.SystemBuffer, pDevExt->pBuffer + ulOffset + pDevExt->ulFilePos, ulLen);
        KeSetEvent(&pDevExt->EventBuffer, IO_NO_INCREMENT, FALSE);
        status = STATUS_SUCCESS;
    }

    Irp->IoStatus.Status = status;
    Irp->IoStatus.Information = ulLen;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    return status;
}

NTSTATUS OnWrite(__in struct _DEVICE_OBJECT*  DeviceObject,
                 __in struct _IRP*  Irp)
{
    NTSTATUS Status = STATUS_SUCCESS;
    PDEVICE_EXTENSION pDevExt = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;
    PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);
    ULONG ulLen = stack->Parameters.Write.Length;
    ULONG ulOffset = (ULONG)stack->Parameters.Write.ByteOffset.QuadPart;

    // DbgBreakPoint();
    if ((ulLen + ulOffset + pDevExt->ulFilePos) > pDevExt->uLenBuffer) {
        Status = STATUS_FILE_INVALID;
        ulLen = 0;
    } else {
        KeWaitForSingleObject(&pDevExt->EventBuffer, Executive, KernelMode, FALSE , NULL);
        memcpy(pDevExt->pBuffer + ulOffset + pDevExt->ulFilePos, Irp->AssociatedIrp.SystemBuffer, ulLen);
        KeSetEvent(&pDevExt->EventBuffer, IO_NO_INCREMENT, FALSE);
        Status = STATUS_SUCCESS;
    }

    Irp->IoStatus.Status = Status;
    Irp->IoStatus.Information = ulLen;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    return Status;
}

NTSTATUS OnControl(__in struct _DEVICE_OBJECT*  DeviceObject,
                   __in struct _IRP*  Irp)
{
    PDEVICE_EXTENSION pDevExt = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;
    PIO_STACK_LOCATION pIrpStack = IoGetCurrentIrpStackLocation(Irp);
    ULONG nCtlCode = pIrpStack->Parameters.DeviceIoControl.IoControlCode;
    PVOID pBuf = Irp->AssociatedIrp.SystemBuffer;
    ULONG nInLength = pIrpStack->Parameters.DeviceIoControl.InputBufferLength;
    ULONG nOutLength = pIrpStack->Parameters.DeviceIoControl.OutputBufferLength;
    NTSTATUS status = STATUS_UNSUCCESSFUL;
    ULONG nLength = 0;

    KdPrint(("DispatchControl\r\n"));
    // DbgBreakPoint();

    switch(nCtlCode)
    {
    case MY_IOCTL_GET_FILE_SIZE:
        KdPrint(("MY_IOCTL_GET_FILE_SIZE\r\n"));
        RtlCopyMemory(pBuf, &pDevExt->uLenBuffer, sizeof(ULONG));
        status = STATUS_SUCCESS;
        nLength = sizeof(ULONG);
        break;
    case MY_IOCTL_SET_FILE_POS:
        KdPrint(("MY_IOCTL_SET_FILE_POS\r\n"));
        RtlCopyMemory(&pDevExt->ulFilePos, pBuf, sizeof(ULONG));
        status = STATUS_SUCCESS;
        nLength = sizeof(ULONG);
        break;
    default:
        break;
    }

    //完成请求
    Irp->IoStatus.Status = status;
    Irp->IoStatus.Information = nLength; 
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    return STATUS_SUCCESS;
}

NTSTATUS OnClose(__in struct _DEVICE_OBJECT*  DeviceObject,
                 __in struct _IRP*  Irp)
{
    PDEVICE_EXTENSION pDevExt = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;
    // DbgBreakPoint();
    ExFreePoolWithTag(pDevExt->pBuffer, POOL_TAG_DEVICE_EXT);
    return STATUS_SUCCESS;
}