Refactor Test Datapath Hooks (#509)

2020-06-16 09:28:33 -07:00 · 2020-06-16 09:28:33 -07:00 · 7408bbc91d
--- a/src/core/binding.c
+++ b/src/core/binding.c
@ -1421,14 +1421,15 @@ QuicBindingReceive(
        Packet->Buffer = Datagram->Buffer;
        Packet->BufferLength = Datagram->BufferLength;

-#if DEBUG
+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
        //
        // The test datapath receive callback allows for test code to modify
        // the datagrams on the receive path, and optionally indicate one or
        // more to be dropped.
        //
-        if (MsQuicLib.TestDatapathHooks != NULL) {
-            if (MsQuicLib.TestDatapathHooks->Receive(Datagram)) {
+        QUIC_TEST_DATAPATH_HOOKS* Hooks = MsQuicLib.TestDatapathHooks;
+        if (Hooks != NULL) {
+            if (Hooks->Receive(Datagram)) {
                *ReleaseChainTail = Datagram;
                ReleaseChainTail = &Datagram->Next;
                QuicPacketLogDrop(Binding, Packet, "Test Dopped");
@ -1551,12 +1552,13 @@ QuicBindingSendTo(
 {
    QUIC_STATUS Status;

-#if DEBUG
-    if (MsQuicLib.TestDatapathHooks != NULL) {
+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
+    QUIC_TEST_DATAPATH_HOOKS* Hooks = MsQuicLib.TestDatapathHooks;
+    if (Hooks != NULL) {

        QUIC_ADDR RemoteAddressCopy = *RemoteAddress;
        BOOLEAN Drop =
-            MsQuicLib.TestDatapathHooks->Send(
+            Hooks->Send(
                &RemoteAddressCopy,
                NULL,
                SendContext);
@ -1596,7 +1598,7 @@ QuicBindingSendTo(
                Binding,
                Status);
        }
-#if DEBUG
+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
    }
 #endif

@ -1614,13 +1616,14 @@ QuicBindingSendFromTo(
 {
    QUIC_STATUS Status;

-#if DEBUG
-    if (MsQuicLib.TestDatapathHooks != NULL) {
+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
+    QUIC_TEST_DATAPATH_HOOKS* Hooks = MsQuicLib.TestDatapathHooks;
+    if (Hooks != NULL) {

        QUIC_ADDR RemoteAddressCopy = *RemoteAddress;
        QUIC_ADDR LocalAddressCopy = *LocalAddress;
        BOOLEAN Drop =
-            MsQuicLib.TestDatapathHooks->Send(
+            Hooks->Send(
                &RemoteAddressCopy,
                &LocalAddressCopy,
                SendContext);
@ -1662,7 +1665,7 @@ QuicBindingSendFromTo(
                Binding,
                Status);
        }
-#if DEBUG
+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
    }
 #endif

--- a/src/core/connection.c
+++ b/src/core/connection.c
@ -4755,6 +4755,13 @@ QuicConnRecvPostProcessing(
        QuicPathSetActive(Connection, *Path);
        *Path = &Connection->Paths[0];

+        QuicTraceEvent(
+            ConnRemoteAddrAdded,
+            "[conn][%p] New Remote IP: %!SOCKADDR!",
+            Connection,
+            LOG_ADDR_LEN(Connection->Paths[0].RemoteAddress),
+            (const uint8_t*)&Connection->Paths[0].RemoteAddress); // TODO - Addr removed event?
+
        QUIC_CONNECTION_EVENT Event;
        Event.Type = QUIC_CONNECTION_EVENT_PEER_ADDRESS_CHANGED;
        Event.PEER_ADDRESS_CHANGED.Address = &(*Path)->RemoteAddress;
--- a/src/core/library.c
+++ b/src/core/library.c
@ -572,7 +572,7 @@ QuicLibrarySetGlobalParam(
        Status = QUIC_STATUS_SUCCESS;
        break;

-#if DEBUG
+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
    case QUIC_PARAM_GLOBAL_TEST_DATAPATH_HOOKS:

        if (BufferLength != sizeof(QUIC_TEST_DATAPATH_HOOKS*)) {
@ -580,14 +580,6 @@ QuicLibrarySetGlobalParam(
            break;
        }

-        if (MsQuicLib.InUse) {
-            QuicTraceLogError(
-                LibraryTestDatapathHooksSetAfterInUse,
-                "[ lib] Tried to change test datapath hooks after library in use!");
-            Status = QUIC_STATUS_INVALID_STATE;
-            break;
-        }
-
        MsQuicLib.TestDatapathHooks = *(QUIC_TEST_DATAPATH_HOOKS**)Buffer;
        QuicTraceLogWarning(
            LibraryTestDatapathHooksSet,
--- a/src/core/library.h
+++ b/src/core/library.h
@ -207,7 +207,7 @@ typedef struct QUIC_LIBRARY {
    //
    QUIC_TOEPLITZ_HASH ToeplitzHash;

-#if DEBUG
+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
    //
    // An optional callback to allow test code to modify the data path.
    //
--- a/src/core/quicdef.h
+++ b/src/core/quicdef.h
@ -307,7 +307,7 @@ QUIC_STATIC_ASSERT(
 // The minimum number of bytes of send allowance we must have before we will
 // send another packet.
 //
-#define QUIC_MIN_SEND_ALLOWANCE                 100
+#define QUIC_MIN_SEND_ALLOWANCE                 75

 //
 // The minimum buffer space that we require before we will pack another
--- a/src/inc/msquicp.h
+++ b/src/inc/msquicp.h
@ -57,6 +57,14 @@ typedef struct QUIC_TEST_DATAPATH_HOOKS {
    QUIC_TEST_DATAPATH_SEND_HOOK Send;
 } QUIC_TEST_DATAPATH_HOOKS;

+#if DEBUG
+//
+// Datapath hooks are currently only enabled on debug builds for functional
+// testing helpers.
+//
+#define QUIC_TEST_DATAPATH_HOOKS_ENABLED 1
+#endif
+
 #define QUIC_PARAM_GLOBAL_ENCRYPTION                    0x80000001  // uint8_t (BOOLEAN)
 #define QUIC_PARAM_GLOBAL_TEST_DATAPATH_HOOKS           0x80000002  // QUIC_TEST_DATAPATH_HOOKS*

--- a/src/test/MsQuicTests.h
+++ b/src/test/MsQuicTests.h
@ -21,6 +21,9 @@ extern QUIC_SEC_CONFIG* SecurityConfig;
 extern "C" {
 #endif

+void QuicTestInitialize();
+void QuicTestUninitialize();
+
 //
 // Parameter Validation Tests
 //
--- a/src/test/bin/quic_gtest.cpp
+++ b/src/test/bin/quic_gtest.cpp
@ -39,6 +39,7 @@ public:
            const QUIC_REGISTRATION_CONFIG RegConfig = { "MsQuicBVT", QUIC_EXECUTION_PROFILE_LOW_LATENCY };
            ASSERT_TRUE(QUIC_SUCCEEDED(MsQuic->RegistrationOpen(&RegConfig, &Registration)));
            ASSERT_TRUE(LoadSecConfig());
+            QuicTestInitialize();
        }
    }
    void TearDown() override {
@ -46,6 +47,7 @@ public:
            DriverClient.Uninitialize();
            DriverService.Uninitialize();
        } else {
+            QuicTestUninitialize();
            MsQuic->SecConfigDelete(SecurityConfig);
            MsQuic->RegistrationClose(Registration);
            MsQuicClose(MsQuic);
@ -398,11 +400,11 @@ TEST_P(WithFamilyArgs, VersionNegotiation) {
    }
 }

+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
 TEST_P(WithFamilyArgs, Rebind) {
    TestLoggerT<ParamType> Logger("QuicTestConnect-Rebind", GetParam());
    if (TestingKernelMode) {
        GTEST_SKIP_(":Unsupported in kernel mode");
-        /* Not supported in kernel mode yet.
        QUIC_RUN_CONNECT_PARAMS Params = {
            GetParam().Family,
            0,  // ServerStatelessRetry
@ -415,7 +417,7 @@ TEST_P(WithFamilyArgs, Rebind) {
            0,  // SessionResumption
            0   // RandomLossPercentage
        };
-        ASSERT_TRUE(DriverClient.Run(IOCTL_QUIC_RUN_CONNECT, Params));*/
+        ASSERT_TRUE(DriverClient.Run(IOCTL_QUIC_RUN_CONNECT, Params));
    } else {
        QuicTestConnect(
            GetParam().Family,
@ -430,6 +432,7 @@ TEST_P(WithFamilyArgs, Rebind) {
            0);     // RandomLossPercentage
    }
 }
+#endif

 TEST_P(WithFamilyArgs, ChangeMaxStreamIDs) {
    TestLoggerT<ParamType> Logger("QuicTestConnect-ChangeMaxStreamIDs", GetParam());
@ -493,6 +496,7 @@ TEST_P(WithHandshakeArgs3, AsyncSecurityConfig) {
    }
 }

+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
 TEST_P(WithHandshakeArgs4, RandomLoss) {
    TestLoggerT<ParamType> Logger("QuicTestConnect-RandomLoss", GetParam());
    if (TestingKernelMode) {
@ -523,6 +527,7 @@ TEST_P(WithHandshakeArgs4, RandomLoss) {
            GetParam().RandomLossPercentage);
    }
 }
+#endif

 TEST_P(WithFamilyArgs, Unreachable) {
    TestLoggerT<ParamType> Logger("QuicTestConnectUnreachable", GetParam());
--- a/src/test/bin/winkernel/driver.cpp
+++ b/src/test/bin/winkernel/driver.cpp
@ -10,6 +10,7 @@ Abstract:
 --*/

 #include <quic_platform.h>
+#include <MsQuicTests.h>

 #include "quic_trace.h"

@ -148,6 +149,8 @@ Return Value:
        goto Error;
    }

+    QuicTestInitialize();
+
    QuicTraceLogInfo(
        TestDriverStarted,
        "[test] Started");
@ -187,6 +190,8 @@ Arguments:
    UNREFERENCED_PARAMETER(Driver);
    NT_ASSERT(KeGetCurrentIrql() == PASSIVE_LEVEL);

+    QuicTestUninitialize();
+
    QuicTestCtlUninitialize();

    QuicTraceLogInfo(
--- a/src/test/lib/DatagramTest.cpp
+++ b/src/test/lib/DatagramTest.cpp
@ -165,74 +165,6 @@ QuicTestDatagramNegotiation(
    }
 }

-struct SelectiveLossHelper
-{
-    static uint32_t DropPacketCount;
-    static QUIC_TEST_DATAPATH_HOOKS DataPathFuncTable;
-    SelectiveLossHelper() {
-        QUIC_TEST_DATAPATH_HOOKS* Value = &DataPathFuncTable;
-        TEST_QUIC_SUCCEEDED(
-            MsQuic->SetParam(
-                nullptr,
-                QUIC_PARAM_LEVEL_GLOBAL,
-                QUIC_PARAM_GLOBAL_TEST_DATAPATH_HOOKS,
-                sizeof(Value),
-                &Value));
-    }
-    ~SelectiveLossHelper() {
-        QUIC_TEST_DATAPATH_HOOKS* Value = nullptr;
-        uint32_t TryCount = 0;
-        while (TryCount++ < 10) {
-            if (QUIC_SUCCEEDED(
-                MsQuic->SetParam(
-                    nullptr,
-                    QUIC_PARAM_LEVEL_GLOBAL,
-                    QUIC_PARAM_GLOBAL_TEST_DATAPATH_HOOKS,
-                    sizeof(Value),
-                    &Value))) {
-                break;
-            }
-            QuicSleep(100); // Let the current datapath queue drain.
-        }
-        if (TryCount == 10) {
-            TEST_FAILURE("Failed to disable test datapath hook");
-        }
-    }
-    void DropPackets(uint32_t Count) { DropPacketCount = Count; }
-    static
-    _IRQL_requires_max_(DISPATCH_LEVEL)
-    BOOLEAN
-    QUIC_API
-    ReceiveCallback(
-        _Inout_ struct QUIC_RECV_DATAGRAM* /* Datagram */
-        )
-    {
-        if (DropPacketCount == 0) {
-            return false;
-        }
-        DropPacketCount--;
-        return true;
-    }
-    static
-    _IRQL_requires_max_(PASSIVE_LEVEL)
-    BOOLEAN
-    QUIC_API
-    SendCallback(
-        _Inout_ QUIC_ADDR* /* RemoteAddress */,
-        _Inout_opt_ QUIC_ADDR* /* LocalAddress */,
-        _Inout_ struct QUIC_DATAPATH_SEND_CONTEXT* /* SendContext */
-        )
-    {
-        return FALSE; // Don't drop
-    }
-};
-
-uint32_t SelectiveLossHelper::DropPacketCount = false;
-QUIC_TEST_DATAPATH_HOOKS SelectiveLossHelper::DataPathFuncTable = {
-    SelectiveLossHelper::ReceiveCallback,
-    SelectiveLossHelper::SendCallback
-};
-
 void
 QuicTestDatagramSend(
    _In_ int Family
@ -309,6 +241,7 @@ QuicTestDatagramSend(

                TEST_EQUAL(1, Client.GetDatagramsAcknowledged());

+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
                LossHelper.DropPackets(1);

                TEST_QUIC_SUCCEEDED(
@ -326,6 +259,7 @@ QuicTestDatagramSend(
                QuicSleep(500);

                TEST_EQUAL(1, Client.GetDatagramsSuspectLost());
+#endif

                Client.Shutdown(QUIC_CONNECTION_SHUTDOWN_FLAG_NONE, QUIC_TEST_NO_ERROR);
                if (!Client.WaitForShutdownComplete()) {
@ -337,11 +271,6 @@ QuicTestDatagramSend(
                TEST_FALSE(Client.GetPeerClosed());
                TEST_FALSE(Client.GetTransportClosed());
            }
-
-#if !QUIC_SEND_FAKE_LOSS
-            TEST_TRUE(Server->GetPeerClosed());
-            TEST_EQUAL(Server->GetPeerCloseErrorCode(), QUIC_TEST_NO_ERROR);
-#endif
        }
    }
 }
--- a/src/test/lib/HandshakeTest.cpp
+++ b/src/test/lib/HandshakeTest.cpp
@ -11,11 +11,22 @@ Abstract:

 #include "precomp.h"

-uint8_t RandomLossHelper::LossPercentage = 0;
-QUIC_TEST_DATAPATH_HOOKS RandomLossHelper::DataPathFuncTable = {
-    RandomLossHelper::ReceiveCallback,
-    RandomLossHelper::SendCallback
+QUIC_TEST_DATAPATH_HOOKS DatapathHooks::FuncTable = {
+    DatapathHooks::ReceiveCallback,
+    DatapathHooks::SendCallback
 };
+DatapathHooks* DatapathHooks::Instance;
+
+void QuicTestInitialize()
+{
+    DatapathHooks::Instance = new DatapathHooks;
+}
+
+void QuicTestUninitialize()
+{
+    delete DatapathHooks::Instance;
+    DatapathHooks::Instance = nullptr;
+}

 _Function_class_(NEW_STREAM_CALLBACK)
 static
@ -227,11 +238,13 @@ QuicTestConnect(
                    Client.GetLocalBidiStreamCount());

                if (ClientRebind) {
-                    QuicAddr NewLocalAddr(QuicAddrFamily);
-                    TEST_QUIC_SUCCEEDED(Client.SetLocalAddr(NewLocalAddr));
+                    QuicAddr OrigLocalAddr;
+                    TEST_QUIC_SUCCEEDED(Client.GetLocalAddr(OrigLocalAddr));
+                    QuicAddr NewLocalAddr(OrigLocalAddr, 1);
                    QuicSleep(100);
-                    TEST_QUIC_SUCCEEDED(Client.GetLocalAddr(NewLocalAddr));
+                    ReplaceAddressHelper AddrHelper(OrigLocalAddr.SockAddr, NewLocalAddr.SockAddr);
                    TEST_FALSE(Client.GetIsShutdown());
+                    Client.SetKeepAlive(25);

                    bool ServerAddressUpdated = false;
                    uint32_t Try = 0;
--- a/src/test/lib/TestConnection.h
+++ b/src/test/lib/TestConnection.h
@ -133,7 +133,7 @@ public:
    uint32_t GetWaitTimeout() const {
        uint32_t WaitTime = TestWaitTimeout;
        if (HasRandomLoss) {
-            WaitTime *= 10; // TODO - Enough?
+            WaitTime *= 20; // TODO - Enough?
        }
        return WaitTime;
    }
--- a/src/test/lib/TestHelpers.h
+++ b/src/test/lib/TestHelpers.h
@ -365,66 +365,282 @@ struct PrivateTransportHelper : QUIC_PRIVATE_TRANSPORT_PARAMETER
    }
 };

-struct RandomLossHelper
+struct DatapathHook
 {
-    static uint8_t LossPercentage;
-    static QUIC_TEST_DATAPATH_HOOKS DataPathFuncTable;
-    RandomLossHelper(uint8_t _LossPercentage) {
-        LossPercentage = _LossPercentage;
-        if (LossPercentage != 0) {
-            QUIC_TEST_DATAPATH_HOOKS* Value = &DataPathFuncTable;
-            TEST_QUIC_SUCCEEDED(
-                MsQuic->SetParam(
-                    nullptr,
-                    QUIC_PARAM_LEVEL_GLOBAL,
-                    QUIC_PARAM_GLOBAL_TEST_DATAPATH_HOOKS,
-                    sizeof(Value),
-                    &Value));
-        }
+    DatapathHook* Next;
+
+    DatapathHook() : Next(nullptr) { }
+
+    virtual
+    _IRQL_requires_max_(DISPATCH_LEVEL)
+    BOOLEAN
+    Receive(
+        _Inout_ struct QUIC_RECV_DATAGRAM* /* Datagram */
+        ) {
+        return FALSE; // Don't drop by default
    }
-    ~RandomLossHelper() {
-        if (LossPercentage != 0) {
-            QUIC_TEST_DATAPATH_HOOKS* Value = nullptr;
-            uint32_t TryCount = 0;
-            while (TryCount++ < 10) {
-                if (QUIC_SUCCEEDED(
-                    MsQuic->SetParam(
-                        nullptr,
-                        QUIC_PARAM_LEVEL_GLOBAL,
-                        QUIC_PARAM_GLOBAL_TEST_DATAPATH_HOOKS,
-                        sizeof(Value),
-                        &Value))) {
-                    break;
-                }
-                QuicSleep(100); // Let the current datapath queue drain.
-            }
-            if (TryCount == 10) {
-                TEST_FAILURE("Failed to disable test datapath hook");
-            }
-        }
+
+    virtual
+    _IRQL_requires_max_(PASSIVE_LEVEL)
+    BOOLEAN
+    Send(
+        _Inout_ QUIC_ADDR* /* RemoteAddress */,
+        _Inout_opt_ QUIC_ADDR* /* LocalAddress */,
+        _Inout_ struct QUIC_DATAPATH_SEND_CONTEXT* /* SendContext */
+        ) {
+        return FALSE; // Don't drop by default
    }
+};
+
+class DatapathHooks
+{
+    static QUIC_TEST_DATAPATH_HOOKS FuncTable;
+
+    DatapathHook* Hooks;
+    QUIC_DISPATCH_LOCK Lock;
+
    static
    _IRQL_requires_max_(DISPATCH_LEVEL)
    BOOLEAN
    QUIC_API
    ReceiveCallback(
-        _Inout_ struct QUIC_RECV_DATAGRAM* /* Datagram */
-        )
-    {
-        uint8_t RandomValue;
-        QuicRandom(sizeof(RandomValue), &RandomValue);
-        return (RandomValue % 100) < LossPercentage;
+        _Inout_ struct QUIC_RECV_DATAGRAM* Datagram
+        ) {
+        return Instance->Receive(Datagram);
    }
+
    static
    _IRQL_requires_max_(PASSIVE_LEVEL)
    BOOLEAN
    QUIC_API
    SendCallback(
-        _Inout_ QUIC_ADDR* /* RemoteAddress */,
-        _Inout_opt_ QUIC_ADDR* /* LocalAddress */,
-        _Inout_ struct QUIC_DATAPATH_SEND_CONTEXT* /* SendContext */
-        )
-    {
-        return FALSE; // Don't drop
+        _Inout_ QUIC_ADDR* RemoteAddress,
+        _Inout_opt_ QUIC_ADDR* LocalAddress,
+        _Inout_ struct QUIC_DATAPATH_SEND_CONTEXT* SendContext
+        ) {
+        return Instance->Send(RemoteAddress, LocalAddress, SendContext);
+    }
+
+    void Register() {
+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
+        QuicTraceLogInfo(
+            TestHookRegister,
+            "[test][hook] Registering");
+        QUIC_TEST_DATAPATH_HOOKS* Value = &FuncTable;
+        TEST_QUIC_SUCCEEDED(
+            MsQuic->SetParam(
+                nullptr,
+                QUIC_PARAM_LEVEL_GLOBAL,
+                QUIC_PARAM_GLOBAL_TEST_DATAPATH_HOOKS,
+                sizeof(Value),
+                &Value));
+#endif
+    }
+
+    void Unregister() {
+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
+        QuicTraceLogInfo(
+            TestHookUnregistering,
+            "[test][hook] Unregistering");
+        QUIC_TEST_DATAPATH_HOOKS* Value = nullptr;
+        uint32_t TryCount = 0;
+        while (TryCount++ < 20) {
+            if (QUIC_SUCCEEDED(
+                MsQuic->SetParam(
+                    nullptr,
+                    QUIC_PARAM_LEVEL_GLOBAL,
+                    QUIC_PARAM_GLOBAL_TEST_DATAPATH_HOOKS,
+                    sizeof(Value),
+                    &Value))) {
+                break;
+            }
+            QuicSleep(100); // Let the current datapath queue drain.
+        }
+        if (TryCount == 20) {
+            TEST_FAILURE("Failed to disable test datapath hook");
+        }
+        QuicTraceLogInfo(
+            TestHookUnregistered,
+            "[test][hook] Unregistered");
+#endif
+    }
+
+    BOOLEAN
+    Receive(
+        _Inout_ struct QUIC_RECV_DATAGRAM* Datagram
+        ) {
+        BOOLEAN Result = FALSE;
+        QuicDispatchLockAcquire(&Lock);
+        DatapathHook* Iter = Hooks;
+        while (Iter) {
+            if (Iter->Receive(Datagram)) {
+                Result = TRUE;
+                break;
+            }
+            Iter = Iter->Next;
+        }
+        QuicDispatchLockRelease(&Lock);
+        return Result;
+    }
+
+    BOOLEAN
+    Send(
+        _Inout_ QUIC_ADDR* RemoteAddress,
+        _Inout_opt_ QUIC_ADDR* LocalAddress,
+        _Inout_ struct QUIC_DATAPATH_SEND_CONTEXT* SendContext
+        ) {
+        BOOLEAN Result = FALSE;
+        QuicDispatchLockAcquire(&Lock);
+        DatapathHook* Iter = Hooks;
+        while (Iter) {
+            if (Iter->Send(RemoteAddress, LocalAddress, SendContext)) {
+                Result = TRUE;
+                break;
+            }
+            Iter = Iter->Next;
+        }
+        QuicDispatchLockRelease(&Lock);
+        return Result;
+    }
+
+public:
+
+    static DatapathHooks* Instance;
+
+    DatapathHooks() : Hooks(nullptr) {
+        QuicDispatchLockInitialize(&Lock);
+    }
+
+    ~DatapathHooks() {
+        QuicDispatchLockUninitialize(&Lock);
+    }
+
+    void AddHook(DatapathHook* Hook) {
+        QuicDispatchLockAcquire(&Lock);
+        DatapathHook** Iter = &Hooks;
+        while (*Iter != nullptr) {
+            Iter = &((*Iter)->Next);
+        }
+        *Iter = Hook;
+        if (Hooks == Hook) {
+            Register();
+        }
+        QuicDispatchLockRelease(&Lock);
+    }
+
+    void RemoveHook(DatapathHook* Hook) {
+        QuicDispatchLockAcquire(&Lock);
+        DatapathHook** Iter = &Hooks;
+        while (*Iter != Hook) {
+            Iter = &((*Iter)->Next);
+        }
+        *Iter = Hook->Next;
+        if (Hooks == nullptr) {
+            Unregister();
+        }
+        QuicDispatchLockRelease(&Lock);
+    }
+};
+
+struct RandomLossHelper : public DatapathHook
+{
+    uint8_t LossPercentage;
+    RandomLossHelper(uint8_t _LossPercentage) : LossPercentage(_LossPercentage) {
+        if (LossPercentage != 0) {
+            DatapathHooks::Instance->AddHook(this);
+        }
+    }
+    ~RandomLossHelper() {
+        if (LossPercentage != 0) {
+            DatapathHooks::Instance->RemoveHook(this);
+        }
+    }
+    _IRQL_requires_max_(DISPATCH_LEVEL)
+    BOOLEAN
+    Receive(
+        _Inout_ struct QUIC_RECV_DATAGRAM* /* Datagram */
+        ) {
+        uint8_t RandomValue;
+        QuicRandom(sizeof(RandomValue), &RandomValue);
+        auto Result = (RandomValue % 100) < LossPercentage;
+        if (Result) {
+            QuicTraceLogVerbose(
+                TestHookDropPacketRandom,
+                "[test][hook] Random packet drop");
+        }
+        return Result;
+    }
+};
+
+struct SelectiveLossHelper : public DatapathHook
+{
+    uint32_t DropPacketCount;
+    SelectiveLossHelper(uint32_t Count = 0) : DropPacketCount(Count) {
+        DatapathHooks::Instance->AddHook(this);
+    }
+    ~SelectiveLossHelper() {
+        DatapathHooks::Instance->RemoveHook(this);
+    }
+    void DropPackets(uint32_t Count) { DropPacketCount = Count; }
+    _IRQL_requires_max_(DISPATCH_LEVEL)
+    BOOLEAN
+    Receive(
+        _Inout_ struct QUIC_RECV_DATAGRAM* /* Datagram */
+        ) {
+        if (DropPacketCount == 0) {
+            return FALSE;
+        }
+        QuicTraceLogVerbose(
+            TestHookDropPacketSelective,
+            "[test][hook] Selective packet drop");
+        DropPacketCount--;
+        return TRUE;
+    }
+};
+
+struct ReplaceAddressHelper : public DatapathHook
+{
+    QUIC_ADDR Original;
+    QUIC_ADDR New;
+    ReplaceAddressHelper(const QUIC_ADDR& OrigAddr, const QUIC_ADDR& NewAddr) :
+        Original(OrigAddr), New(NewAddr) {
+        DatapathHooks::Instance->AddHook(this);
+    }
+    ~ReplaceAddressHelper() {
+        DatapathHooks::Instance->RemoveHook(this);
+    }
+    _IRQL_requires_max_(DISPATCH_LEVEL)
+    BOOLEAN
+    Receive(
+        _Inout_ struct QUIC_RECV_DATAGRAM* Datagram
+        ) {
+        if (QuicAddrCompare(
+                &Datagram->Tuple->RemoteAddress,
+                &Original)) {
+            Datagram->Tuple->RemoteAddress = New;
+            QuicTraceLogVerbose(
+                TestHookReplaceAddrRecv,
+                "[test][hook] Recv Addr :%hu => :%hu",
+                QuicAddrGetPort(&Original),
+                QuicAddrGetPort(&New));
+        }
+        return FALSE;
+    }
+    _IRQL_requires_max_(PASSIVE_LEVEL)
+    BOOLEAN
+    Send(
+        _Inout_ QUIC_ADDR* RemoteAddress,
+        _Inout_opt_ QUIC_ADDR* /* LocalAddress */,
+        _Inout_ struct QUIC_DATAPATH_SEND_CONTEXT* /* SendContext */
+        ) {
+        if (QuicAddrCompare(RemoteAddress, &New)) {
+            *RemoteAddress = Original;
+            QuicTraceLogVerbose(
+                TestHookReplaceAddrSend,
+                "[test][hook] Send Addr :%hu => :%hu",
+                QuicAddrGetPort(&New),
+                QuicAddrGetPort(&Original));
+        }
+        return FALSE;
    }
 };