mirror of
https://git.zx2c4.com/wireguard-nt
synced 2024-09-22 06:11:35 +00:00
f0a85143d7
Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>
299 lines
9.9 KiB
C
299 lines
9.9 KiB
C
/* SPDX-License-Identifier: GPL-2.0
|
|
*
|
|
* Copyright (C) 2015-2021 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved.
|
|
*/
|
|
|
|
#include "memory.h"
|
|
#include "messages.h"
|
|
|
|
static CONST ULONG PacketCacheSizes[] = { 192, 512, 1024, 1500, 9000 };
|
|
static NDIS_HANDLE LooseNbPool, LooseNblPool;
|
|
static NDIS_HANDLE NbDataPools[ARRAYSIZE(PacketCacheSizes)], NblDataPools[ARRAYSIZE(PacketCacheSizes)];
|
|
|
|
#pragma warning(suppress : 28195) /* IoAllocateMdl allocates, even if missing the SAL annotation. */
|
|
_Use_decl_annotations_
|
|
MDL *
|
|
MemAllocateDataAndMdlChain(ULONG BufferSize)
|
|
{
|
|
NT_ASSERT(BufferSize <= (MAXULONG - PAGE_SIZE));
|
|
VOID *Memory = MemAllocate(BufferSize);
|
|
if (!Memory)
|
|
return NULL;
|
|
MDL *Mdl = IoAllocateMdl(Memory, BufferSize, FALSE, FALSE, NULL);
|
|
if (!Mdl)
|
|
{
|
|
MemFree(Memory);
|
|
return NULL;
|
|
}
|
|
MmBuildMdlForNonPagedPool(Mdl);
|
|
return Mdl;
|
|
}
|
|
|
|
#pragma warning(suppress : 6014) /* IoFreeMdl frees, even if missing the SAL annotation. */
|
|
_Use_decl_annotations_
|
|
VOID
|
|
MemFreeDataAndMdlChain(MDL *Mdl)
|
|
{
|
|
while (Mdl)
|
|
{
|
|
MDL *Next = Mdl->Next;
|
|
VOID *Memory = MmGetMdlVirtualAddress(Mdl);
|
|
IoFreeMdl(Mdl);
|
|
MemFree(Memory);
|
|
Mdl = Next;
|
|
}
|
|
}
|
|
|
|
_Use_decl_annotations_
|
|
NET_BUFFER_LIST *
|
|
MemAllocateNetBufferList(ULONG SpaceBefore, ULONG Size, ULONG SpaceAfter)
|
|
{
|
|
ULONG Sum = Size;
|
|
if (!NT_SUCCESS(RtlULongAdd(Sum, SpaceBefore, &Sum) || !NT_SUCCESS(RtlULongAdd(Sum, SpaceAfter, &Sum))) ||
|
|
Sum > MTU_MAX)
|
|
return NULL;
|
|
for (ULONG i = 0; i < ARRAYSIZE(PacketCacheSizes); ++i)
|
|
{
|
|
if (PacketCacheSizes[i] >= Sum)
|
|
{
|
|
NET_BUFFER_LIST *Nbl = NdisAllocateNetBufferList(NblDataPools[i], 0, 0);
|
|
if (!Nbl)
|
|
return NULL;
|
|
NET_BUFFER_DATA_LENGTH(NET_BUFFER_LIST_FIRST_NB(Nbl)) = Size;
|
|
NET_BUFFER_DATA_OFFSET(NET_BUFFER_LIST_FIRST_NB(Nbl)) =
|
|
NET_BUFFER_CURRENT_MDL_OFFSET(NET_BUFFER_LIST_FIRST_NB(Nbl)) = SpaceBefore;
|
|
return Nbl;
|
|
}
|
|
}
|
|
|
|
#pragma warning(suppress : 6014) /* MDL is aliased in Nbl or freed on failure. */
|
|
MDL *Mdl = MemAllocateDataAndMdlChain(Sum);
|
|
if (!Mdl)
|
|
return NULL;
|
|
NET_BUFFER *Nb = NdisAllocateNetBuffer(LooseNbPool, Mdl, SpaceBefore, Size);
|
|
if (!Nb)
|
|
goto cleanupMdl;
|
|
NET_BUFFER_LIST *Nbl = NdisAllocateNetBufferList(LooseNblPool, 0, 0);
|
|
if (!Nbl)
|
|
goto cleanupNb;
|
|
NET_BUFFER_LIST_FIRST_NB(Nbl) = Nb;
|
|
return Nbl;
|
|
|
|
cleanupNb:
|
|
NdisFreeNetBuffer(Nb);
|
|
cleanupMdl:
|
|
MemFreeDataAndMdlChain(Mdl);
|
|
return NULL;
|
|
}
|
|
|
|
_Use_decl_annotations_
|
|
VOID
|
|
MemFreeNetBufferList(NET_BUFFER_LIST *Nbl)
|
|
{
|
|
if (Nbl->NdisPoolHandle == LooseNblPool)
|
|
{
|
|
while (NET_BUFFER_LIST_FIRST_NB(Nbl))
|
|
{
|
|
NET_BUFFER *Nb = NET_BUFFER_LIST_FIRST_NB(Nbl);
|
|
NET_BUFFER_LIST_FIRST_NB(Nbl) = NET_BUFFER_NEXT_NB(NET_BUFFER_LIST_FIRST_NB(Nbl));
|
|
if (Nb->NdisPoolHandle == LooseNbPool)
|
|
MemFreeDataAndMdlChain(NET_BUFFER_FIRST_MDL(Nb));
|
|
NdisFreeNetBuffer(Nb);
|
|
}
|
|
}
|
|
NdisFreeNetBufferList(Nbl);
|
|
}
|
|
|
|
#pragma warning(suppress : 28195) /* NdisAllocateNetBufferList & co allocate. */
|
|
_Use_decl_annotations_
|
|
NET_BUFFER_LIST *
|
|
MemAllocateNetBufferListWithClonedGeometry(NET_BUFFER_LIST *Original, ULONG AdditionalBytesPerNb)
|
|
{
|
|
if (NET_BUFFER_LIST_FIRST_NB(Original) && !NET_BUFFER_NEXT_NB(NET_BUFFER_LIST_FIRST_NB(Original)))
|
|
{
|
|
ULONG Length = NET_BUFFER_DATA_LENGTH(NET_BUFFER_LIST_FIRST_NB(Original));
|
|
if (Length > MTU_MAX ||
|
|
!NT_SUCCESS(RtlULongAdd(Length, AdditionalBytesPerNb, &Length)))
|
|
return NULL;
|
|
NET_BUFFER_LIST *Clone = MemAllocateNetBufferList(0, Length, 0);
|
|
if (!Clone)
|
|
return NULL;
|
|
NET_BUFFER_LIST_INFO(Clone, NetBufferListProtocolId) = NET_BUFFER_LIST_INFO(Original, NetBufferListProtocolId);
|
|
Clone->ParentNetBufferList = Original;
|
|
return Clone;
|
|
}
|
|
|
|
NET_BUFFER_LIST *Clone = NdisAllocateNetBufferList(LooseNblPool, 0, 0);
|
|
if (!Clone)
|
|
return NULL;
|
|
NET_BUFFER_LIST_INFO(Clone, NetBufferListProtocolId) = NET_BUFFER_LIST_INFO(Original, NetBufferListProtocolId);
|
|
NET_BUFFER **CloneNb = &NET_BUFFER_LIST_FIRST_NB(Clone);
|
|
for (NET_BUFFER *Nb = NET_BUFFER_LIST_FIRST_NB(Original); Nb; Nb = NET_BUFFER_NEXT_NB(Nb))
|
|
{
|
|
ULONG Length;
|
|
if (NET_BUFFER_DATA_LENGTH(Nb) > MTU_MAX ||
|
|
!NT_SUCCESS(RtlULongAdd(NET_BUFFER_DATA_LENGTH(Nb), AdditionalBytesPerNb, &Length)))
|
|
goto cleanupClone;
|
|
for (ULONG i = 0; i < ARRAYSIZE(PacketCacheSizes); ++i)
|
|
{
|
|
if (PacketCacheSizes[i] >= Length)
|
|
{
|
|
*CloneNb = NdisAllocateNetBufferMdlAndData(NbDataPools[i]);
|
|
if (!*CloneNb)
|
|
goto cleanupClone;
|
|
NET_BUFFER_DATA_LENGTH(*CloneNb) = Length;
|
|
goto linkNb;
|
|
}
|
|
}
|
|
#pragma warning(suppress : 6014) /* `CloneMdl` is aliased in NdisAllocateNetBuffer or freed on failure. */
|
|
MDL *CloneMdl = MemAllocateDataAndMdlChain(Length);
|
|
if (!CloneMdl)
|
|
goto cleanupClone;
|
|
#pragma warning(suppress : 6014) /* `*CloneNb` is aliased in Clone or freed on failure. */
|
|
*CloneNb = NdisAllocateNetBuffer(LooseNbPool, CloneMdl, 0, 0);
|
|
if (!*CloneNb)
|
|
{
|
|
MemFreeDataAndMdlChain(CloneMdl);
|
|
goto cleanupClone;
|
|
}
|
|
linkNb:
|
|
CloneNb = &NET_BUFFER_NEXT_NB(*CloneNb);
|
|
}
|
|
Clone->ParentNetBufferList = Original;
|
|
return Clone;
|
|
|
|
cleanupClone:
|
|
MemFreeNetBufferList(Clone);
|
|
return NULL;
|
|
}
|
|
|
|
_Use_decl_annotations_
|
|
BOOLEAN
|
|
MemNetBufferListIsOurs(NET_BUFFER_LIST *Nbl)
|
|
{
|
|
if (Nbl->NdisPoolHandle == LooseNblPool)
|
|
return TRUE;
|
|
for (ULONG i = 0; i < ARRAYSIZE(PacketCacheSizes); ++i)
|
|
{
|
|
if (Nbl->NdisPoolHandle == NblDataPools[i])
|
|
return TRUE;
|
|
}
|
|
return FALSE;
|
|
}
|
|
|
|
_Use_decl_annotations_
|
|
NTSTATUS
|
|
MemCopyFromMdl(VOID *Dst, MDL *Src, ULONG Offset, ULONG Size)
|
|
{
|
|
if (!Src)
|
|
return STATUS_BUFFER_TOO_SMALL;
|
|
UCHAR *DstBuf = Dst;
|
|
while (Offset >= MmGetMdlByteCount(Src))
|
|
{
|
|
Offset -= MmGetMdlByteCount(Src);
|
|
Src = Src->Next;
|
|
if (!Src)
|
|
return STATUS_BUFFER_TOO_SMALL;
|
|
}
|
|
for (ULONG CurSize; Size; Src = Src->Next, Size -= CurSize, DstBuf += CurSize)
|
|
{
|
|
if (!Src)
|
|
return STATUS_BUFFER_TOO_SMALL;
|
|
UCHAR *SrcBuf = MmGetSystemAddressForMdlSafe(Src, NormalPagePriority | MdlMappingNoExecute | MdlMappingNoWrite);
|
|
if (!SrcBuf)
|
|
return STATUS_INSUFFICIENT_RESOURCES;
|
|
CurSize = min(MmGetMdlByteCount(Src) - Offset, Size);
|
|
RtlCopyMemory(DstBuf, SrcBuf + Offset, CurSize);
|
|
Offset = 0;
|
|
}
|
|
return STATUS_SUCCESS;
|
|
}
|
|
|
|
#ifdef ALLOC_PRAGMA
|
|
# pragma alloc_text(INIT, MemDriverEntry)
|
|
#endif
|
|
_Use_decl_annotations_
|
|
NTSTATUS
|
|
MemDriverEntry(VOID)
|
|
{
|
|
for (ULONG i = 0; i < ARRAYSIZE(PacketCacheSizes); ++i)
|
|
{
|
|
NET_BUFFER_LIST_POOL_PARAMETERS NblDataPoolParameters = {
|
|
.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,
|
|
.fAllocateNetBuffer = TRUE,
|
|
.DataSize = PacketCacheSizes[i]
|
|
};
|
|
NblDataPools[i] = NdisAllocateNetBufferListPool(NULL, &NblDataPoolParameters);
|
|
if (!NblDataPools[i])
|
|
{
|
|
for (ULONG j = 0; j < i; ++j)
|
|
NdisFreeNetBufferListPool(NblDataPools[j]);
|
|
goto cleanup;
|
|
}
|
|
}
|
|
for (ULONG i = 0; i < ARRAYSIZE(PacketCacheSizes); ++i)
|
|
{
|
|
NET_BUFFER_POOL_PARAMETERS NbDataPoolParameters = {
|
|
.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,
|
|
.DataSize = PacketCacheSizes[i]
|
|
};
|
|
NbDataPools[i] = NdisAllocateNetBufferPool(NULL, &NbDataPoolParameters);
|
|
if (!NbDataPools[i])
|
|
{
|
|
for (ULONG j = 0; j < i; ++j)
|
|
NdisFreeNetBufferPool(NbDataPools[j]);
|
|
goto cleanupNblDataPools;
|
|
}
|
|
}
|
|
NET_BUFFER_LIST_POOL_PARAMETERS LooseNblPoolParameters = {
|
|
.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
|
|
};
|
|
LooseNblPool = NdisAllocateNetBufferListPool(NULL, &LooseNblPoolParameters);
|
|
if (!LooseNblPool)
|
|
goto cleanupNbDataPools;
|
|
NET_BUFFER_POOL_PARAMETERS LooseNbPoolParameters = {
|
|
.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
|
|
};
|
|
LooseNbPool = NdisAllocateNetBufferPool(NULL, &LooseNbPoolParameters);
|
|
if (!LooseNbPool)
|
|
goto cleanupLooseNblPool;
|
|
return STATUS_SUCCESS;
|
|
|
|
cleanupLooseNblPool:
|
|
NdisFreeNetBufferListPool(LooseNblPool);
|
|
cleanupNbDataPools:
|
|
for (ULONG i = 0; i < ARRAYSIZE(PacketCacheSizes); ++i)
|
|
NdisFreeNetBufferPool(NbDataPools[i]);
|
|
cleanupNblDataPools:
|
|
for (ULONG i = 0; i < ARRAYSIZE(PacketCacheSizes); ++i)
|
|
NdisFreeNetBufferListPool(NblDataPools[i]);
|
|
cleanup:
|
|
return STATUS_INSUFFICIENT_RESOURCES;
|
|
}
|
|
|
|
_Use_decl_annotations_
|
|
VOID MemUnload(VOID)
|
|
{
|
|
NdisFreeNetBufferPool(LooseNbPool);
|
|
NdisFreeNetBufferListPool(LooseNblPool);
|
|
for (ULONG i = 0; i < ARRAYSIZE(PacketCacheSizes); ++i)
|
|
{
|
|
NdisFreeNetBufferPool(NbDataPools[i]);
|
|
NdisFreeNetBufferListPool(NblDataPools[i]);
|
|
}
|
|
}
|