基于AFD驱动的进程流量控制

基于AFD驱动的进程流量控制

摘要：目前有些软件可以监控进程流量，功能实现的都很多错的。对于进程流量的控制很很多种方案，每一种方案也都有其缺点。比如有应用层基于LSP来做的，也有通过TDI和NDIS中间层来做的。其实现的效果和复杂度也各不相同。而我们要讨论的是一种基于AFD驱动来做的一种进程流量控制方案。

    关键词：AFD 进程流量  控制驱动

1.   总述

目前有些软件可以监控进程流量，功能实现的都很多错的。对于进程流量的控制很很多种方案，每一种方案也都有其缺点。比如有应用层基于LSP来做的，也有通过TDI和NDIS中间层来做的。其实现的效果和复杂度也各不相同。而我们要讨论的是一种基于AFD驱动来做的一种进程流量控制方案。当然这里说的流量控制和传统意义上的流量控制不一样，这里讨论的进程流量控制是按照用户的设置来放慢进程发送和接收数据的速度从而减少带宽的占用。

2.   AFD驱动

首先来介绍一下AFD驱动。大家可能对NDIS及TDI了解的比较多，AFD驱动相对会少一些。NDIS及TDI是微软提供的内核网络驱动模型中的编程接口规范，而AFD驱动是微软Windows操作一个驱动部件。它是上层SOCKET在内核心中的实现。我这里画了一个图例来说明它们之间的关系：

简来说AFD驱动向上与SOCKET应用接口约定了接口来实现SOCKET，AFD驱动实际上是一个TDI客户端，它通过TDI接口调用微软件的另一个网络部件TCPIP驱动来完成功能。AFD并没有官方的资料说明它的接口，但是在网上还是可以找到很关于AFD驱动的资料的。这里笔者参考了ReactOS-0.3.4-REL-src源代码中的AFD的内容，它于官方的AFD驱动实现还是有一些驱动别的，比如未实现FASTIO接口等。

在AFD中它主要处理如下IO控制码，上层的应用也主要是通过它们来完成SOCKET各种操作。

/* IOCTL Generation */

#define FSCTL_AFD_BASE                  FILE_DEVICE_NETWORK

#define _AFD_CONTROL_CODE(Operation,Method) /

((FSCTL_AFD_BASE)<<12 | (Operation<<2) | Method)

/* AFD Commands */

#define AFD_BIND            0

#define AFD_CONNECT              1

#define AFD_START_LISTEN         2

#define AFD_WAIT_FOR_LISTEN      3

#define AFD_ACCEPT          4

#define AFD_RECV            5

#define AFD_RECV_DATAGRAM        6

#define AFD_SEND            7

#define AFD_SEND_DATAGRAM        8

#define AFD_SELECT          9

#define AFD_DISCONNECT           10

#define AFD_GET_SOCK_NAME        11

#define AFD_GET_PEER_NAME               12

#define AFD_GET_TDI_HANDLES      13

#define AFD_SET_INFO             14

#define AFD_GET_CONTEXT          16

#define AFD_SET_CONTEXT          17

#define AFD_SET_CONNECT_DATA         18

#define AFD_SET_CONNECT_OPTIONS      19

#define AFD_SET_DISCONNECT_DATA      20

#define AFD_SET_DISCONNECT_OPTIONS   21

#define AFD_GET_CONNECT_DATA         22

#define AFD_GET_CONNECT_OPTIONS      23

#define AFD_GET_DISCONNECT_DATA      24

#define AFD_GET_DISCONNECT_OPTIONS   25

#define AFD_SET_CONNECT_DATA_SIZE       26

#define AFD_SET_CONNECT_OPTIONS_SIZE    27

#define AFD_SET_DISCONNECT_DATA_SIZE    28

#define AFD_SET_DISCONNECT_OPTIONS_SIZE 29

#define AFD_GET_INFO             30

#define AFD_EVENT_SELECT         33

#define AFD_ENUM_NETWORK_EVENTS         34

#define AFD_DEFER_ACCEPT         35

#define AFD_GET_PENDING_CONNECT_DATA 41

/* AFD IOCTLs */

#define IOCTL_AFD_BIND /

_AFD_CONTROL_CODE(AFD_BIND, METHOD_NEITHER)

#define IOCTL_AFD_CONNECT /

_AFD_CONTROL_CODE(AFD_CONNECT, METHOD_NEITHER)

#define IOCTL_AFD_START_LISTEN /

_AFD_CONTROL_CODE(AFD_START_LISTEN, METHOD_NEITHER)

#define IOCTL_AFD_WAIT_FOR_LISTEN /

_AFD_CONTROL_CODE(AFD_WAIT_FOR_LISTEN, METHOD_BUFFERED )

#define IOCTL_AFD_ACCEPT /

_AFD_CONTROL_CODE(AFD_ACCEPT, METHOD_BUFFERED )

#define IOCTL_AFD_RECV /

_AFD_CONTROL_CODE(AFD_RECV, METHOD_NEITHER)

#define IOCTL_AFD_RECV_DATAGRAM /

_AFD_CONTROL_CODE(AFD_RECV_DATAGRAM, METHOD_NEITHER)

#define IOCTL_AFD_SEND /

_AFD_CONTROL_CODE(AFD_SEND, METHOD_NEITHER)

#define IOCTL_AFD_SEND_DATAGRAM /

_AFD_CONTROL_CODE(AFD_SEND_DATAGRAM, METHOD_NEITHER)

#define IOCTL_AFD_SELECT /

_AFD_CONTROL_CODE(AFD_SELECT, METHOD_BUFFERED )

#define IOCTL_AFD_DISCONNECT /

_AFD_CONTROL_CODE(AFD_DISCONNECT, METHOD_NEITHER)

#define IOCTL_AFD_GET_SOCK_NAME /

_AFD_CONTROL_CODE(AFD_GET_SOCK_NAME, METHOD_NEITHER)

#define IOCTL_AFD_GET_PEER_NAME /

_AFD_CONTROL_CODE(AFD_GET_PEER_NAME, METHOD_NEITHER)

#define IOCTL_AFD_GET_TDI_HANDLES /

_AFD_CONTROL_CODE(AFD_GET_TDI_HANDLES, METHOD_NEITHER)

#define IOCTL_AFD_SET_INFO /

_AFD_CONTROL_CODE(AFD_SET_INFO, METHOD_NEITHER)

#define IOCTL_AFD_GET_CONTEXT /

_AFD_CONTROL_CODE(AFD_GET_CONTEXT, METHOD_NEITHER)

#define IOCTL_AFD_SET_CONTEXT /

_AFD_CONTROL_CODE(AFD_SET_CONTEXT, METHOD_NEITHER)

#define IOCTL_AFD_SET_CONNECT_DATA /

_AFD_CONTROL_CODE(AFD_SET_CONNECT_DATA, METHOD_NEITHER)

#define IOCTL_AFD_SET_CONNECT_OPTIONS /

_AFD_CONTROL_CODE(AFD_SET_CONNECT_OPTIONS, METHOD_NEITHER)

#define IOCTL_AFD_SET_DISCONNECT_DATA /

_AFD_CONTROL_CODE(AFD_SET_DISCONNECT_DATA, METHOD_NEITHER)

#define IOCTL_AFD_SET_DISCONNECT_OPTIONS /

_AFD_CONTROL_CODE(AFD_SET_DISCONNECT_OPTIONS, METHOD_NEITHER)

#define IOCTL_AFD_GET_CONNECT_DATA /

_AFD_CONTROL_CODE(AFD_GET_CONNECT_DATA, METHOD_NEITHER)

#define IOCTL_AFD_GET_CONNECT_OPTIONS /

_AFD_CONTROL_CODE(AFD_GET_CONNECT_OPTIONS, METHOD_NEITHER)

#define IOCTL_AFD_GET_DISCONNECT_DATA /

_AFD_CONTROL_CODE(AFD_GET_DISCONNECT_DATA, METHOD_NEITHER)

#define IOCTL_AFD_GET_DISCONNECT_OPTIONS /

_AFD_CONTROL_CODE(AFD_GET_DISCONNECT_OPTIONS, METHOD_NEITHER)

#define IOCTL_AFD_SET_CONNECT_DATA_SIZE /

_AFD_CONTROL_CODE(AFD_SET_CONNECT_DATA_SIZE, METHOD_NEITHER)

#define IOCTL_AFD_SET_CONNECT_OPTIONS_SIZE /

_AFD_CONTROL_CODE(AFD_SET_CONNECT_OPTIONS_SIZE, METHOD_NEITHER)

#define IOCTL_AFD_SET_DISCONNECT_DATA_SIZE /

_AFD_CONTROL_CODE(AFD_SET_DISCONNECT_DATA_SIZE, METHOD_NEITHER)

#define IOCTL_AFD_SET_DISCONNECT_OPTIONS_SIZE /

_AFD_CONTROL_CODE(AFD_SET_DISCONNECT_OPTIONS_SIZE, METHOD_NEITHER)

#define IOCTL_AFD_GET_INFO /

_AFD_CONTROL_CODE(AFD_GET_INFO, METHOD_NEITHER)

#define IOCTL_AFD_EVENT_SELECT /

_AFD_CONTROL_CODE(AFD_EVENT_SELECT, METHOD_NEITHER)

#define IOCTL_AFD_DEFER_ACCEPT /

_AFD_CONTROL_CODE(AFD_DEFER_ACCEPT, METHOD_NEITHER)

#define IOCTL_AFD_GET_PENDING_CONNECT_DATA /

_AFD_CONTROL_CODE(AFD_GET_PENDING_CONNECT_DATA, METHOD_NEITHER)

#define IOCTL_AFD_ENUM_NETWORK_EVENTS /

_AFD_CONTROL_CODE(AFD_ENUM_NETWORK_EVENTS, METHOD_NEITHER)

可以通过名称看出这些IO控制码，对应了我们常用的SOCKET API。而我们要关注仅仅是的其中的一部分，下面就对我们完成流量控制要关注的IO控制码做一下简单介绍。在这些IO控制码中和发送和接收相关的有四个，它会分别是IOCTL_AFD_SEND、IOCTL_AFD_SEND_DATAGRAM、IOCTL_AFD_RECV、IOCTL_AFD_RECV_DATAGRAM。它们分别用于发送带链接的数据、发送非面向链接的数据包、接收面向链接的数据包、接收非面向链接的数据包。这四个IO控制码进行数据发送和接收时还涉及一些数据结构。这些在后面介绍流量控制时会用到。先在这里说明一下。

typedef struct _AFD_MAPBUF {

PVOID BufferAddress;

PMDL  Mdl;

} AFD_MAPBUF, *PAFD_MAPBUF;

typedef struct _AFD_WSABUF {

UINT  len;

PCHAR buf;

} AFD_WSABUF, *PAFD_WSABUF;

typedef struct  _AFD_RECV_INFO {

PAFD_WSABUF                 BufferArray;

ULONG                       BufferCount;

ULONG                       AfdFlags;

ULONG                       TdiFlags;

} AFD_RECV_INFO , *PAFD_RECV_INFO ;

typedef struct _AFD_RECV_INFO_UDP {

PAFD_WSABUF                 BufferArray;

ULONG                       BufferCount;

ULONG                       AfdFlags;

ULONG                       TdiFlags;

PVOID                       Address;

PINT                   AddressLength;

} AFD_RECV_INFO_UDP, *PAFD_RECV_INFO_UDP;

typedef struct  _AFD_SEND_INFO {

PAFD_WSABUF                 BufferArray;

ULONG                       BufferCount;

ULONG                       AfdFlags;

ULONG                       TdiFlags;

} AFD_SEND_INFO , *PAFD_SEND_INFO ;

typedef struct _AFD_SEND_INFO_UDP {

PAFD_WSABUF                 BufferArray;

ULONG                       BufferCount;

ULONG                       AfdFlags;

ULONG                       Padding[9];

ULONG                       SizeOfRemoteAddress;

PVOID                       RemoteAddress;

} AFD_SEND_INFO_UDP, *PAFD_SEND_INFO_UDP;

其中IOCTL_AFD_SEND，IOCTL_AFD_RECV对应用结构为AFD_SEND_INFO和AFD_RECV_INFO，了解这两个结构我们仅仅是为了知道某次发送或接收的数据长度是多少。这两个结构中都有两个相同的域BufferArray和BufferCount，它们标识了发送或接收的缓冲区信息。AFD_WSABUF结构中len成员标识了一个缓冲区的长度，只要把某次发送或接收操作的BufferCount个缓冲区的长度累加就是本次发送和接收操作数据的总长度。

其中IOCTL_AFD_SEND_DATAGRAM，IOCTL_AFD_RECV_DATAGRAM对应用结构为AFD_SEND_INFO_UDP和AFD_RECV_INFO_UDP。这两个结构中也是都有两个相同的域BufferArray和BufferCount，它们标识了发送或接收的缓冲区信息。AFD_WSABUF结构中len成员标识了一个缓冲区的长度，只要把某次发送或接收操作的BufferCount个缓冲区的长度累加就是本次发送和接收操作数据的总长度。这和IOCTL_AFD_SEND，IOCTL_AFD_RECV对的结构一样的情况。

这节主要介绍了一下Windows中Socket是实现模型及AFD驱动和数据发送与接收相关的四个控制码。下面这一节主要来阐述一下如何对进程流量进行控制的方法。

3.   控制方法

对进程流量控制的方法有很多，总的来说我们想要控制进程的通信流量，最直接的方法就是放慢进程发送与接收操作的速度。那如何放慢进程数据发送和接收的速度呢？如果我们能HOOK到所要进程的发送和接收例程，那么我们就可以统计出每一个进程的流量，如果流量超出限制我们就放慢例程的返回。通过上面的第二节的内容我们知道，所有SOCKET应用的实现都是通过AFD在内核实现的。也就是说要完成发送或接收操作必须与内核层的AFD驱动交互。

上层的SOCKET应用要与内核层的AFD驱动交互主要是通过读写请求和IO控制请求来完成。如果一个上层应用要对一个驱动发起一个IO控制操作如发送数据时向AFD驱动发送一个IOCTL_AFD_SEND控制码并填写一个AFD_SEND_INFO结构给AFD驱动，然而这个操作要通过Kernel32导出的DeviceIoControl函数来完成。这个函数的原型如下：

BOOL DeviceIoControl(
  HANDLE hDevice,              // handle to device of interest
  DWORD dwIoControlCode,       // control code of operation to perform
  LPVOID lpInBuffer,           // pointer to buffer to supply input data
  DWORD nInBufferSize,         // size, in bytes, of input buffer
  LPVOID lpOutBuffer,          // pointer to buffer to receive output data
  DWORD nOutBufferSize,        // size, in bytes, of output buffer
  LPDWORD lpBytesReturned,     // pointer to variable to receive byte count
  LPOVERLAPPED lpOverlapped    // pointer to structure for asynchronous operation
);

dwIoControlCode这个参数就是要向驱动请求的控制码，如果我们HOOK这个函数我们就可在依据这个参数在这里统计和限制进程的流量。但是如果在应用层做HOOK的话会比较麻烦，这里我们采三内核SSDT表HOOK的方式来做。这个函数其实对应用内核的一个服务例程ZwDeviceIoControlFile，这个函数所上面函数参数类似原型如下：

NTSTATUS

ZwDeviceIoControlFile(

IN HANDLE  FileHandle,

IN HANDLE  Event,

IN PIO_APC_ROUTINE  ApcRoutine,

IN PVOID  ApcContext,

OUT PIO_STATUS_BLOCK  IoStatusBlock,

IN ULONG  IoControlCode,

IN PVOID  InputBuffer,

IN ULONG  InputBufferLength,

OUT PVOID  OutputBuffer,

IN ULONG  OutputBufferLength

)

我们只要HOOK这个函数就可以截获上层用应用程序向所有内核驱动发送的所有IO控制码，而我们只关注上层应用程序向AFD驱动发送的IOCTL_AFD_SEND、IOCTL_AFD_SEND_DATAGRAM、IOCTL_AFD_RECV、IOCTL_AFD_RECV_DATAGRAM这个四个IO控制码。当应用程序通过SOCKET发送或接收网络数据时就会发送出上述四个请求中的一个，我们只要分析出每次请求操作多少长度的数据并统计，如果发现流量超标那么我们就计算一个暂停的时间间隔并暂停一下。此外为了实现我们控制的粒度尽量的小和精度尽量的高，我们还需要在转发请求前对请求的数据做一下修正，比如每一次请求的数据长度不能超过策略限定进程每秒流量的峰值。如果我们不这样做当应用程序一次向下请求超出流量峰值很多倍的数据时，将会使程序等待过长的时间这样会引发很多的问题，例如程序无法正常退出等等。下节详细介绍实现的过程。

4.   程序实现

4.1.      初始化

首先我们要在驱动程序的入口函数中初始化进程队列及HOOK服务例程ZwDeviceIoControlFile，下面入口函数的代码。

NTSTATUS

DriverEntry(__in PDRIVER_OBJECT DriverObject,

__in PUNICODE_STRING RegistryPath)

{

UNICODE_STRING DeviceName;

UNICODE_STRING DeviceLinkName;

NTSTATUS Status;

PDEVICE_OBJECT DeviceObject;

DriverObject->MajorFunction[IRP_MJ_CREATE] = OnCreate;

DriverObject->MajorFunction[IRP_MJ_CLOSE] = OnClose;

DriverObject->MajorFunction[IRP_MJ_CLEANUP] = OnCleanup;

DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = OnDeviceIoControl;

DriverObject->DriverUnload = DriverUnload;

do

{

RtlInitUnicodeString(&DeviceName,DEVICE_NAME);

Status = IoCreateDevice(DriverObject,

0,

&DeviceName,

FILE_DEVICE_UNKNOWN,

0,

TRUE,

&DeviceObject);

if (!NT_SUCCESS(Status))

{

break;

}

DeviceObject->Flags |= DO_BUFFERED_IO;

RtlInitUnicodeString(&DeviceLinkName,DEVICE_LINK_NAME);

Status = IoCreateSymbolicLink(&DeviceLinkName,&DeviceName);

if (!NT_SUCCESS(Status))

{

IoDeleteDevice(DeviceObject);

break;

}

// 获得SSDT服务表

KeServiceTablePointers = RegmonMapServiceTable( &HookDescriptors );

if (!KeServiceTablePointers)

{

IoDeleteDevice(DeviceObject);

IoDeleteSymbolicLink(&DeviceLinkName);

Status = STATUS_INSUFFICIENT_RESOURCES;

break;

}

// 初始化进程数据

ProcInitInfoArray();

// HOOK ZwDeviceIoControlFile函数

Hook();

} while(FALSE);

return Status;

}

4.2.      HOOK函数处理

我们在HOOK函数里要做三个操作：1、对请求的数据做一些预处理，因为上层的应用程序可能一次发送的数据会很大，如果超过我们限定流量峰值很多倍的话就会造成暂停时间过长而应用程序无法正常退出的情况，所以我们做一下预处理。2、调用原始的函数。3、获得实际发送或接收的数据长度统计流量并判定流量是否超标，如果超标就计算一下应该暂停之长时间并暂停一下再返回调用。

NTSTATUS

HookZwDeviceIoControlFile(

IN HANDLE  FileHandle,

IN HANDLE  Event,

IN PIO_APC_ROUTINE  ApcRoutine,

IN PVOID  ApcContext,

OUT PIO_STATUS_BLOCK  IoStatusBlock,

IN ULONG  IoControlCode,

IN PVOID  InputBuffer,

IN ULONG  InputBufferLength,

OUT PVOID  OutputBuffer,

IN ULONG  OutputBufferLength

)

{

NTSTATUS Status;

// Hook的前处理函数

DeviceIoControlFileProcessPrec(

FileHandle,

Event,

ApcRoutine,

ApcContext,

IoStatusBlock,

IoControlCode,

InputBuffer,

InputBufferLength,

OutputBuffer,

OutputBufferLength);

Status = RealZwDeviceIoControlFile(

FileHandle,

Event,

ApcRoutine,

ApcContext,

IoStatusBlock,

IoControlCode,

InputBuffer,

InputBufferLength,

OutputBuffer,

OutputBufferLength);

// HOOK的后处理函数

DeviceIoControlFileProcessPost(

FileHandle,

Event,

ApcRoutine,

ApcContext,

IoStatusBlock,

IoControlCode,

InputBuffer,

InputBufferLength,

OutputBuffer,

OutputBufferLength);

return Status;

}

4.3.      预处理

预处理代码稍多一点但并不复杂。首先根据当前进程Id去查询一下是不是有它的流量控制策略，如果有就根据流量控制策略规定的流量峰值求出一个最大传输长度，算法也很简单如果流量峰值小于等于64KB就最大数据传输长度就是峰值的一半（也可以是三分二总之要小于峰值），大于64KB就是就是64KB。这样做是为了不让请求暂停的时间过长而无法退出程序。其次每一个请求操作会下传送一个结构，结构中有一个缓冲区数组；那我们就必须根据最大传输长度，去适当的修改相前的构结。总结一下这一步的目的就是让AFD驱动一次发送或接收不超过峰值长度的数据。

VOID

DeviceIoControlFileProcessPrec(

IN HANDLE  FileHandle,

IN HANDLE  Event,

IN PIO_APC_ROUTINE  ApcRoutine,

IN PVOID  ApcContext,

OUT PIO_STATUS_BLOCK  IoStatusBlock,

IN ULONG  IoControlCode,

IN PVOID  InputBuffer,

IN ULONG  InputBufferLength,

OUT PVOID  OutputBuffer,

IN ULONG  OutputBufferLength

)

{

if (UnloadFlags == FALSE)

{

InterlockedIncrement(&HookCallNumber);

switch(IoControlCode)

{

case IOCTL_AFD_SEND:

{

PAFD_SEND_INFO Info = (PAFD_SEND_INFO)InputBuffer;

ULONG DataSize = 0;

PROCESS_MAX_TRAN_SIZE MaxTranSize;

if (!FlwGetProcessMaxTranSize(&MaxTranSize,PsGetCurrentProcessId()))

{

break;

}

else if (MaxTranSize.MaxTranLength[FLAGS_SEND] < MIN_TRAN_SIZE/2)

{

break;

}

if (Info != NULL &&

InputBufferLength >= sizeof(AFD_SEND_INFO))

{

ULONG i = 0;

if (Info->BufferArray != NULL)

{

for (i = 0; i<Info->BufferCount; i++)

{

if (DataSize ==  MaxTranSize.MaxTranLength[FLAGS_SEND])

{

Info->BufferArray[i].len = 0;

}

else if (DataSize+Info->BufferArray[i].len

> MaxTranSize.MaxTranLength[FLAGS_SEND])

{

Info->BufferArray[i].len = (UINT)

MaxTranSize.MaxTranLength[FLAGS_SEND]- DataSize;

DataSize = (UINT)MaxTranSize.MaxTranLength[FLAGS_SEND];

}

else

{

DataSize += Info->BufferArray[i].len;

}

}

}

}

}

break;

case IOCTL_AFD_SEND_DATAGRAM:

{

PAFD_SEND_INFO_UDP Info = (PAFD_SEND_INFO_UDP)InputBuffer;

ULONG DataSize = 0;

PROCESS_MAX_TRAN_SIZE MaxTranSize;

if (!FlwGetProcessMaxTranSize(&MaxTranSize,PsGetCurrentProcessId()))

{

break;

}

else if (MaxTranSize.MaxTranLength[FLAGS_SEND] < MIN_FLOW_DATA/2)

{

break;

}

if (Info != NULL &&

InputBufferLength >= sizeof(AFD_SEND_INFO_UDP))

{

ULONG i = 0;

if (Info->BufferArray != NULL)

{

for (i = 0; i<Info->BufferCount; i++)

{

if (DataSize ==  MaxTranSize.MaxTranLength[FLAGS_SEND])

{

Info->BufferArray[i].len = 0;

}

else if (DataSize+Info->BufferArray[i].len

> MaxTranSize.MaxTranLength[FLAGS_SEND])

{

Info->BufferArray[i].len = (UINT)

MaxTranSize.MaxTranLength[FLAGS_SEND]- DataSize;

DataSize = (UINT)MaxTranSize.MaxTranLength[FLAGS_SEND];

break;

}

else

{

DataSize += Info->BufferArray[i].len;

}

}

}

}

}

break;

case IOCTL_AFD_RECV:

{

PAFD_RECV_INFO Info = (PAFD_RECV_INFO)InputBuffer;

ULONG DataSize = 0;

PROCESS_MAX_TRAN_SIZE MaxTranSize;

if (!FlwGetProcessMaxTranSize(&MaxTranSize,PsGetCurrentProcessId()))

{

break;

}

else if (MaxTranSize.MaxTranLength[FLAGS_RECV] < MIN_FLOW_DATA/2)

{

break;

}

if (Info != NULL &&

InputBufferLength >= sizeof(PAFD_RECV_INFO))

{

ULONG i = 0;

if (Info->BufferArray != NULL)

{

for (i = 0; i<Info->BufferCount; i++)

{

if (DataSize == MaxTranSize.MaxTranLength[FLAGS_RECV])

{

Info->BufferArray[i].len = 0;

}

else if (DataSize+Info->BufferArray[i].len

> MaxTranSize.MaxTranLength[FLAGS_RECV])

{

Info->BufferArray[i].len = (UINT)

MaxTranSize.MaxTranLength[FLAGS_RECV]- DataSize;

DataSize = (UINT)MaxTranSize.MaxTranLength[FLAGS_RECV];

}

else

{

DataSize += Info->BufferArray[i].len;

}

}

}

}

}

break;

case IOCTL_AFD_RECV_DATAGRAM:

{

PAFD_RECV_INFO_UDP Info = (PAFD_RECV_INFO_UDP)InputBuffer;

ULONG DataSize = 0;

PROCESS_MAX_TRAN_SIZE MaxTranSize;

if (!FlwGetProcessMaxTranSize(&MaxTranSize,PsGetCurrentProcessId()))

{

break;

}

else if (MaxTranSize.MaxTranLength[FLAGS_RECV] < MIN_FLOW_DATA/2)

{

break;

}

if (Info != NULL &&

InputBufferLength >= sizeof(PAFD_RECV_INFO_UDP))

{

ULONG i = 0;

if (Info->BufferArray != NULL)

{

for (i = 0; i<Info->BufferCount; i++)

{

if (DataSize ==  MaxTranSize.MaxTranLength[FLAGS_RECV])

{

Info->BufferArray[i].len = 0;

}

else if (DataSize+Info->BufferArray[i].len

> MaxTranSize.MaxTranLength[FLAGS_RECV])

{

Info->BufferArray[i].len = (UINT)

MaxTranSize.MaxTranLength[FLAGS_RECV]- DataSize;

DataSize = (UINT)MaxTranSize.MaxTranLength[FLAGS_RECV];

}

else

{

DataSize += Info->BufferArray[i].len;

}

}

}

}

}

break;

default:

break;

}

InterlockedDecrement(&HookCallNumber);

}

}

4.4.      后处理

后处理比较简单当请求完成后发送或接收了多少长度的数据会反应在IoStatusBlock参数的information中。因此我们只要记录这个值就可以了。具体的流量统计及暂停操作在FlwCtrlProcessFlowForSize中完成。

VOID

DeviceIoControlFileProcessPost(

IN HANDLE  FileHandle,

IN HANDLE  Event,

IN PIO_APC_ROUTINE  ApcRoutine,

IN PVOID  ApcContext,

OUT PIO_STATUS_BLOCK  IoStatusBlock,

IN ULONG  IoControlCode,

IN PVOID  InputBuffer,

IN ULONG  InputBufferLength,

OUT PVOID  OutputBuffer,

IN ULONG  OutputBufferLength

)

{

if (NT_SUCCESS(IoStatusBlock->Status) &&

UnloadFlags == FALSE)

{

InterlockedIncrement(&HookCallNumber);

switch(IoControlCode)

{

case IOCTL_AFD_SEND:

case IOCTL_AFD_SEND_DATAGRAM:

{

ULONG DataSize = 0;

DataSize = IoStatusBlock->Information;

if (DataSize > 0)

{

FlwCtrlProcessFlowForSize(DataSize,

FLAGS_SEND,

PsGetCurrentProcessId());

}

}

break;

case IOCTL_AFD_RECV:

case IOCTL_AFD_RECV_DATAGRAM:

{

ULONG DataSize = 0;

DataSize = IoStatusBlock->Information;

if (DataSize > 0)

{

FlwCtrlProcessFlowForSize(DataSize,

FLAGS_RECV,

PsGetCurrentProcessId());

}

}

break;

default:

break;

}

InterlockedDecrement(&HookCallNumber);

}

}

4.5.      其它辅助操作

4.5.1.         最大传输值计算

这里限制对进程每秒流量可控制的最小峰值，如果设定的进程流量峰值小于程序中可控制的最小峰值最大传输值统一为峰值的一半。MAX_TRAN_SIZE为64KB，如果设定的进程每秒流量峰值大于这个长度，统一的最大传输值为64KB。

BOOLEAN

FlwGetProcessMaxTranSize(LPPROCESS_MAX_TRAN_SIZE MaxTranSize,

HANDLE ProcessId)

{

FLOW_POLICY_ENTRY Policy;

if (FlwGetProcessPolicy(&Policy,ProcessId))

{

if (Policy.BytesRate[FLAGS_SEND] > MIN_FLOW_DATA)

{

MaxTranSize->MaxTranLength[FLAGS_SEND] =

Policy.BytesRate[FLAGS_SEND]>MAX_TRAN_SIZE?MAX_TRAN_SIZE/2:MIN_TRAN_SIZE;

}

else

{

MaxTranSize->MaxTranLength[FLAGS_SEND] = Policy.BytesRate[FLAGS_SEND]/2;

}

if (Policy.BytesRate[FLAGS_RECV] > MIN_FLOW_DATA)

{

MaxTranSize->MaxTranLength[FLAGS_RECV] =

Policy.BytesRate[FLAGS_RECV]>MAX_TRAN_SIZE?MAX_TRAN_SIZE:MIN_TRAN_SIZE;

}

else

{

MaxTranSize->MaxTranLength[FLAGS_RECV] = Policy.BytesRate[FLAGS_RECV]/2;

}

return TRUE;

}

return FALSE;

}

4.5.2.         流量统计

在这个函数里统计了进程每秒中发送和接收的数据流量及总共发送和接收的数据流量，同时这个函数还返回上一秒到本次统计流量经历的毫秒数。这个返回值是在计算需要暂停的时间时有用。

ULONG64

TkFlowStatForSizeEx(LONG    PacketSize,

LPFLOW_DATA pFlow,

BOOLEAN     isSend,

ULONG64     LastTime)

{

ULONG64 nSecond;

ULONG64 nLimit;

nSecond = TkGetCurrMillisecond();

if (LastTime == 0)

{

nLimit = nSecond - pFlow->LastTime[isSend];

}

else

{

nLimit = nSecond - LastTime;

nSecond = LastTime;

}

pFlow->SumBytes[isSend] += PacketSize;

if (nLimit < 1000)

{

pFlow->LastBytes[isSend] += PacketSize;

pFlow->BytesRate[isSend] = pFlow->LastBytes[isSend];

}

else

{

pFlow->BytesRate[isSend] = (pFlow->LastBytes[isSend]+PacketSize)/((nLimit+1000)/1000);

pFlow->LastTime[isSend] =  nSecond;

pFlow->LastBytes[isSend] = PacketSize;

nLimit = 0;

}

return nLimit;

}

4.5.3.         策略应用

基本的流量是先根据当前进程Id查出进程的流量策略和流量数据的存储条目，计算统计流量。如果超标就计算一个需要暂停的毫秒数。计算方法为：（统计出的流量*1000）/流量峰值-本次发送距离上一秒的毫秒数（即TkFlowStatForSizeEx的返回值）。

BOOLEAN

FlwCtrlProcessFlowForSize(LONG Data_Size,

BOOLEAN IsSend,

HANDLE ProcessId)

{

if (ProcessId == (HANDLE)-1)

{

ProcessId = PsGetCurrentProcessId();

}

return FlwCtrlProcessFlow(Data_Size,IsSend,ProcessId);

}

BOOLEAN

_Wait(LONG64 millisecond)

{

LARGE_INTEGER  TimeOut;

LONG64 startTime = 0;

NTSTATUS Status;

do

{

startTime = TkGetCurrMillisecond();

TimeOut.QuadPart = (LONGLONG)(millisecond * (-10000));

Status = KeDelayExecutionThread(KernelMode,FALSE,&TimeOut);

startTime = TkGetCurrMillisecond() - startTime;

if (startTime >= millisecond ||

UnloadFlags)

{

break;

}

millisecond -= startTime;

} while(TRUE);

return TRUE;

}

BOOLEAN

FlwCtrlProcessFlow(ULONG DataSize,

BOOLEAN IsSend,

HANDLE ProcessId)

{

KIRQL    oldIRQL;

ULONG    i = 0;

BOOLEAN  bRet = FALSE;

FW_PROCESS_LOAD_INFO   ProcessInfo;

FLOW_POLICY_ENTRY      Policy;

ULONG64 LastTime =  0;

ULONG64 nLimit = 0;

int loop  = 3;

do

{

KeAcquireSpinLock( &g_FwProcessFlowPolicyLock, &oldIRQL );

if (NULL != g_pFwProcessFlowPolicy &&

ProcGetInfoForArray(&ProcessInfo,ProcessId))

{

for (i=0;i<g_pFwProcessFlowPolicy->count;i++)

{

if (!_wcsnicmp(ProcessInfo.ProcessInfo.wProcFullPath,

g_pFwProcessFlowPolicy->Entry[i].wProcFullPath,

MAX_PATH))

{

Policy = g_pFwProcessFlowPolicy->Entry[i];

bRet = TRUE;

break;

}

}

}

KeReleaseSpinLock( &g_FwProcessFlowPolicyLock, oldIRQL );

if (bRet)

{

nLimit = TkFlowStatForSizeEx(DataSize,

&ProcessInfo.FlowData,IsSend,LastTime);

if (ProcessInfo.FlowData.LastBytes[IsSend] > Policy.BytesRate[IsSend] &&

Policy.BytesRate[IsSend] >= MIN_FLOW_DATA)

{

if (!UnloadFlags && KeGetCurrentIrql() <= APC_LEVEL)

{

nLimit = ((ProcessInfo.FlowData.LastBytes[IsSend]*1000)/

Policy.BytesRate[IsSend])-nLimit;

_Wait((LONG64)nLimit);

loop --;

}

else

{

break;

}

}

else

{

break;

}

}

else

{

break;

}

} while(loop>0);

// 返回是否控制成功

FlwStatisticProcFlow(DataSize,

IsSend,

ProcessId,

LastTime);

return bRet;

}

5.   总结

这种方法比较简单但是测试中一个问题比较难解决就是一个程序有多个收发线程时比较难处理，上面的阐述中并没有处理这种情况。同时我们为了缩小暂停的时间在请求转发前对数据做了预处理，这样如果限定的流量太小而进程又没有检查实际发送和接收的数据长度时就会让程序表现出异常的行为，例如在对飞秋的测试中会出现截断消息内容的情况。当然这些都可以通过进一步的改进去解决，这里只是阐述了一种方向。还有一个方法是通过TDI来获得进程与端口的对应关系再通过联合NDIS中间层驱动方式来控制进程流量，但在测试过程序中发现控制的粒度不是很好。最后，谢谢大家！希望本文能给大家带来一些帮助。

基于AFD驱动的进程流量控制,布布扣,bubuko.com

时间： 2024-10-08 15:18:59

基于AFD驱动的进程流量控制的相关文章

AFD驱动

首先来介绍一下AFD驱动.大家可能对NDIS及TDI了解的比较多,AFD驱动相对会少一些.NDIS及TDI是微软提供的内核网络驱动模型中的编程接口规范,而AFD驱动是微软Windows操作一个驱动部件.它是上层SOCKET在内核心中的实现.我这里画了一个图例来说明它们之间的关系: 简来说AFD驱动向上与SOCKET应用接口约定了接口来实现SOCKET,AFD驱动实际上是一个TDI客户端,它通过TDI接口调用微软件的另一个网络部件TCPIP驱动来完成功能.AFD并没有官方的资料说明它的接口,但是在

ACE框架 基于共享内存的进程间通讯

ACE框架将基于共享内存的进程间通讯功能,如其它IO组件或IPC组件一样,设计成三个组件.流操作组件ACE_MEM_Stream,连接器组件ACE_MEM_Connector,以及接收连接组件ACE_MEM_Accpter.ACE框架为基于共享内存的进程间通讯提供了两种数据传输(分发deliver)策略.一种是使用生产者-消费者队列的一对多的多用户MT策略,另一种是使用socket流的可以使用反应器响应数据接收事件的Reactor策略.不论哪一种策略都要通过socket进行TCP连接,并进行进程

使用 Spring 2.5 基于注解驱动的 Spring MVC

使用 Spring 2.5 基于注解驱动的 Spring MVC 原文链接:http://www.ibm.com/developerworks/cn/java/j-lo-spring25-mvc/ 基于注解的配置有越来越流行的趋势,Spring 2.5 顺应这种趋势,为 Spring MVC 提供了完全基于注解的配置.本文将介绍 Spring 2.5 新增的 Sping MVC 注解功能,讲述如何使用注解配置替换传统的基于 XML 的 Spring MVC 配置. 概述 继 Spring 2.0

新年第一弹：基于领域驱动的Java开发工具包Common项目分享

项目地址:https://github.com/xuliugen/common 喜欢的请加星. 先上图: 本项目主要包含:constant.domain.exception.util这四个主要部分. Constant 主要包含了一些项目中的常用常量. 示例: (1)标点符号常量类 public interface ConstPunctuation { /** * : 冒号 */ String COLON = ":"; /** * - 中划线 */ String MINUS = &quo

tornado项目之基于领域驱动模型架构设计的京东用户管理后台

本博文将一步步揭秘京东等大型网站的领域驱动模型,致力于让读者完全掌握这种网络架构中的“高富帅”. 一.预备知识: 1.接口: python中并没有类似java等其它语言中的接口类型,但是python中有抽象类和抽象方法.如果一个抽象类有抽象方法,那么继承它的子类必须实现抽象类的所有方法,因此,我们基于python的抽象类和抽象方法实现接口功能. 示例代码: from abc import ABCMeta from abc import abstractmethod #导入抽象方法 class F

使用 Spring 2.5 基于注解驱动的 Spring MVC--转

概述 继 Spring 2.0 对 Spring MVC 进行重大升级后,Spring 2.5 又为 Spring MVC 引入了注解驱动功能.现在你无须让 Controller 继承任何接口,无需在 XML 配置文件中定义请求和 Controller 的映射关系,仅仅使用注解就可以让一个 POJO 具有 Controller 的绝大部分功能 —— Spring MVC 框架的易用性得到了进一步的增强.在框架灵活性.易用性和扩展性上,Spring MVC 已经全面超越了其它的 MVC 框架,伴随

【tornado】系列项目（二）基于领域驱动模型的区域后台管理+前端easyui实现

本项目是一个系列项目,最终的目的是开发出一个类似京东商城的网站.本文主要介绍后台管理中的区域管理,以及前端基于easyui插件的使用.本次增删改查因数据量少,因此采用模态对话框方式进行,关于数据量大采用跳转方式修改,详见博主后续博文. 后台界面展示: 地区管理包含省市县的管理.详见下文. 一.数据库设计 ? 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 class Province(Base):  

在.net下打造mongoDb基于官方驱动最新版本

还是一如既往先把结构图放出来,上上个版本添加了redis的缓存,但是不满足我的需求,因为公司有项目要求是分布式所以呢,这里我就增加了mongoDb进行缓存分布式,好了先看结构图. 总的来说比较蛋疼,因为从来没有使用过mongoDB,从安装,到转为windows服务,设置权限等等,好吧这都是题外话. 在写这个MongoDB版本的时候遇到的一些问题,我先总结下: 1.MongoDb版本是官网最新版3.4.4,官方驱动为2.4.3,首先我的项目是以GUID做为主键,在往MongonDB中插入时遇到的是

Windows Minifilter驱动 - 获取进程ID, 进程名字和线程ID （5）

在minifilter里面可能有好几种获取调用进程id,名字和线程的办法.我这里有一种: 使用 PsSetCreateProcessNotifyRoutine 和 PsSetLoadImageNotifyRoutine 这是两个API,我们可以借助它们获取进程信息.具体看:http://msdn.microsoft.com/en-us/library/windows/hardware/ff559951(v=vs.85).aspx PsSetLoadImageNotifyRoutine 可以使用这