driver: memory: move NBL pools to global scope

This is preparation for the next commit, which will attempt to allocate
everything at once for the RX path.

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

9 files changed, 64 insertions, 88 deletions

diff --git a/driver/device.c b/driver/device.c
index 899c4a2..500346f 100644
--- a/driver/device.c
+++ b/driver/device.c
@@ -160,7 +160,7 @@ SendNetBufferLists(
         }
 
         NET_BUFFER_LIST *CloneNbl = MemAllocateNetBufferListWithClonedGeometry(
-            Wg->NblPool, Wg->NbPool, Nbl, sizeof(MESSAGE_DATA) + NoiseEncryptedLen(0) + MESSAGE_PADDING_MULTIPLE - 1);
+            Nbl, sizeof(MESSAGE_DATA) + NoiseEncryptedLen(0) + MESSAGE_PADDING_MULTIPLE - 1);
         if (!CloneNbl)
         {
             NET_BUFFER_LIST_STATUS(Nbl) = NDIS_STATUS_RESOURCES;
@@ -255,15 +255,6 @@ CancelSend(NDIS_HANDLE MiniportAdapterContext, PVOID CancelId)
 {
 }
 
-static MINIPORT_RETURN_NET_BUFFER_LISTS ReturnNetBufferLists;
-_Use_decl_annotations_
-static VOID
-ReturnNetBufferLists(NDIS_HANDLE MiniportAdapterContext, PNET_BUFFER_LIST NetBufferLists, ULONG ReturnFlags)
-{
-    WG_DEVICE *Wg = (WG_DEVICE *)MiniportAdapterContext;
-    FreeReceiveNetBufferList(Wg, NetBufferLists);
-}
-
 static EX_CALLBACK_FUNCTION MtuRegistryChange;
 _Use_decl_annotations_
 static NTSTATUS
@@ -396,8 +387,6 @@ HaltEx(NDIS_HANDLE MiniportAdapterContext, NDIS_HALT_ACTION HaltAction)
     MemFree(Wg->PeerHashtable);
     MuReleasePushLockExclusive(&Wg->DeviceUpdateLock);
 
-    NdisFreeNetBufferPool(Wg->NbPool);
-    NdisFreeNetBufferListPool(Wg->NblPool);
     WritePointerNoFence(&Wg->MiniportAdapterHandle, NULL);
     LogInfo(Wg, "Interface destroyed");
     MemFree(Wg);
@@ -567,30 +556,6 @@ InitializeEx(
      * Revisit this when we drop support for old Windows versions. */
     Wg->FunctionalDeviceObject->Reserved = Wg;
 
-    NET_BUFFER_LIST_POOL_PARAMETERS NblPoolParameters = {
-        .Header = { .Type = NDIS_OBJECT_TYPE_DEFAULT,
-                    .Revision = NET_BUFFER_LIST_POOL_PARAMETERS_REVISION_1,
-                    .Size = NDIS_SIZEOF_NET_BUFFER_LIST_POOL_PARAMETERS_REVISION_1 },
-        .ProtocolId = NDIS_PROTOCOL_ID_DEFAULT,
-        .PoolTag = MEMORY_TAG
-    };
-    Wg->NblPool = NdisAllocateNetBufferListPool(MiniportAdapterHandle, &NblPoolParameters);
-    if (!Wg->NblPool)
-    {
-        Status = STATUS_INSUFFICIENT_RESOURCES;
-        goto cleanupWg;
-    }
-    NET_BUFFER_POOL_PARAMETERS NbPoolParameters = { .Header = { .Type = NDIS_OBJECT_TYPE_DEFAULT,
-                                                                .Revision = NET_BUFFER_POOL_PARAMETERS_REVISION_1,
-                                                                .Size =
-                                                                    NDIS_SIZEOF_NET_BUFFER_POOL_PARAMETERS_REVISION_1 },
-                                                    .PoolTag = MEMORY_TAG };
-    Wg->NbPool = NdisAllocateNetBufferPool(MiniportAdapterHandle, &NbPoolParameters);
-    if (!Wg->NbPool)
-    {
-        Status = STATUS_INSUFFICIENT_RESOURCES;
-        goto cleanupNblPool;
-    }
     ExInitializeRundownProtection(&Wg->ItemsInFlight);
     ExRundownCompleted(&Wg->ItemsInFlight); /* Wait until Restart is called to mark this active. */
 
@@ -609,7 +574,7 @@ InitializeEx(
 
     Wg->PeerHashtable = PubkeyHashtableAlloc();
     if (!Wg->PeerHashtable)
-        goto cleanupNbPool;
+        goto cleanupWg;
 
     Wg->IndexHashtable = IndexHashtableAlloc();
     if (!Wg->IndexHashtable)
@@ -689,10 +654,6 @@ cleanupIndexHashtable:
     MemFree(Wg->IndexHashtable);
 cleanupPeerHashtable:
     MemFree(Wg->PeerHashtable);
-cleanupNbPool:
-    NdisFreeNetBufferPool(Wg->NbPool);
-cleanupNblPool:
-    NdisFreeNetBufferListPool(Wg->NblPool);
 cleanupWg:
     MemFree(Wg);
     if (Status == STATUS_INSUFFICIENT_RESOURCES)
diff --git a/driver/device.h b/driver/device.h
index 83704f3..fb402ce 100644
--- a/driver/device.h
+++ b/driver/device.h
@@ -73,7 +73,6 @@ typedef struct _WG_DEVICE
     NDIS_HANDLE MiniportAdapterHandle; /* This is actually a pointer to NDIS_MINIPORT_BLOCK struct. */
     DEVICE_OBJECT *FunctionalDeviceObject;
     NDIS_STATISTICS_INFO Statistics;
-    NDIS_HANDLE NblPool, NbPool;
     EX_RUNDOWN_REF ItemsInFlight;
     PTR_RING EncryptQueue, DecryptQueue;
     PEER_SERIAL TxQueue, RxQueue, HandshakeTxQueue;
diff --git a/driver/memory.c b/driver/memory.c
index 4718909..afdc260 100644
--- a/driver/memory.c
+++ b/driver/memory.c
@@ -8,6 +8,7 @@
 
 static CONST ULONG PacketCacheSizes[] = { 192, 512, 1024, 1500, 9000 };
 static LOOKASIDE_ALIGN LOOKASIDE_LIST_EX PacketCaches[ARRAYSIZE(PacketCacheSizes)];
+static NDIS_HANDLE NblPool, NbPool;
 
 _Must_inspect_result_
 _IRQL_requires_max_(DISPATCH_LEVEL)
@@ -76,7 +77,7 @@ MemFreeDataAndMdlChain(MDL *Mdl)
 
 _Use_decl_annotations_
 NET_BUFFER_LIST *
-MemAllocateNetBufferList(NDIS_HANDLE NblPool, NDIS_HANDLE NbPool, ULONG SpaceBefore, ULONG Size, ULONG SpaceAfter)
+MemAllocateNetBufferList(ULONG SpaceBefore, ULONG Size, ULONG SpaceAfter)
 {
     ULONG Sum = Size;
     if (!NT_SUCCESS(RtlULongAdd(Sum, SpaceBefore, &Sum) || !NT_SUCCESS(RtlULongAdd(Sum, SpaceAfter, &Sum))) ||
@@ -120,8 +121,6 @@ MemFreeNetBufferList(NET_BUFFER_LIST *Nbl)
 _Use_decl_annotations_
 NET_BUFFER_LIST *
 MemAllocateNetBufferListWithClonedGeometry(
-    NDIS_HANDLE NblPool,
-    NDIS_HANDLE NbPool,
     NET_BUFFER_LIST *Original,
     ULONG AdditionalBytesPerNb)
 {
@@ -157,6 +156,13 @@ cleanupClone:
 }
 
 _Use_decl_annotations_
+BOOLEAN
+MemNetBufferListIsOurs(NET_BUFFER_LIST *Nbl)
+{
+    return Nbl->NdisPoolHandle == NblPool;
+}
+
+_Use_decl_annotations_
 NTSTATUS
 MemCopyFromMdl(VOID *Dst, MDL *Src, ULONG Offset, ULONG Size)
 {
@@ -202,12 +208,40 @@ MemDriverEntry(VOID)
             return Status;
         }
     }
+
+    NET_BUFFER_LIST_POOL_PARAMETERS NblPoolParameters = {
+        .Header = { .Type = NDIS_OBJECT_TYPE_DEFAULT,
+                    .Revision = NET_BUFFER_LIST_POOL_PARAMETERS_REVISION_1,
+                    .Size = NDIS_SIZEOF_NET_BUFFER_LIST_POOL_PARAMETERS_REVISION_1 },
+        .ProtocolId = NDIS_PROTOCOL_ID_DEFAULT,
+        .PoolTag = MEMORY_TAG
+    };
+    NblPool = NdisAllocateNetBufferListPool(NULL, &NblPoolParameters);
+    if (!NblPool)
+        goto cleanupPacketCaches;
+    NET_BUFFER_POOL_PARAMETERS NbPoolParameters = { .Header = { .Type = NDIS_OBJECT_TYPE_DEFAULT,
+                                                                .Revision = NET_BUFFER_POOL_PARAMETERS_REVISION_1,
+                                                                .Size =
+                                                                    NDIS_SIZEOF_NET_BUFFER_POOL_PARAMETERS_REVISION_1 },
+                                                    .PoolTag = MEMORY_TAG };
+    NbPool = NdisAllocateNetBufferPool(NULL, &NbPoolParameters);
+    if (!NbPool)
+        goto cleanupNblPool;
     return STATUS_SUCCESS;
+
+cleanupNblPool:
+    NdisFreeNetBufferListPool(NblPool);
+cleanupPacketCaches:
+    for (ULONG i = 0; i < ARRAYSIZE(PacketCacheSizes); ++i)
+        ExDeleteLookasideListEx(&PacketCaches[i]);
+    return STATUS_INSUFFICIENT_RESOURCES;
 }
 
 _Use_decl_annotations_
 VOID MemUnload(VOID)
 {
+    NdisFreeNetBufferPool(NbPool);
+    NdisFreeNetBufferListPool(NblPool);
     for (ULONG i = 0; i < ARRAYSIZE(PacketCacheSizes); ++i)
         ExDeleteLookasideListEx(&PacketCaches[i]);
 }
diff --git a/driver/memory.h b/driver/memory.h
index 99a691e..a51d87d 100644
--- a/driver/memory.h
+++ b/driver/memory.h
@@ -115,8 +115,6 @@ _Return_type_success_(return != NULL)
 __drv_allocatesMem(mem)
 NET_BUFFER_LIST *
 MemAllocateNetBufferList(
-    _In_ NDIS_HANDLE NblPool,
-    _In_ NDIS_HANDLE NbPool,
     _In_ ULONG SpaceBefore,
     _In_ ULONG Size,
     _In_ ULONG SpaceAfter);
@@ -125,11 +123,12 @@ _IRQL_requires_max_(DISPATCH_LEVEL)
 __drv_allocatesMem(mem)
 NET_BUFFER_LIST *
 MemAllocateNetBufferListWithClonedGeometry(
-    _In_ NDIS_HANDLE NblPool,
-    _In_ NDIS_HANDLE NbPool,
     _In_ NET_BUFFER_LIST *Original,
     _In_ ULONG AdditionalBytesPerNb);
 
+BOOLEAN
+MemNetBufferListIsOurs(_In_ NET_BUFFER_LIST *Nbl);
+
 _IRQL_requires_max_(DISPATCH_LEVEL)
 VOID
 MemFreeNetBufferList(__drv_freesMem(mem) _In_ NET_BUFFER_LIST *Nbl);
diff --git a/driver/queueing.h b/driver/queueing.h
index ea87ca3..3deead6 100644
--- a/driver/queueing.h
+++ b/driver/queueing.h
@@ -125,9 +125,15 @@ _IRQL_requires_max_(DISPATCH_LEVEL)
 VOID
 PacketReceive(_Inout_ WG_DEVICE *Wg, _In_ __drv_aliasesMem NET_BUFFER_LIST *First);
 
+MINIPORT_RETURN_NET_BUFFER_LISTS ReturnNetBufferLists;
+
 _IRQL_requires_max_(DISPATCH_LEVEL)
-VOID
-FreeReceiveNetBufferList(_Inout_ WG_DEVICE *Wg, __drv_freesMem(Mem) _In_ NET_BUFFER_LIST *First);
+static inline VOID
+FreeReceiveNetBufferList(_In_ NET_BUFFER_LIST *First)
+{
+    if (First)
+        ReturnNetBufferLists(First->SourceHandle, First, 0);
+}
 
 _IRQL_requires_max_(DISPATCH_LEVEL)
 VOID
diff --git a/driver/receive.c b/driver/receive.c
index cda511a..a7bed9b 100644
--- a/driver/receive.c
+++ b/driver/receive.c
@@ -152,7 +152,7 @@ PacketHandshakeRxWorker(MULTICORE_WORKQUEUE *WorkQueue)
     while ((Nbl = NetBufferListInterlockedDequeue(&Wg->HandshakeRxQueue)) != NULL)
     {
         ReceiveHandshakePacket(Wg, Nbl);
-        FreeReceiveNetBufferList(Wg, Nbl);
+        FreeReceiveNetBufferList(Nbl);
     }
 }
 
@@ -367,7 +367,7 @@ falsePacket:
     ++Peer->Device->Statistics.ifInErrors;
     ++Peer->Device->Statistics.ifInDiscards;
 packetProcessed:
-    FreeReceiveNetBufferList(Peer->Device, Nbl);
+    FreeReceiveNetBufferList(Nbl);
     return FALSE;
 }
 
@@ -419,7 +419,7 @@ PacketPeerRxWork(_Inout_ WG_PEER *Peer, _In_ ULONG Budget)
     next:
         NoiseKeypairPut(Keypair, FALSE);
         if (Free)
-            FreeReceiveNetBufferList(Peer->Device, Nbl);
+            FreeReceiveNetBufferList(Nbl);
         ExReleaseRundownProtection(&Peer->InUse);
         PeerPut(Peer);
     }
@@ -503,7 +503,7 @@ PacketConsumeData(_Inout_ WG_DEVICE *Wg, _Inout_ __drv_aliasesMem NET_BUFFER_LIS
     cleanupKeypair:
         NoiseKeypairPut(Keypair, FALSE);
     cleanupNbl:
-        FreeReceiveNetBufferList(Wg, Nbl);
+        FreeReceiveNetBufferList(Nbl);
         PeerPut(Peer);
     }
     if (FirstForDevice && !QueueEnqueuePerDevice(&Wg->DecryptQueue, &Wg->DecryptThreads, FirstForDevice))
@@ -600,7 +600,7 @@ PacketReceive(WG_DEVICE *Wg, NET_BUFFER_LIST *First)
         continue;
 
     cleanup:
-        FreeReceiveNetBufferList(Wg, Nbl);
+        FreeReceiveNetBufferList(Nbl);
     }
 
     if (FirstData)
@@ -609,7 +609,7 @@ PacketReceive(WG_DEVICE *Wg, NET_BUFFER_LIST *First)
 
 _Use_decl_annotations_
 VOID
-FreeReceiveNetBufferList(WG_DEVICE *Wg, NET_BUFFER_LIST *First)
+ReturnNetBufferLists(NDIS_HANDLE MiniportAdapterContext, PNET_BUFFER_LIST First, ULONG ReturnFlags)
 {
     for (NET_BUFFER_LIST *Nbl = First, *NextNbl; Nbl; Nbl = NextNbl)
     {
@@ -630,5 +630,5 @@ FreeIncomingHandshakes(WG_DEVICE *Wg)
 {
     NET_BUFFER_LIST *Nbl;
     while ((Nbl = NetBufferListInterlockedDequeue(&Wg->HandshakeRxQueue)) != NULL)
-        FreeReceiveNetBufferList(Wg, Nbl);
+        FreeReceiveNetBufferList(Nbl);
 }
diff --git a/driver/selftest/ratelimiter.c b/driver/selftest/ratelimiter.c
index 2859103..b430093 100644
--- a/driver/selftest/ratelimiter.c
+++ b/driver/selftest/ratelimiter.c
@@ -146,42 +146,24 @@ RatelimiterSelftest(VOID)
     IPV4HDR *Hdr4;
     IPV6HDR *Hdr6 = NULL;
 
-    NET_BUFFER_LIST_POOL_PARAMETERS NblPoolParameters = {
-        .Header = { .Type = NDIS_OBJECT_TYPE_DEFAULT,
-                    .Revision = NET_BUFFER_LIST_POOL_PARAMETERS_REVISION_1,
-                    .Size = NDIS_SIZEOF_NET_BUFFER_LIST_POOL_PARAMETERS_REVISION_1 },
-        .ProtocolId = NDIS_PROTOCOL_ID_DEFAULT,
-        .PoolTag = MEMORY_TAG
-    };
-    NDIS_HANDLE NblPool = NdisAllocateNetBufferListPool(NULL, &NblPoolParameters);
-    if (!NblPool)
-        goto out;
-    NET_BUFFER_POOL_PARAMETERS NbPoolParameters = { .Header = { .Type = NDIS_OBJECT_TYPE_DEFAULT,
-                                                                .Revision = NET_BUFFER_POOL_PARAMETERS_REVISION_1,
-                                                                .Size =
-                                                                    NDIS_SIZEOF_NET_BUFFER_POOL_PARAMETERS_REVISION_1 },
-                                                    .PoolTag = MEMORY_TAG };
-    NDIS_HANDLE NbPool = NdisAllocateNetBufferPool(NULL, &NbPoolParameters);
-    if (!NbPool)
-        goto cleanupNblPool;
     ++Test;
     ++Test;
     ++Test;
 
-    Nbl4 = MemAllocateNetBufferList(NblPool, NbPool, 0, sizeof(*Hdr4), 0);
+    Nbl4 = MemAllocateNetBufferList(0, sizeof(*Hdr4), 0);
     if (!Nbl4)
-        goto cleanupNofree;
+        goto out;
     NdisSetNblFlag(Nbl4, NDIS_NBL_FLAGS_IS_IPV4);
     NET_BUFFER_LIST_INFO(Nbl4, NetBufferListProtocolId) = (VOID *)Htons(NDIS_ETH_TYPE_IPV4);
     Hdr4 = MemGetValidatedNetBufferListData(Nbl4);
     Hdr4->Saddr = Htonl(8182);
     ++Test;
 
-    Nbl6 = MemAllocateNetBufferList(NblPool, NbPool, 0, sizeof(*Hdr6), 0);
+    Nbl6 = MemAllocateNetBufferList(0, sizeof(*Hdr6), 0);
     if (!Nbl6)
     {
         MemFreeNetBufferList(Nbl4);
-        goto cleanupNofree;
+        goto out;
     }
     NdisSetNblFlag(Nbl6, NDIS_NBL_FLAGS_IS_IPV6);
     NET_BUFFER_LIST_INFO(Nbl6, NetBufferListProtocolId) = (VOID *)Htons(NDIS_ETH_TYPE_IPV6);
@@ -241,10 +223,6 @@ RatelimiterSelftest(VOID)
 cleanup:
     MemFreeNetBufferList(Nbl4);
     MemFreeNetBufferList(Nbl6);
-cleanupNofree:
-    NdisFreeNetBufferPool(NbPool);
-cleanupNblPool:
-    NdisFreeNetBufferListPool(NblPool);
 out:
     if (Success)
         LogDebug("ratelimiter self-tests: pass");
diff --git a/driver/send.c b/driver/send.c
index 568b332..c2fadca 100644
--- a/driver/send.c
+++ b/driver/send.c
@@ -247,8 +247,7 @@ PacketSendKeepalive(WG_PEER *Peer)
 
     if (NetBufferListIsQueueEmpty(&Peer->StagedPacketQueue))
     {
-        Nbl = MemAllocateNetBufferList(
-            Peer->Device->NblPool, Peer->Device->NbPool, 0, 0, sizeof(MESSAGE_DATA) + NoiseEncryptedLen(0));
+        Nbl = MemAllocateNetBufferList(0, 0, sizeof(MESSAGE_DATA) + NoiseEncryptedLen(0));
         if (!Nbl)
             return;
         Nbl->ParentNetBufferList = Nbl;
@@ -506,7 +505,7 @@ FreeSendNetBufferList(WG_DEVICE *Wg, NET_BUFFER_LIST *FirstNbl, ULONG SendComple
     {
         NextNbl = NET_BUFFER_LIST_NEXT_NBL(Nbl);
         NET_BUFFER_LIST_NEXT_NBL(Nbl) = NULL;
-        if (Nbl->NdisPoolHandle == Wg->NblPool)
+        if (MemNetBufferListIsOurs(Nbl))
         {
             if (Nbl->ParentNetBufferList != Nbl)
             {
diff --git a/driver/socket.c b/driver/socket.c
index 836e344..4cd38f6 100644
--- a/driver/socket.c
+++ b/driver/socket.c
@@ -613,7 +613,7 @@ Receive(_In_opt_ PVOID SocketContext, _In_ ULONG Flags, _In_opt_ WSK_DATAGRAM_IN
         ULONG Length;
         if (!NT_SUCCESS(RtlSIZETToULong(DataIndication->Buffer.Length, &Length)))
             goto skipDatagramIndication;
-        Nbl = MemAllocateNetBufferList(Wg->NblPool, Wg->NbPool, 0, Length, 0);
+        Nbl = MemAllocateNetBufferList(0, Length, 0);
         if (!Nbl || !ReadBooleanNoFence(&Wg->IsUp) || !ExAcquireRundownProtection(&Socket->ItemsInFlight))
             goto skipDatagramIndication;
         NET_BUFFER_LIST_DATAGRAM_INDICATION(Nbl) = DataIndication;