Windows驱动之IO_REMOVE_LOCK

Windows驱动之IO_REMOVE_LOCK

在设备驱动中，我们很多情况下会遇到这样的问题，各种PNP请求不是同步请求过来的；例如当我们在Read、Write设备的时候，可能设备这个时候接受到了IRP_MN_REMOVE_DEVICE移除设备的请求，如果这个时候删除资源，那么其他IRP就可能出错了。为了避免这种问题的出现，Windows引入了IO_REMOVE_LOCK移除锁。

IO_REMOVE_LOCK的操作函数如下：

VOID IoInitializeRemoveLock(
  _In_ PIO_REMOVE_LOCK Lock,
  _In_ ULONG           AllocateTag,
  _In_ ULONG           MaxLockedMinutes,
  _In_ ULONG           HighWatermark
);

NTSTATUS IoAcquireRemoveLock(
  _In_     PIO_REMOVE_LOCK RemoveLock,
  _In_opt_ PVOID           Tag
);

VOID IoReleaseRemoveLock(
  _In_ PIO_REMOVE_LOCK RemoveLock,
  _In_ PVOID           Tag
);

VOID IoReleaseRemoveLockAndWait(
  _In_ PIO_REMOVE_LOCK RemoveLock,
  _In_ PVOID           Tag
);

下面我们来看下这些函数的操作过程.

1. IO_REMOVE_LOCK

typedef struct _IO_REMOVE_LOCK_COMMON_BLOCK {
    BOOLEAN     Removed;
    BOOLEAN     Reserved [3];
    LONG        IoCount;
    KEVENT      RemoveEvent;

} IO_REMOVE_LOCK_COMMON_BLOCK;

typedef struct _IO_REMOVE_LOCK {
    IO_REMOVE_LOCK_COMMON_BLOCK Common;
#if DBG
    IO_REMOVE_LOCK_DBG_BLOCK Dbg;
#endif
} IO_REMOVE_LOCK, *PIO_REMOVE_LOCK;

这个成员主要有如下的作用：

1. Removed ： 锁对象是否被移除。
2. IoCount ： 被请求的次数。
3. RemoveEvent ： 用来通知所有对象都被移除。

2. IoInitializeRemoveLock

在使用之前，必须先初始化IO_REMOVE_LOCK，如下：

NTSTATUS AddDevice(PDRIVER_OBJECT DriverObject, PDEVICE_OBJECT pdo)
{
    //...
    IoInitializeRemoveLock(&pdx->RemoveLock, 0, 0, 0);
    //...
}

其中IoInitializeRemoveLock的实现如下：

VOID
NTAPI
IoInitializeRemoveLockEx(IN PIO_REMOVE_LOCK RemoveLock,
                         IN ULONG AllocateTag,
                         IN ULONG MaxLockedMinutes,
                         IN ULONG HighWatermark,
                         IN ULONG RemlockSize)
{
    PEXTENDED_IO_REMOVE_LOCK Lock = (PEXTENDED_IO_REMOVE_LOCK)RemoveLock;
    {
        return;
    }

    switch (RemlockSize)
    {
        case (sizeof(IO_REMOVE_LOCK_DBG_BLOCK) + sizeof(IO_REMOVE_LOCK_COMMON_BLOCK)):
            //调试信息

        case sizeof(IO_REMOVE_LOCK_COMMON_BLOCK):
            Lock->Common.Removed = FALSE;
            Lock->Common.IoCount = 1;
            KeInitializeEvent(&Lock->Common.RemoveEvent,
                              SynchronizationEvent,
                              FALSE);
    }
}

这个函数的主要作用是初始化PIO_REMOVE_LOCK,如下：

1. Lock->Common.Removed = FALSE;初始化的时候，移除对象没有移除，当前正常。
2. Lock->Common.IoCount = 1; : 拥有的IO次数为1.
3. KeInitializeEvent(&Lock->Common.RemoveEvent : 初始化移除通知事件。

3. IoAcquireRemoveLock

当我们在操作IRP的时候，可以先请求移除锁，例如如下：

NTSTATUS DispatchSomething(PDEVICE_OBJECT fdo, PIRP Irp)
{
    PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
    NTSTATUS status = IoAcquireRemoveLock(&pdx->RemoveLock, Irp);
    if (!NT_SUCCESS(status))
        return CompleteRequest(Irp, status, 0);
    //...
    IoReleaseRemoveLock(&pdx->RemoveLock, Irp);
    return CompleteRequest(Irp, <some code>, <info value>);
}

在请求IPR的时候，先调用IoAcquireRemoveLock占用删除锁，防止设备被删除，IRP完成之后，再调用IoReleaseRemoveLock释放删除锁。

需要注意的是IoAcquireRemoveLock当锁被删除之后，返回STATUS_DELETE_PENDING, 函数实现如下：

NTSTATUS
NTAPI
IoAcquireRemoveLockEx(IN PIO_REMOVE_LOCK RemoveLock,
                      IN OPTIONAL PVOID Tag,
                      IN LPCSTR File,
                      IN ULONG Line,
                      IN ULONG RemlockSize)
{
    KIRQL OldIrql;
    LONG LockValue;
    PIO_REMOVE_LOCK_TRACKING_BLOCK TrackingBlock;
    PEXTENDED_IO_REMOVE_LOCK Lock = (PEXTENDED_IO_REMOVE_LOCK)RemoveLock;

    LockValue = InterlockedIncrement(&(Lock->Common.IoCount));
    ASSERT(LockValue > 0);
    if (!Lock->Common.Removed)
    {
        //...
    }
    else
    {
        if (!InterlockedDecrement(&(Lock->Common.IoCount)))
        {
            KeSetEvent(&(Lock->Common.RemoveEvent), IO_NO_INCREMENT, FALSE);
        }
        return STATUS_DELETE_PENDING;
    }
    return STATUS_SUCCESS;
}

1. InterlockedIncrement(&(Lock->Common.IoCount)); 递增锁的计数。
2. 如果Lock->Common.Removed已经被移除了，那么InterlockedDecrement(&(Lock->Common.IoCount))并释放引用计数，然后返回STATUS_DELETE_PENDING。
3. 如果锁没有被释放（Lock->Common.Removed为FALSE），那么直接请求占用锁。

4. IoReleaseRemoveLockEx

占用删除锁之后，使用IoReleaseRemoveLockEx来释放删除锁，这个代码如下：

VOID
NTAPI
IoReleaseRemoveLockEx(IN PIO_REMOVE_LOCK RemoveLock,
                      IN PVOID Tag,
                      IN ULONG RemlockSize)
{
    LONG LockValue;

    //...

    LockValue = InterlockedDecrement(&(Lock->Common.IoCount));

    if (!LockValue)
    {
        KeSetEvent(&(Lock->Common.RemoveEvent), IO_NO_INCREMENT, FALSE);
    }
}

这里一个比较简的逻辑：

1. InterlockedDecrement(&(Lock->Common.IoCount))递减IO引用计数。
2. 如果计数递减到0，那么说明有人在等待这个锁被删除，设置事件。KeSetEvent(&(Lock->Common.RemoveEvent), IO_NO_INCREMENT, FALSE);

其中等待PIO_REMOVE_LOCK的流程为IoReleaseRemoveLockAndWait.

5. IoReleaseRemoveLockAndWait

如果当设备要删除的时候，我们调用：

1. IoAcquireRemoveLock : 占用移除锁。
2. IoReleaseRemoveLockAndWait : 释放锁，并等待所有的删除锁占用被释放。

NTSTATUS DispatchPnp(PDEVICE_OBJECT fdo, PIRP Irp)
{
    PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
    NTSTATUS status = IoAcquireRemoveLock(&pdx->RemoveLock, Irp);
    if (!NT_SUCCESS(status))
        return CompleteRequest(Irp, status, 0);
    //...
    status = (*fcntab[fcn](fdo, Irp);
    if (fcn != IRP_MN_REMOVE_DEVICE)
        IoReleaseRemoveLock(&pdx->RemoveLock, Irp);
    return 
}

NTSTATUS HandleRemoveDevice(PDEVICE_OBJECT fdo, PIRP Irp)
{
    Irp->IoStatus.Status = STATUS_SUCCESS;
    PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
    AbortRequests(&pdx->dqReadWrite, STATUS_DELETE_PENDING);				
    DeregisterAllInterfaces(pdx);
    StopDevice(fdo, pdx->state == WORKING);
    pdx->state = REMOVED;
    NTSTATUS status = DefaultPnpHandler(pdx->LowerDeviceObject, Irp);			
    IoReleaseRemoveLockAndWait(&pdx->RemoveLock, Irp);				
    RemoveDevice(fdo);
    return status;
}

其中IoReleaseRemoveLockAndWait这个函数的代码如下：

VOID
NTAPI
IoReleaseRemoveLockAndWaitEx(IN PIO_REMOVE_LOCK RemoveLock,
                             IN PVOID Tag,
                             IN ULONG RemlockSize)
{
    LONG LockValue;
    PIO_REMOVE_LOCK_TRACKING_BLOCK TrackingBlock;
    PEXTENDED_IO_REMOVE_LOCK Lock = (PEXTENDED_IO_REMOVE_LOCK)RemoveLock;
    PAGED_CODE();

    Lock->Common.Removed = TRUE;
    LockValue = InterlockedDecrement(&Lock->Common.IoCount);
    if (InterlockedDecrement(&Lock->Common.IoCount) > 0)
    {
        KeWaitForSingleObject(&(Lock->Common.RemoveEvent),
                              Executive,
                              KernelMode,
                              FALSE,
                              NULL);
    }

    //...
}

从上面流程我们知道：

1. Lock->Common.Removed = TRUE; : 当IoReleaseRemoveLockAndWaitEx调用的时候，删除锁被释放，也就是说其他地方已经不能在处理请求了，对象被删除了。
2. LockValue = InterlockedDecrement(&Lock->Common.IoCount); : 释放自己IoAcquireRemoveLock占用的引用。
3. if (InterlockedDecrement(&Lock->Common.IoCount) > 0) : 这条语句比较重要，因为IoReleaseRemoveLockAndWaitEx这个函数除了释放IoAcquireRemoveLock占用的引用计数之外，还需要释放整个状态（也就是说等待所有其他占用锁的完成）, 在IoInitializeRemoveLock的时候IoCount初始化为1，所以这里要调用这个来判断是不是还有其他线程在调用IoAcquireRemoveLock占用锁，如果是Lock->Common.IoCount大于0，那么有其他线程请求了改锁，需要等待完成。

6. 总结

其实整理来说IO_REMOVE_LOCK使用比较简单，总结为如下：

1. IoInitializeRemoveLock ： 初始化IO_REMOVE_LOCK。
2. 如果要进入禁止删除操作的代码段调用IoAcquireRemoveLock, 代码调用完成之后，调用IoReleaseRemoveLock.
3. 如果需要开始删除关键数据，需要先IoReleaseRemoveLock获取锁，然后IoReleaseRemoveLockAndWaitEx释放锁，并等待其他线程代码段执行完成。