
#pragma once
#include "afxmt.h"

using namespace std;

static const UINT64 c_uiFilterWeight = 0xFFFFFFFFFFFFFFFE;
static const UINT16 c_uiSublayerWeight = 0xFFFE;

enum TEnTimerType
{
    enIpV4SingleAddr,
    enIpV4Subnet,
    enIpV6SingleAddr,
    enExternalClientBlocked
};
struct TTimerDescriptor
{
    TEnTimerType        enTimerType;
    UINT32              ulIpV4SingleAddr;
    IP4SUBNET_STRUCT    stIpV4Subnet;
    FWP_BYTE_ARRAY16    stIpV6SingleAddr;
    int                 nIpV4SingleAddrIndex;
    int                 nIpV6SingleAddrIndex;
    int                 nIpV4SubnetIndex;
};
typedef TTimerDescriptor* PTTimerDescriptor;
typedef unordered_map<HANDLE, PTTimerDescriptor> TTimerArray;
typedef pair<TTimerArray::iterator, bool> TTimerArrayInsertPair;
typedef TTimerArray::value_type TTimerArrayValueType;

// CDnsStateThread
class CDnsMessageWB;
class CDnsFirewallThread;
class CDnsStateThread : public CWinThread
{
	DECLARE_DYNCREATE(CDnsStateThread)

public:
	CDnsStateThread();
	virtual ~CDnsStateThread();

public:
	virtual BOOL InitInstance();
	virtual int ExitInstance();

protected:
	DECLARE_MESSAGE_MAP()
public:
    FWP_BYTE_BLOB* m_pDnsServerAppID;
    int m_nRetCode;
    HANDLE m_hWfpEngineHandle;
    UINT32 m_nInboundTransportV4Id;
    UINT32 m_nInboundTransportV6Id;
    UINT32 m_nAleAuthRecvAcceptV4CalloutId;
    UINT32 m_nAleAuthRecvAcceptV6CalloutId;
    UINT32 m_nAleFlowEstablishedV4CalloutId;
    UINT32 m_nAleFlowEstablishedV6CalloutId;
    UINT32 m_nStreamV4CalloutId;
    UINT32 m_nStreamV6CalloutId;
    UINT32 m_nDatagramsV4CalloutId;
    UINT32 m_nDatagramsV6CalloutId;
    UINT64 m_nInboundTransportV4FilterId;
    UINT64 m_nInboundTransportV6FilterId;
    UINT64 m_nAleAuthRecvAcceptV4FilterID;
    UINT64 m_nAleAuthRecvAcceptV6FilterID;
    UINT64 m_nAleFlowEstablishedV4FilterID;
    UINT64 m_nAleFlowEstablishedV6FilterID;
    UINT64 m_nStreamV4FilterID;
    UINT64 m_nStreamV6FilterID;
    UINT64 m_nDatagramV4DataFilterID;
    UINT64 m_nDatagramV6DataFilterID;
    HANDLE m_hDnsFirewallDevice;
    FLOW_DATA m_aFowContextArray[DNSFIREWALL_FILE_COUNT];
    virtual BOOL OnIdle(LONG lCount);
    TTimerArray m_oTimerArray;
    void ResetFirewallTimer(
        PTTimerDescriptor pTimerDescriptor
        );
    bool m_bFireWallInitialized;
    CDnsFirewallThread* m_aDnsFirewallThreadsArray[DNSFIREWALL_THREAD_COUNT];
    static void CALLBACK sfTimerProc(
        _In_ PVOID   lpParameter,
        _In_ BOOLEAN TimerOrWaitFired
        );
    void mfTimerProc(
        PTTimerDescriptor pTimerDescriptor
        );
    CCriticalSection m_oTimersLock;
    void ResetFirewallIpV4Config(UINT32 ulIpAddress);
    void ResetFirewallIpV6Config(FWP_BYTE_ARRAY16& rRemoteAddressV6);
    HANDLE m_hTimerQueue;
    FIREWALL_CONFIG m_stFirewallConfig;
    bool m_bServiceRebootNeed;
    CCriticalSection m_oThreadLock;
    CEvent m_oDnsSvcRestartNeedEvent;
    bool m_bDriverStarted;
    OVERLAPPED m_stWaitRestartOverlapped;
    CEvent m_oRestartQueryEvent;
    DWORD m_dwRestartStatus;
};
