首页 > 未分类 > TDI 过滤驱动监视进程的网络流量

TDI 过滤驱动监视进程的网络流量

要控制流量大概是在 IoCompletion 里面返回 STATUS_MORE_PROCESSING_REQUIRED 然后投递一个 WorkItem 到某个工作线程里面去,然后在那工作线程里面多 Sleep 一会儿再 IoCompleteRequest?

[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id:  280, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id:  280, Sent:        413 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:         57 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id:    4, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:         57 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id:    4, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id:  280, Sent:        413 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:         57 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id:    4, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id:  280, Sent:        596 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:        239 bytes / s Received:      70948 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id:    4, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:        239 bytes / s Received:    8927920 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id:    4, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   12346060 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:    8163876 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:    9287056 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:    3076616 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:    7026796 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:    9108028 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:    3158352 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:    5285124 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   10047148 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   11329012 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   13408224 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   13795512 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   16626816 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id:    4, Sent:        201 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   13855816 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   20182288 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   14285852 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   14093656 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   17039600 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   18342860 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   16822336 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   14925796 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   17185912 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   17422724 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   14773376 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   13452628 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:   20832756 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id:  280, Sent:         17 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:    2674734 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1132, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:    2674734 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id: 1588, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id:    4, Sent:          0 bytes / s Received:          0 bytes / s.
[FlowWatcher] DumpNetworkUsage: Process Id:    4, Sent:          0 bytes / s Received:          0 bytes / s.

#include <ntddk.h>
#include <tdikrnl.h>

//
// 驱动名称。
//

#define PROJECT_NAME "FlowWatcher"

//
// 输出网络使用情况的间隔,单位毫秒。
//
// 由于 MILLISECOND_TO_SECOND 在将毫秒转换为秒时可能会有精度损失
// 因此此处最好是整数秒。
//

#define NETWORK_USAGE_DUMP_INTERVAL 1000

//
// 毫秒转换为内核下的等待时长。
//
// 内核下的等待时长为负数时表示超时间隔为相对时长,单位是纳秒。
//
// 有关于内核下的等待时长(即 KeWaitForSingleObject 的 Timeout 参数),参考
// http://msdn.microsoft.com/en-us/library/windows/hardware/ff553350(v=vs.85).aspx。
//

#define MILLISECOND_TO_KERNEL_DELAY(Millisecond) (-1LL * (Millisecond) * 1000 * 10)

//
// 毫秒转换为秒。
//
// 此处是整除,若不是整数秒,则会有精度损失。
//

#define MILLISECOND_TO_SECOND(Millisecond) ((Millisecond) / 1000)

//
// 获取用于附加操作的设备的设备扩展。
//

#define ATTACHING_DEVICE_EXTENSION(DeviceObject) ((PATTACHING_DEVICE_EXTENSION)((DeviceObject)->DeviceExtension))

//
// 用于记录网络使用信息的通用表的节点的缓冲池标记。
//

#define NETWORK_USAGE_GENERIC_TABLE_NODE_TAG 'NUGT'

//
// TDI 的 TCP 设备名。
//

#define TCP_DEVICE_NAME L"\\Device\\Tcp"

//
// TDI 的 UDP 设备名。
//

#define UDP_DEVICE_NAME L"\\Device\\Udp"

//
// TDI 的 RAW IP 设备名。
//

#define RAWIP_DEVICE_NAME L"\\Device\\RawIp"

//
// 计算静态数组中的元素个数。
//

#define CountOf(Array) ((sizeof(Array) / sizeof((Array)[0])))

//
// 单进程的网络使用的统计信息。
//

typedef struct _NETWORK_USAGE_INFORMATION {

  //
  // 自从上次输出统计信息后发送的字节数。
  //
  // 每次输出统计信息都会清空发送字节数。
  //

  ULONGLONG BytesSent;

  //
  // 自从上次输出统计信息后接受的字节数。
  //
  // 每次输出统计信息都会清空接受字节数。
  //

  ULONGLONG BytesReceived;
} NETWORK_USAGE_INFORMATION, *PNETWORK_USAGE_INFORMATION;

//
// 储存在通用表中的节点。
//

typedef struct _NETWORK_USAGE_GENERIC_TABLE_NODE {

  //
  // 进程的 PCB 指针。
  //
  // 在通用表中存放时,用于比较结构体相对大小使用的就是 Process 成员。
  //

  PEPROCESS Process;

  //
  // 进程的网络使用信息。
  //

  NETWORK_USAGE_INFORMATION NetworkUsageInformation;
} NETWORK_USAGE_GENERIC_TABLE_NODE, *PNETWORK_USAGE_GENERIC_TABLE_NODE;

//
// 用于统计进程网络使用信息的表。
//

typedef struct _NETWORK_USAGE_GENERIC_TABLE {

  //
  // 用于记录信息的通用表。
  //

  RTL_GENERIC_TABLE Table;

  //
  // 用于同步访问的自旋锁。
  //

  KSPIN_LOCK TableLock;
} NETWORK_USAGE_GENERIC_TABLE, *PNETWORK_USAGE_GENERIC_TABLE;

//
// 输出网络使用情况的工作线程控制块。
//

typedef struct _NETWORK_USAGE_DUMPER_CONTROL_BLOCK {

  //
  // 工作线程的指针。
  //
  
  PETHREAD Worker;

  //
  // 为 TRUE 时要求工作线程退出。
  //

  BOOLEAN Exit;
  
  //
  // 输出网络使用情况的工作线程每隔 NETWORK_USAGE_DUMP_INTERVAL 毫秒
  // 输出一次。如果直接使用 KeDelayExecutionThread,那么在 DriverUnload
  // 中等待工作线程退出最多会需要大约 NETWORK_USAGE_DUMP_INTERVAL 
  // 毫秒的时间。当 NETWORK_USAGE_DUMP_INTERVAL 增大时,DriverUnload
  // 可能会需要相当长的等待时间。
  //
  // 考虑到 DriverUnload 是由系统的工作线程调用的,因此应当尽快返回。
  // 所以此处令工作线程等待 SleepingEvent。SleepingEvent 除非在 Exit
  // 为 TRUE 的时候才会置为信号态。这样,在驱动退出之前。等待 SleepingEvent
  // 的结果总是会超时,这样也就实现了 KeDelayExecutionThread 的效果,
  // 而 DriverUnload 中将 SleepingEvent 置为信号态,这样等待就会立即
  // 结束,从而让工作线程立即退出。
  //
  // 如果将来 WDK 提供了类似 KeAlertThread 的例程,则可以考虑使用
  // KeDelayExecutionThread + KeAlertThread 的组合。
  //
  // 同样,通过投递 APC 也可以唤醒工作线程,但是为了唤醒线程而投递
  // APC 并非一个好的选择。
  //

  KEVENT SleepingEvent;
} NETWORK_USAGE_DUMPER_CONTROL_BLOCK, *PNETWORK_USAGE_DUMPER_CONTROL_BLOCK;

//
// 附加到设备的操作的状态。
//

typedef struct _DEVICE_ATTACH_STATE {

  //
  // 用于附加操作所创建的设备。
  //

  PDEVICE_OBJECT AttachingDevice;

  //
  // 附加到的设备。
  //
  // 由于我们是附加到目标设备之上,因此此设备也是紧挨我们下层的设备。
  //

  PDEVICE_OBJECT AttachedDevice;
} DEVICE_ATTACH_STATE, *PDEVICE_ATTACH_STATE;

//
// 用于附加操作所创建的设备的设备扩展。
//

typedef struct _ATTACHING_DEVICE_EXTENSION {

  //
  // 用于描述附加操作的状态。
  //

  DEVICE_ATTACH_STATE AttachState;
} ATTACHING_DEVICE_EXTENSION, *PATTACHING_DEVICE_EXTENSION;

//
// 用于描述附加操作的数据结构。
//

typedef struct _DEVICE_ATTACHMENT {

  //
  // 要附加到的设备的设备名。
  //

  UNICODE_STRING TargetDeviceName;

  //
  // 为附加操作创建的设备对象。
  //

  PDEVICE_OBJECT AttachingDevice;
} DEVICE_ATTACHMENT, *PDEVICE_ATTACHMENT;

//
// 输出调试信息。
//

#define DbgWrite(Format, ...) DbgPrint("[" PROJECT_NAME "] " __FUNCTION__ ": " Format, __VA_ARGS__)
#define DbgWriteLine(Format, ...) DbgWrite(Format "\n", __VA_ARGS__)

//
// 输出网络使用情况的工作线程的控制块。
//

NETWORK_USAGE_DUMPER_CONTROL_BLOCK NetworkUsageDumperControlBlock;

//
// 统计网络使用状况的表。
//

NETWORK_USAGE_GENERIC_TABLE NetworkUsageGenericTable;

//
// 需要进行的附加驱动的操作。
//

DEVICE_ATTACHMENT DeviceAttachments[] = {
  {RTL_CONSTANT_STRING(TCP_DEVICE_NAME), 0},
  {RTL_CONSTANT_STRING(UDP_DEVICE_NAME), 0},
  {RTL_CONSTANT_STRING(RAWIP_DEVICE_NAME), 0}
};

//
// 驱动对象。
//

PDRIVER_OBJECT FlowWatcherDriverObject;

//
// 函数原型声明。
//

RTL_GENERIC_COMPARE_ROUTINE NetworkUsageGenericTableComparaRoutine;
RTL_GENERIC_ALLOCATE_ROUTINE NetworkUsageGenericTableAllocateRoutine;
RTL_GENERIC_FREE_ROUTINE NetworkUsageGenericTableFreeRoutine;
DRIVER_DISPATCH GeneralPassthroughDispath;
DRIVER_INITIALIZE DriverEntry;
DRIVER_UNLOAD DriverUnload;
KSTART_ROUTINE NetworkUsageDumper;

VOID InitializeNetworkUsageGenericTable(
  VOID
  );

NTSTATUS AttachDevice(
  __out PDEVICE_ATTACH_STATE DeviceAttachState,
  __in PDEVICE_OBJECT AttachingDevice,
  __in PUNICODE_STRING DeviceName
  );

NTSTATUS CreateAndAttachDevice(
  __out PDEVICE_OBJECT *AttachingDevice,
  __in PUNICODE_STRING DeviceName
  );

VOID DetachDevice(
  __in PDEVICE_ATTACH_STATE DeviceAttachState
  );

VOID DeleteAndDetachDevice(
  __in PDEVICE_ATTACH_STATE DeviceAttachState
  );

NTSTATUS AttachDevices(
  VOID
  );

VOID DetachDevices(
  VOID
  );

NTSTATUS CreateNetworkUsageDumper(
  VOID
  );

VOID TerminateNetworkUsageDumper(
  VOID
  );

//
// 将函数关联到对应的代码段。
//

#ifdef ALLOC_PRAGMA

#pragma alloc_text(INIT, DriverEntry)
#pragma alloc_text(INIT, InitializeNetworkUsageGenericTable)
#pragma alloc_text(INIT, CreateNetworkUsageDumper)
#pragma alloc_text(INIT, AttachDevice)
#pragma alloc_text(INIT, CreateAndAttachDevice)
#pragma alloc_text(INIT, AttachDevices)

#pragma alloc_text(PAGE, DriverUnload)
#pragma alloc_text(PAGE, TerminateNetworkUsageDumper)
#pragma alloc_text(PAGE, NetworkUsageDumper)
#pragma alloc_text(PAGE, DetachDevice)
#pragma alloc_text(PAGE, DeleteAndDetachDevice)
#pragma alloc_text(PAGE, DetachDevices)

#endif

//
// 用于比较两个 NETWORK_USAGE_GENERIC_TABLE_NODE 的比较函数。
//

RTL_GENERIC_COMPARE_RESULTS NTAPI NetworkUsageGenericTableComparaRoutine(
  __in PRTL_GENERIC_TABLE GenericTable,
  __in PVOID FirstObject,
  __in PVOID SecondObject
  )
{
  PNETWORK_USAGE_GENERIC_TABLE_NODE LeftNode, RightNode;

  LeftNode = (PNETWORK_USAGE_GENERIC_TABLE_NODE)FirstObject;
  RightNode = (PNETWORK_USAGE_GENERIC_TABLE_NODE)SecondObject;

  //
  // 此处使用 NETWORK_USAGE_GENERIC_TABLE_NODE::Process 来比较结构体的大小。
  //

  if (LeftNode->Process < RightNode->Process) {
    return GenericLessThan;
  } else if (LeftNode->Process > RightNode->Process) {
    return GenericGreaterThan;
  } else {
    return GenericEqual;
  }
}

//
// 用于分配 NETWORK_USAGE_GENERIC_TABLE_NODE 的函数。
//

PVOID NTAPI NetworkUsageGenericTableAllocateRoutine(
  __in PRTL_GENERIC_TABLE GenericTable,
  __in CLONG BufferSize
  )
{
  return ExAllocatePoolWithTag(NonPagedPool, BufferSize, NETWORK_USAGE_GENERIC_TABLE_NODE_TAG);
}

//
// 用于释放 NETWORK_USAGE_GENERIC_TABLE_NODE 的函数。
//

VOID NTAPI NetworkUsageGenericTableFreeRoutine(
  __in PRTL_GENERIC_TABLE GenericTable,
  __in PVOID Buffer
  )
{
  ExFreePoolWithTag(Buffer, NETWORK_USAGE_GENERIC_TABLE_NODE_TAG);

  return;
}

//
// 初始化统计网络使用情况的通用表。
//

VOID InitializeNetworkUsageGenericTable(
  VOID
  )
{
  //
  // 初始化用于同步访问的自旋锁。
  //

  KeInitializeSpinLock(&NetworkUsageGenericTable.TableLock);

  //
  // 初始化用于记录信息的通用表。
  //

  RtlInitializeGenericTable(&NetworkUsageGenericTable.Table, 
    NetworkUsageGenericTableComparaRoutine,
    NetworkUsageGenericTableAllocateRoutine, 
    NetworkUsageGenericTableFreeRoutine,
    NULL);

  return;
}

//
// 释放统计网络使用情况的通用表(不上锁)。
//

VOID DeleteNetworkUsageGenericTableUnsafe(
  VOID
  )
{  
  PVOID CurrentEntry;

  //
  // 逐个删除元素。
  //

  for (CurrentEntry = RtlEnumerateGenericTable(&NetworkUsageGenericTable.Table, TRUE);
    CurrentEntry;
    CurrentEntry = RtlEnumerateGenericTable(&NetworkUsageGenericTable.Table, FALSE)) {
    BOOLEAN DeleteStatus;

    //
    // 删除元素。
    //

    DeleteStatus = RtlDeleteElementGenericTable(&NetworkUsageGenericTable.Table, CurrentEntry);

    ASSERT(DeleteStatus == TRUE);
  }

  return;
}

//
// 释放统计网络使用情况的通用表。
//

VOID DeleteNetworkUsageGenericTable(
  VOID
  )
{  
  KLOCK_QUEUE_HANDLE LockHandle;
  PVOID CurrentEntry;

  //
  // 获取自旋锁。
  //

  KeAcquireInStackQueuedSpinLock(&NetworkUsageGenericTable.TableLock,
    &LockHandle);

  //
  // 清空通用表。
  //

  DeleteNetworkUsageGenericTableUnsafe();

  //
  // 释放自旋锁。
  //

  KeReleaseInStackQueuedSpinLock(&LockHandle);

  return;
}

//
// 获取某进程的网络使用情况(不上锁)。
//

BOOLEAN QueryProcessNetworkUsageUnsafe(
  __in PEPROCESS Process,
  __out PNETWORK_USAGE_INFORMATION NetworkUsageInformation
  )
{
  NETWORK_USAGE_GENERIC_TABLE_NODE TableNode;
  PVOID EntryFound;

  TableNode.Process = Process;

  //
  // 查询使用情况。
  //

  EntryFound = RtlLookupElementGenericTable(&NetworkUsageGenericTable.Table, 
    (PVOID)&TableNode);

  if (!EntryFound) {
    return FALSE;
  }

  //
  // 返回信息。
  //

  *NetworkUsageInformation = ((PNETWORK_USAGE_GENERIC_TABLE_NODE)EntryFound)->NetworkUsageInformation;

  return TRUE;
}

//
// 获取某进程的网络使用情况。
//

BOOLEAN QueryProcessNetworkUsage(
  __in PEPROCESS Process,
  __out PNETWORK_USAGE_INFORMATION NetworkUsageInformation
  )
{
  KLOCK_QUEUE_HANDLE LockHandle;
  BOOLEAN Status;

  //
  // 获取自旋锁。
  //

  KeAcquireInStackQueuedSpinLock(&NetworkUsageGenericTable.TableLock,
    &LockHandle);

  //
  // 查询使用情况。
  //

  Status = QueryProcessNetworkUsageUnsafe(Process, NetworkUsageInformation);

  //
  // 释放自旋锁。
  //

  KeReleaseInStackQueuedSpinLock(&LockHandle);

  if (Status != TRUE) {
    return Status;
  }

  return Status;
}

//
// 更新某进程的网络使用情况(不上锁)。
//

VOID UpdateProcessNetworkUsageUnsafe(
  __in PEPROCESS Process,
  __in PNETWORK_USAGE_INFORMATION NetworkUsageInformation
  )
{
  NETWORK_USAGE_GENERIC_TABLE_NODE TableNode;
  PVOID EntryFound;

  TableNode.Process = Process;
  TableNode.NetworkUsageInformation = *NetworkUsageInformation;

  //
  // 删除已有节点。
  //

  EntryFound = RtlLookupElementGenericTable(&NetworkUsageGenericTable.Table, 
    (PVOID)&TableNode);

  if (EntryFound) {
    BOOLEAN Status;

    Status = RtlDeleteElementGenericTable(&NetworkUsageGenericTable.Table, EntryFound);

    ASSERT(Status == TRUE);
  }

  //
  // 插入新节点。
  //

  (VOID)RtlInsertElementGenericTable(&NetworkUsageGenericTable.Table, 
    (PVOID)&TableNode, sizeof(NETWORK_USAGE_GENERIC_TABLE_NODE), NULL);

  return;
}

//
// 更新某进程的网络使用情况。
//

VOID UpdateProcessNetworkUsage(
  __in PEPROCESS Process,
  __in PNETWORK_USAGE_INFORMATION NetworkUsageInformation
  )
{
  KLOCK_QUEUE_HANDLE LockHandle;

  //
  // 获取自旋锁。
  //

  KeAcquireInStackQueuedSpinLock(&NetworkUsageGenericTable.TableLock,
    &LockHandle);

  //
  // 更新网络使用情况。
  //

  UpdateProcessNetworkUsageUnsafe(Process, NetworkUsageInformation);

  //
  // 释放自旋锁。
  //

  KeReleaseInStackQueuedSpinLock(&LockHandle);

  return;
}

//
// 输出网络使用情况的统计信息。
//

VOID DumpNetworkUsage(
  VOID
  )
{
  KLOCK_QUEUE_HANDLE LockHandle;
  PVOID CurrentEntry;

  //
  // 获取自旋锁。
  //

  KeAcquireInStackQueuedSpinLock(&NetworkUsageGenericTable.TableLock,
    &LockHandle);

  //
  // 遍历统计网络使用情况的通用表。
  //

  for (CurrentEntry = RtlEnumerateGenericTable(&NetworkUsageGenericTable.Table, TRUE);
    CurrentEntry;
    CurrentEntry = RtlEnumerateGenericTable(&NetworkUsageGenericTable.Table, FALSE)) {
    PNETWORK_USAGE_GENERIC_TABLE_NODE TableNode;
    PEPROCESS Process;
    NETWORK_USAGE_INFORMATION NetworkUsageInformation;

    TableNode = (PNETWORK_USAGE_GENERIC_TABLE_NODE)CurrentEntry;

    Process = TableNode->Process;
    NetworkUsageInformation = TableNode->NetworkUsageInformation;

    //
    // 输出信息。
    //
    // DbgPrint 在输出 64 位整数时似乎不能识别 %lld(至少不能识别 %10lld),因此
    // 此处使用 %I64d。
    //

    DbgWriteLine("Process Id: %4d, Sent: %10I64d bytes / s Received: %10I64d bytes / s.",
      PsGetProcessId(Process), 
      NetworkUsageInformation.BytesSent / MILLISECOND_TO_SECOND(NETWORK_USAGE_DUMP_INTERVAL),
      NetworkUsageInformation.BytesReceived / MILLISECOND_TO_SECOND(NETWORK_USAGE_DUMP_INTERVAL));
  }

  //
  // 清空统计数据。
  //

  DeleteNetworkUsageGenericTableUnsafe();

  //
  // 释放自旋锁。
  //

  KeReleaseInStackQueuedSpinLock(&LockHandle);

  return;
}

//
// 用于输出当前网络使用情况的工作线程。
//

VOID NetworkUsageDumper(
  __in PVOID Context
  )
{
  PNETWORK_USAGE_DUMPER_CONTROL_BLOCK ControlBlock;

  PAGED_CODE();

  //
  // 获得控制块指针。
  //

  ControlBlock = (PNETWORK_USAGE_DUMPER_CONTROL_BLOCK)Context;

  while (!ControlBlock->Exit) {
    LARGE_INTEGER DumpInterval;
    
    //
    // 每隔 NETWORK_USAGE_DUMP_INTERVAL 毫秒输出一次统计信息。
    //
    // 有关于 SleepingEvent 的使用,参考 NETWORK_USAGE_DUMPER_CONTROL_BLOCK
    // 的定义。
    //

    DumpInterval.QuadPart = MILLISECOND_TO_KERNEL_DELAY(NETWORK_USAGE_DUMP_INTERVAL);

    (VOID)KeWaitForSingleObject((PVOID)&ControlBlock->SleepingEvent, Executive, KernelMode,
      FALSE, &DumpInterval);

    //
    // 输出统计信息。
    //

    DumpNetworkUsage();
  }

  return;
}

//
// 附加到指定设备。
//

NTSTATUS AttachDevice(
  __out PDEVICE_ATTACH_STATE DeviceAttachState,
  __in PDEVICE_OBJECT AttachingDevice,
  __in PUNICODE_STRING DeviceName
  )
{
  NTSTATUS Status;

  //
  // 附加到指定设备。
  //

  Status = IoAttachDevice(AttachingDevice, DeviceName, &DeviceAttachState->AttachedDevice);

  if (!NT_SUCCESS(Status)) {
    return Status;
  }

  //
  // 填充附加状态信息。
  //

  DeviceAttachState->AttachingDevice = AttachingDevice;

  return Status;
}

//
// 创建一个设备对象并将其附加到指定的设备上。
//

NTSTATUS CreateAndAttachDevice(
  __out PDEVICE_OBJECT *AttachingDevice,
  __in PUNICODE_STRING DeviceName
  )
{
  PATTACHING_DEVICE_EXTENSION DeviceExtension;
  PDEVICE_OBJECT LocalAttachingDevice;
  NTSTATUS Status;

  //
  // 创建用于附加操作的设备对象。
  //

  Status = IoCreateDevice(FlowWatcherDriverObject,
    sizeof(ATTACHING_DEVICE_EXTENSION),  NULL, FILE_DEVICE_NETWORK, 0, FALSE,
    &LocalAttachingDevice);

  if (!NT_SUCCESS(Status)) {
    return Status;
  }

  DeviceExtension = (PATTACHING_DEVICE_EXTENSION)LocalAttachingDevice->DeviceExtension;

  //
  // 将设备附加到指定的设备上。
  //

  Status = AttachDevice(&DeviceExtension->AttachState, LocalAttachingDevice, DeviceName);

  if (!NT_SUCCESS(Status)) {
    IoDeleteDevice(LocalAttachingDevice);

    return Status;
  }

  //
  // 将新创建的设备对象返回给调用方。
  //

  *AttachingDevice = LocalAttachingDevice;

  return Status;
}

//
// 从之前附加到的设备上分离。
//

VOID DetachDevice(
  __in PDEVICE_ATTACH_STATE DeviceAttachState
  )
{
  PAGED_CODE();

  //
  // 分离。
  //

  IoDetachDevice(DeviceAttachState->AttachedDevice);

  return;
}

//
// 从之前附加到的设备上分离并释放为附加操作创建的设备对象。
//

VOID DeleteAndDetachDevice(
  __in PDEVICE_ATTACH_STATE DeviceAttachState
  )
{
  PAGED_CODE();

  //
  // 分离设备。
  //

  DetachDevice(DeviceAttachState);

  //
  // 删除设备。
  //

  IoDeleteDevice(DeviceAttachState->AttachingDevice);

  return;
}

//
// 附加到各个(TCP / UDP / RAW IP) TDI 驱动的设备对象上。
//

NTSTATUS AttachDevices(
  VOID
  )
{
  ULONG DeviceIndex;
  NTSTATUS Status;  

  //
  // 附加到 TDI 的设备上。
  //

  for (DeviceIndex = 0; DeviceIndex < CountOf(DeviceAttachments); DeviceIndex++) {
    ULONG DeviceIndex2;

    //
    // 附加到该设备上。
    //

    Status = CreateAndAttachDevice(&DeviceAttachments[DeviceIndex].AttachingDevice,
      &DeviceAttachments[DeviceIndex].TargetDeviceName);

    if (!NT_SUCCESS(Status)) {

      //
      // 撤销已经成功的附加操作。
      //

      for (DeviceIndex2 = 0; DeviceIndex2 < DeviceIndex; DeviceIndex2++) {

        //
        // 分离设备并删除创建的设备。
        //

        DeleteAndDetachDevice(&ATTACHING_DEVICE_EXTENSION(
          DeviceAttachments[DeviceIndex2].AttachingDevice)->AttachState);
      }

      break;
    }
  }

  return Status;
}

//
// 从 TDI 的各个设备对象上分离。
//

VOID DetachDevices(
  VOID
  )
{
  ULONG DeviceIndex;

  PAGED_CODE();
  
  //
  // 撤销附加到 TDI 的设备的操作。
  //
  
  for (DeviceIndex = 0; DeviceIndex < CountOf(DeviceAttachments); DeviceIndex++) {

    //
    // 分离设备并删除创建的设备。
    //

    DeleteAndDetachDevice(&ATTACHING_DEVICE_EXTENSION(
      DeviceAttachments[DeviceIndex].AttachingDevice)->AttachState);
  }

  return;
}

//
// 创建用于输出网络使用情况的线程。
//

NTSTATUS CreateNetworkUsageDumper(
  VOID
  )
{
  HANDLE ThreadHandle;
  NTSTATUS Status;

  //
  // 初始化线程控制块。
  //

  NetworkUsageDumperControlBlock.Exit = FALSE;

  //
  // 用于线程休眠的事件。
  //

  KeInitializeEvent(&NetworkUsageDumperControlBlock.SleepingEvent, NotificationEvent, FALSE);

  //
  // 创建线程。
  //

  Status = PsCreateSystemThread(&ThreadHandle, 0, NULL, 0, NULL, NetworkUsageDumper, 
    (PVOID)&NetworkUsageDumperControlBlock);

  if (!NT_SUCCESS(Status)) {
    return Status;
  }

  //
  // 获得线程的 ETHREAD。
  //

  Status = ObReferenceObjectByHandle(ThreadHandle, 0, *PsThreadType, KernelMode, 
    (PVOID *)&NetworkUsageDumperControlBlock.Worker, NULL);

  ZwClose(ThreadHandle);

  if (!NT_SUCCESS(Status)) {
    NetworkUsageDumperControlBlock.Exit = TRUE;
    KeSetEvent(&NetworkUsageDumperControlBlock.SleepingEvent, 0, FALSE);

    return Status;
  }

  return Status;
}

//
// 结束用于输出网络使用情况的线程。
//

VOID TerminateNetworkUsageDumper(
  VOID
  )
{
  PAGED_CODE();

  //
  // 通知线程退出。
  //

  NetworkUsageDumperControlBlock.Exit = TRUE;

  //
  // 唤醒线程。
  //

  KeSetEvent(&NetworkUsageDumperControlBlock.SleepingEvent, 0, TRUE);

  //
  // 等待线程退出。
  //

  KeWaitForSingleObject((PVOID)NetworkUsageDumperControlBlock.Worker, Executive, KernelMode,
    FALSE, NULL);

  return;
}

//
// 用于统计进程的网络使用情况。
//
// 此处只能处理 IRP_MJ_INTERNAL_DEVICE_CONTROL 类型的 IRP,对于
// IRP_MJ_DEVICE_CONTROL,需要使用 TdiMapUserRequest 将其映射为
// IRP_MJ_INTERNAL_DEVICE_CONTROL 类型的 IRP。
//

VOID AnalyzeNetworkIrp(
  __in PIRP Irp
  )
{
  PEPROCESS Process;
  PIO_STACK_LOCATION IoStackLocation;
  NETWORK_USAGE_INFORMATION NetworkUsageInformation;

  IoStackLocation = IoGetCurrentIrpStackLocation(Irp);

  ASSERT(IoStackLocation->MajorFunction == IRP_MJ_INTERNAL_DEVICE_CONTROL);

  Process = PsGetCurrentProcess();

  if (QueryProcessNetworkUsage(Process, &NetworkUsageInformation) == FALSE) {

    //
    // 若查询不到当前进程的网络使用情况,则表示通用表中还没有这个进程
    // 相关的表项,那么就将该进程的网络使用情况初始化成 0。
    //

    NetworkUsageInformation.BytesReceived = 0;
    NetworkUsageInformation.BytesSent = 0;
  }

  if (IoStackLocation->MinorFunction == TDI_RECEIVE) {
    PTDI_REQUEST_KERNEL_RECEIVE Parameters;

    Parameters = (PTDI_REQUEST_KERNEL_RECEIVE)&IoStackLocation->Parameters;

    //
    // 统计接受的字节数。
    //

    NetworkUsageInformation.BytesReceived += Parameters->ReceiveLength;
  } else if (IoStackLocation->MinorFunction == TDI_SEND) {
    PTDI_REQUEST_KERNEL_SEND Parameters;

    Parameters = (PTDI_REQUEST_KERNEL_SEND)&IoStackLocation->Parameters;

    //
    // 统计发送的字节数。
    //

    NetworkUsageInformation.BytesSent += Parameters->SendLength;
  } else if (IoStackLocation->MinorFunction == TDI_RECEIVE_DATAGRAM) {
    PTDI_REQUEST_KERNEL_RECEIVEDG Parameters;

    Parameters = (PTDI_REQUEST_KERNEL_RECEIVEDG)&IoStackLocation->Parameters;

    //
    // 统计接受的字节数。
    //

    NetworkUsageInformation.BytesReceived += Parameters->ReceiveLength;
  } else if (IoStackLocation->MinorFunction == TDI_SEND_DATAGRAM) {
    PTDI_REQUEST_KERNEL_SENDDG Parameters;

    Parameters = (PTDI_REQUEST_KERNEL_SENDDG)&IoStackLocation->Parameters;

    //
    // 统计发送的字节数。
    //

    NetworkUsageInformation.BytesSent += Parameters->SendLength;
  }

  //
  // 更新进程的网络使用情况。
  //

  UpdateProcessNetworkUsage(Process, &NetworkUsageInformation);

  return;
}

//
// 通用的传递 IRP 到下层驱动的派遣函数。
//
 
NTSTATUS GeneralPassthroughDispath(
  __in PDEVICE_OBJECT DeviceObject,
  __inout PIRP Irp
  )
{
  PDEVICE_OBJECT LowerDeviceObject;
  PIO_STACK_LOCATION IoStackLocation;

  ASSERT(ATTACHING_DEVICE_EXTENSION(DeviceObject));

  LowerDeviceObject = ATTACHING_DEVICE_EXTENSION(DeviceObject)->AttachState.AttachedDevice;

  IoStackLocation = IoGetCurrentIrpStackLocation(Irp);

  //
  // 对于 IRP_MJ_DEVICE_CONTROL 类型请求,将其转换为 IRP_MJ_INTERNAL_DEVICE_CONTROL
  // 类型后再处理。
  //

  if (IoStackLocation->MajorFunction == IRP_MJ_DEVICE_CONTROL) {

    //
    // 若转换失败,则 IRP 依然是 IRP_MJ_DEVICE_CONTROL 类型,那么就不会被提交给
    // AnalyzeNetworkIrp 做进一步处理。
    //
    // 由于稍后会再次判断 IRP 的类型,因此此处可以忽略 TdiMapUserRequest 的返回值。
    //

    (VOID)TdiMapUserRequest(DeviceObject, Irp, IoStackLocation);
  }

  //
  // 若是 IRP_MJ_INTERNAL_DEVICE_CONTROL 类型请求,则记录进程的网络使用情况。
  //

  if (IoStackLocation->MajorFunction == IRP_MJ_INTERNAL_DEVICE_CONTROL) {
    AnalyzeNetworkIrp(Irp);
  }

  //
  // 跳过当前 IRP 参数栈。
  //

  IoSkipCurrentIrpStackLocation(Irp);

  //
  // 下发 IRP。
  //

  return IoCallDriver(LowerDeviceObject, Irp);
}

//
// 驱动入口点。
//

NTSTATUS DriverEntry(
  __inout PDRIVER_OBJECT DriverObject,
  __in PUNICODE_STRING RegistryPath
  )
{
  ULONG MajorFunctionIndex;
  NTSTATUS Status;

  FlowWatcherDriverObject = DriverObject;

  //
  // 填充回调表。
  //

  DriverObject->DriverUnload = DriverUnload;

  for (MajorFunctionIndex = 0; MajorFunctionIndex <= IRP_MJ_MAXIMUM_FUNCTION;
    MajorFunctionIndex++) {
    DriverObject->MajorFunction[MajorFunctionIndex] = GeneralPassthroughDispath;
  }

  //
  // 初始化记录进程网络使用情况的通用表。
  //

  InitializeNetworkUsageGenericTable();

  //
  // 附加到 TDI 的设备对象上。
  //

  Status = AttachDevices();

  if (!NT_SUCCESS(Status)) {
    return Status;
  }

  //
  // 创建输出网络使用信息的线程。
  //

  Status = CreateNetworkUsageDumper();

  if (!NT_SUCCESS(Status)) {
    DetachDevices();

    return Status;
  }
  
  return Status;
}

//
// 驱动卸载回调。
//

VOID DriverUnload(
  __in PDRIVER_OBJECT DriverObject
  )
{
  PAGED_CODE();

  //
  // 结束输出网络使用情况的线程。
  //

  TerminateNetworkUsageDumper();

  //
  // 从 TDI 的设备对象上分离。
  //

  DetachDevices();

  //
  // 释放用于统计进程网络使用情况的通用表。
  //

  DeleteNetworkUsageGenericTable();

  return;
}

Update at 19:24 2012/8/21:

似乎应该在 IoAttachDevice 之前去掉 DeviceObject->Flags 中的 DO_DEVICE_INITIALIZING,同时要把 TDI 的 DeviceObject 的 Flags 复制过来。
另外用 Timer Object 要比用 Event 来的合理。

  1. 本文目前尚无任何评论.