基于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 */

#defineFSCTL_AFD_BASE                 FILE_DEVICE_NETWORK

#define_AFD_CONTROL_CODE(Operation,Method) /

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

 

/* AFD Commands */

#defineAFD_BIND            0

#defineAFD_CONNECT              1

#defineAFD_START_LISTEN         2

#defineAFD_WAIT_FOR_LISTEN      3

#defineAFD_ACCEPT          4

#defineAFD_RECV            5

#defineAFD_RECV_DATAGRAM        6

#defineAFD_SEND            7

#defineAFD_SEND_DATAGRAM        8

#defineAFD_SELECT          9

#defineAFD_DISCONNECT           10

#defineAFD_GET_SOCK_NAME        11

#defineAFD_GET_PEER_NAME              12

#defineAFD_GET_TDI_HANDLES      13

#defineAFD_SET_INFO             14

#defineAFD_GET_CONTEXT          16

#defineAFD_SET_CONTEXT          17

#defineAFD_SET_CONNECT_DATA        18

#defineAFD_SET_CONNECT_OPTIONS     19

#defineAFD_SET_DISCONNECT_DATA     20

#defineAFD_SET_DISCONNECT_OPTIONS  21

#defineAFD_GET_CONNECT_DATA        22

#defineAFD_GET_CONNECT_OPTIONS     23

#defineAFD_GET_DISCONNECT_DATA     24

#defineAFD_GET_DISCONNECT_OPTIONS  25

#defineAFD_SET_CONNECT_DATA_SIZE      26

#defineAFD_SET_CONNECT_OPTIONS_SIZE   27

#defineAFD_SET_DISCONNECT_DATA_SIZE   28

#defineAFD_SET_DISCONNECT_OPTIONS_SIZE 29

#defineAFD_GET_INFO             30

#defineAFD_EVENT_SELECT         33

#defineAFD_ENUM_NETWORK_EVENTS        34

#defineAFD_DEFER_ACCEPT         35

#defineAFD_GET_PENDING_CONNECT_DATA41

 

/* AFD IOCTLs */

 

#defineIOCTL_AFD_BIND /

 _AFD_CONTROL_CODE(AFD_BIND,METHOD_NEITHER)

#defineIOCTL_AFD_CONNECT /

 _AFD_CONTROL_CODE(AFD_CONNECT,METHOD_NEITHER)

#defineIOCTL_AFD_START_LISTEN /

 _AFD_CONTROL_CODE(AFD_START_LISTEN,METHOD_NEITHER)

#defineIOCTL_AFD_WAIT_FOR_LISTEN /

 _AFD_CONTROL_CODE(AFD_WAIT_FOR_LISTEN,METHOD_BUFFERED )

#defineIOCTL_AFD_ACCEPT /

 _AFD_CONTROL_CODE(AFD_ACCEPT,METHOD_BUFFERED )

#defineIOCTL_AFD_RECV /

 _AFD_CONTROL_CODE(AFD_RECV,METHOD_NEITHER)

#defineIOCTL_AFD_RECV_DATAGRAM /

 _AFD_CONTROL_CODE(AFD_RECV_DATAGRAM,METHOD_NEITHER)

#defineIOCTL_AFD_SEND /

 _AFD_CONTROL_CODE(AFD_SEND,METHOD_NEITHER)

#defineIOCTL_AFD_SEND_DATAGRAM /

 _AFD_CONTROL_CODE(AFD_SEND_DATAGRAM,METHOD_NEITHER)

#defineIOCTL_AFD_SELECT /

 _AFD_CONTROL_CODE(AFD_SELECT,METHOD_BUFFERED )

#defineIOCTL_AFD_DISCONNECT /

 _AFD_CONTROL_CODE(AFD_DISCONNECT,METHOD_NEITHER)

#defineIOCTL_AFD_GET_SOCK_NAME /

 _AFD_CONTROL_CODE(AFD_GET_SOCK_NAME,METHOD_NEITHER)

#defineIOCTL_AFD_GET_PEER_NAME /

 _AFD_CONTROL_CODE(AFD_GET_PEER_NAME,METHOD_NEITHER)

#defineIOCTL_AFD_GET_TDI_HANDLES /

 _AFD_CONTROL_CODE(AFD_GET_TDI_HANDLES,METHOD_NEITHER)

#defineIOCTL_AFD_SET_INFO /

 _AFD_CONTROL_CODE(AFD_SET_INFO,METHOD_NEITHER)

#defineIOCTL_AFD_GET_CONTEXT /

 _AFD_CONTROL_CODE(AFD_GET_CONTEXT,METHOD_NEITHER)

#defineIOCTL_AFD_SET_CONTEXT /

 _AFD_CONTROL_CODE(AFD_SET_CONTEXT,METHOD_NEITHER)

#defineIOCTL_AFD_SET_CONNECT_DATA /

 _AFD_CONTROL_CODE(AFD_SET_CONNECT_DATA,METHOD_NEITHER)

#defineIOCTL_AFD_SET_CONNECT_OPTIONS /

 _AFD_CONTROL_CODE(AFD_SET_CONNECT_OPTIONS,METHOD_NEITHER)

#defineIOCTL_AFD_SET_DISCONNECT_DATA /

 _AFD_CONTROL_CODE(AFD_SET_DISCONNECT_DATA,METHOD_NEITHER)

#defineIOCTL_AFD_SET_DISCONNECT_OPTIONS /

 _AFD_CONTROL_CODE(AFD_SET_DISCONNECT_OPTIONS,METHOD_NEITHER)

#defineIOCTL_AFD_GET_CONNECT_DATA /

 _AFD_CONTROL_CODE(AFD_GET_CONNECT_DATA,METHOD_NEITHER)

#defineIOCTL_AFD_GET_CONNECT_OPTIONS /

 _AFD_CONTROL_CODE(AFD_GET_CONNECT_OPTIONS,METHOD_NEITHER)

#defineIOCTL_AFD_GET_DISCONNECT_DATA /

 _AFD_CONTROL_CODE(AFD_GET_DISCONNECT_DATA,METHOD_NEITHER)

#defineIOCTL_AFD_GET_DISCONNECT_OPTIONS /

 _AFD_CONTROL_CODE(AFD_GET_DISCONNECT_OPTIONS,METHOD_NEITHER)

#defineIOCTL_AFD_SET_CONNECT_DATA_SIZE /

 _AFD_CONTROL_CODE(AFD_SET_CONNECT_DATA_SIZE,METHOD_NEITHER)

#defineIOCTL_AFD_SET_CONNECT_OPTIONS_SIZE /

 _AFD_CONTROL_CODE(AFD_SET_CONNECT_OPTIONS_SIZE,METHOD_NEITHER)

#defineIOCTL_AFD_SET_DISCONNECT_DATA_SIZE /

 _AFD_CONTROL_CODE(AFD_SET_DISCONNECT_DATA_SIZE,METHOD_NEITHER)

#defineIOCTL_AFD_SET_DISCONNECT_OPTIONS_SIZE /

 _AFD_CONTROL_CODE(AFD_SET_DISCONNECT_OPTIONS_SIZE,METHOD_NEITHER)

#defineIOCTL_AFD_GET_INFO /

 _AFD_CONTROL_CODE(AFD_GET_INFO,METHOD_NEITHER)

#defineIOCTL_AFD_EVENT_SELECT /

 _AFD_CONTROL_CODE(AFD_EVENT_SELECT,METHOD_NEITHER)

#defineIOCTL_AFD_DEFER_ACCEPT /

 _AFD_CONTROL_CODE(AFD_DEFER_ACCEPT,METHOD_NEITHER)

#defineIOCTL_AFD_GET_PENDING_CONNECT_DATA /

 _AFD_CONTROL_CODE(AFD_GET_PENDING_CONNECT_DATA,METHOD_NEITHER)

#defineIOCTL_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控制码进行数据发送和接收时还涉及一些数据结构。这些在后面介绍流量控制时会用到。先在这里说明一下。

typedefstruct _AFD_MAPBUF {

    PVOID BufferAddress;

    PMDL  Mdl;

}AFD_MAPBUF, *PAFD_MAPBUF;

 

typedefstruct _AFD_WSABUF {

    UINT  len;

    PCHAR buf;

}AFD_WSABUF, *PAFD_WSABUF;

 

typedefstruct  _AFD_RECV_INFO {

    PAFD_WSABUF                BufferArray;

    ULONG                      BufferCount;

    ULONG                      AfdFlags;

    ULONG                      TdiFlags;

}AFD_RECV_INFO , *PAFD_RECV_INFO ;

 

typedefstruct _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;

 

typedefstruct  _AFD_SEND_INFO {

    PAFD_WSABUF                 BufferArray;

    ULONG                      BufferCount;

    ULONG                      AfdFlags;

    ULONG                      TdiFlags;

}AFD_SEND_INFO , *PAFD_SEND_INFO ;

 

typedefstruct _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,

                        OUTPVOID  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(__inPDRIVER_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_SIZEMaxTranSize,

                            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_DATApFlow,

                    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(LONGData_Size,

                         BOOLEAN IsSend,

                         HANDLE ProcessId)

{

    if (ProcessId == (HANDLE)-1)

    {

        ProcessId = PsGetCurrentProcessId();

    }

    return FlwCtrlProcessFlow(Data_Size,IsSend,ProcessId);

}

BOOLEAN

_Wait(LONG64millisecond)

{

    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(ULONGDataSize,

                    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中间层驱动方式来控制进程流量，但在测试过程序中发现控制的粒度不是很好。最后，谢谢大家！希望本文能给大家带来一些帮助。

文章来源：http://blog.csdn.net/lziog/article/details/6530501