yttrium-virtio-gpu/NetKVM/Common/ParaNdis_TX.cpp

#include "ndis56common.h"
#include "kdebugprint.h"
#include "Trace.h"
#ifdef NETKVM_WPP_ENABLED
#include "ParaNdis_TX.tmh"
#endif

CNBL::CNBL(PNET_BUFFER_LIST NBL,
           PPARANDIS_ADAPTER Context,
           CParaNdisTX &ParentTXPath,
           CAllocationHelper<CNBL> *NBLAllocator,
           CAllocationHelper<CNB> *NBAllocator)
    : m_NBL(NBL), m_Context(Context), m_ParentTXPath(&ParentTXPath), CNdisAllocatableViaHelper<CNBL>(NBLAllocator),
      m_NBAllocator(NBAllocator)
{
    m_NBL->Scratch = this;
    m_LsoInfo.Value = NET_BUFFER_LIST_INFO(m_NBL, TcpLargeSendNetBufferListInfo);
    m_CsoInfo.Value = NET_BUFFER_LIST_INFO(m_NBL, TcpIpChecksumNetBufferListInfo);
#if PARANDIS_SUPPORT_USO
    m_UsoInfo.Value = NET_BUFFER_LIST_INFO(m_NBL, UdpSegmentationOffloadInfo);
#endif
    ParaNdis_DebugNBLIn(NBL, m_LogIndex);
}

CNBL::~CNBL()
{
    CDpcIrqlRaiser OnDpc;

    auto head = m_Chain.m_FirstInChain;

    if (head)
    {
        head->Release();
    }

    ParaNdis_DebugHistory(this, eHistoryLogOperation::hopNBLDestructor, m_NBL, 0, 0, 0);

    m_Buffers.ForEachDetached([this](CNB *NB) { CNB::Destroy(NB); });

    if (m_NBL)
    {
        auto NBL = DetachInternalObject();
        NETKVM_ASSERT(NET_BUFFER_LIST_NEXT_NBL(NBL) == nullptr);
        if (CallCompletionForNBL(m_Context, NBL))
        {
            m_ParentTXPath->CompleteOutstandingNBLChain(NBL);
        }
        else
        {
            m_ParentTXPath->CompleteOutstandingInternalNBL(NBL);
        }
    }

    m_History.ForEachDetached([this](NBLHistory *Entry) { NBLHistory::Destroy(Entry, m_Context->MiniportHandle); });
}

bool CNBL::ParsePriority()
{
    NDIS_NET_BUFFER_LIST_8021Q_INFO priorityInfo;
    priorityInfo.Value = m_Context->ulPriorityVlanSetting ? NET_BUFFER_LIST_INFO(m_NBL, Ieee8021QNetBufferListInfo)
                                                          : nullptr;

    if (!priorityInfo.TagHeader.VlanId)
    {
        priorityInfo.TagHeader.VlanId = m_Context->VlanId;
    }

    if (priorityInfo.TagHeader.CanonicalFormatId || !IsValidVlanId(m_Context, priorityInfo.TagHeader.VlanId))
    {
        DPrintf(0, "Discarded invalid priority tag %p", priorityInfo.Value);
        return false;
    }
    else if (priorityInfo.Value)
    {
        // ignore priority, if configured
        if (!IsPrioritySupported(m_Context))
        {
            priorityInfo.TagHeader.UserPriority = 0;
        }
        // ignore VlanId, if specified
        if (!IsVlanSupported(m_Context))
        {
            priorityInfo.TagHeader.VlanId = 0;
        }
        if (priorityInfo.Value)
        {
            m_TCI = static_cast<UINT16>(priorityInfo.TagHeader.UserPriority << 13 | priorityInfo.TagHeader.VlanId);
            DPrintf(1, "Populated priority tag %p", priorityInfo.Value);
        }
    }

    return true;
}

void CNBL::RegisterNB(CNB *NB)
{
    m_Buffers.PushBack(NB);
    m_BuffersNumber++;
}

void CNBL::RegisterMappedNB(CNB *NB)
{
    UNREFERENCED_PARAMETER(NB);
    if (m_BuffersNumber == (ULONG)m_BuffersMapped.AddRef())
    {
        AddHistory(__FUNCTION__, "", "", NULL, "Buffers", m_BuffersNumber);
        m_ParentTXPath->NBLMappingDone(this);
    }
}

bool CNBL::ParseBuffers()
{
    m_MaxDataLength = 0;
    CAllocationHelper<CNB> *pNBAllocator = this;

    for (auto NB = NET_BUFFER_LIST_FIRST_NB(m_NBL); NB != nullptr; NB = NET_BUFFER_NEXT_NB(NB))
    {
        CNB *NBHolder = new (pNBAllocator) CNB(NB, this, m_Context, pNBAllocator);
        pNBAllocator = m_NBAllocator;
        if (!NBHolder || !NBHolder->IsValid())
        {
            return false;
        }
        RegisterNB(NBHolder);
        m_MaxDataLength = max(m_MaxDataLength, NBHolder->GetDataLength());
    }

    if (m_MaxDataLength == 0)
    {
        DPrintf(0, "Empty NBL (%p) dropped", m_NBL);
        return false;
    }

    return true;
}

bool CNBL::NeedsLSO()
{
    return m_MaxDataLength > m_Context->MaxPacketSize.nMaxFullSizeOS;
}

bool CNBL::FitsLSO()
{
    return (m_MaxDataLength <= PARANDIS_MAX_LSO_SIZE + LsoTcpHeaderOffset() + MAX_TCP_HEADER_SIZE);
}

bool CNBL::ParseLSO()
{
    NETKVM_ASSERT(IsLSO());

    if (m_LsoInfo.LsoV1Transmit.Type != NDIS_TCP_LARGE_SEND_OFFLOAD_V1_TYPE &&
        m_LsoInfo.LsoV2Transmit.Type != NDIS_TCP_LARGE_SEND_OFFLOAD_V2_TYPE)
    {
        return false;
    }

    if (NeedsLSO() && (!m_LsoInfo.LsoV2Transmit.MSS || !m_LsoInfo.LsoV2Transmit.TcpHeaderOffset))
    {
        return false;
    }

    if (!FitsLSO())
    {
        return false;
    }

    if (!LsoTcpHeaderOffset() != !MSS())
    {
        return false;
    }

    if ((!m_Context->Offload.flags.fTxLso || !m_Context->bOffloadv4Enabled) &&
        m_LsoInfo.LsoV2Transmit.IPVersion == NDIS_TCP_LARGE_SEND_OFFLOAD_IPv4)
    {
        return false;
    }

    if (m_LsoInfo.LsoV2Transmit.Type == NDIS_TCP_LARGE_SEND_OFFLOAD_V2_TYPE &&
        m_LsoInfo.LsoV2Transmit.IPVersion == NDIS_TCP_LARGE_SEND_OFFLOAD_IPv6 &&
        (!m_Context->Offload.flags.fTxLsov6 || !m_Context->bOffloadv6Enabled))
    {
        return false;
    }

    return true;
}

#if PARANDIS_SUPPORT_USO
bool CNBL::ParseUSO()
{
    if (m_MaxDataLength > PARANDIS_MAX_USO_SIZE + (ULONGLONG)m_UsoInfo.Transmit.UdpHeaderOffset + sizeof(UDPHeader))
    {
        return false;
    }
    if (!m_UsoInfo.Transmit.MSS || !m_UsoInfo.Transmit.UdpHeaderOffset)
    {
        return false;
    }
    if (m_UsoInfo.Transmit.IPVersion == NDIS_UDP_SEGMENTATION_OFFLOAD_IPV4 && !m_Context->Offload.flags.fUsov4)
    {
        return false;
    }
    if (m_UsoInfo.Transmit.IPVersion == NDIS_UDP_SEGMENTATION_OFFLOAD_IPV6 && !m_Context->Offload.flags.fUsov6)
    {
        return false;
    }
    return true;
}
#else
bool CNBL::ParseUSO()
{
    return false;
}
#endif

template <typename TClassPred, typename TOffloadPred, typename TSupportedPred>
bool CNBL::ParseCSO(TClassPred IsClass, TOffloadPred IsOffload, TSupportedPred IsSupported, LPSTR OffloadName)
{
    NETKVM_ASSERT(IsClass());
    UNREFERENCED_PARAMETER(IsClass);

    if (IsOffload())
    {
        if (!IsSupported())
        {
            DPrintf(0, "%s request when it is not supported", OffloadName);
#if FAIL_UNEXPECTED
            // ignore unexpected CS requests while this passes WHQL
            return false;
#endif
        }
    }
    return true;
}

bool CNBL::ParseOffloads()
{
    if (IsLSO())
    {
        if (!ParseLSO())
        {
            return false;
        }
    }
    else if (IsUSO())
    {
        if (!ParseUSO())
        {
            return false;
        }
    }
    else if (IsIP4CSO())
    {
        if (!ParseCSO([this]() -> bool { return IsIP4CSO(); },
                      [this]() -> bool { return m_CsoInfo.Transmit.TcpChecksum; },
                      [this]() -> bool {
                          return m_Context->Offload.flags.fTxTCPChecksum && m_Context->bOffloadv4Enabled;
                      },
                      "TCP4 CSO"))
        {
            return false;
        }
        else if (!ParseCSO([this]() -> bool { return IsIP4CSO(); },
                           [this]() -> bool { return m_CsoInfo.Transmit.UdpChecksum; },
                           [this]() -> bool {
                               return m_Context->Offload.flags.fTxUDPChecksum && m_Context->bOffloadv4Enabled;
                           },
                           "UDP4 CSO"))
        {
            return false;
        }

        if (!ParseCSO([this]() -> bool { return IsIP4CSO(); },
                      [this]() -> bool { return m_CsoInfo.Transmit.IpHeaderChecksum; },
                      [this]() -> bool {
                          return m_Context->Offload.flags.fTxIPChecksum && m_Context->bOffloadv4Enabled;
                      },
                      "IP4 CSO"))
        {
            return false;
        }
    }
    else if (IsIP6CSO())
    {
        if (!ParseCSO([this]() -> bool { return IsIP6CSO(); },
                      [this]() -> bool { return m_CsoInfo.Transmit.TcpChecksum; },
                      [this]() -> bool {
                          return m_Context->Offload.flags.fTxTCPv6Checksum && m_Context->bOffloadv6Enabled;
                      },
                      "TCP6 CSO"))
        {
            return false;
        }
        else if (!ParseCSO([this]() -> bool { return IsIP6CSO(); },
                           [this]() -> bool { return m_CsoInfo.Transmit.UdpChecksum; },
                           [this]() -> bool {
                               return m_Context->Offload.flags.fTxUDPv6Checksum && m_Context->bOffloadv6Enabled;
                           },
                           "UDP6 CSO"))
        {
            return false;
        }
    }

    return true;
}

void CNBL::StartMapping()
{
    CDpcIrqlRaiser OnDpc;

    AddRef();

    m_Buffers.ForEach([this](CNB *NB) {
        if (!NB->ScheduleBuildSGListForTx())
        {
            m_HaveFailedMappings = true;
            NB->MappingDone(nullptr);
        }
    });

    Release();
}

void CNBL::OnLastReferenceGone()
{
    Destroy(this);
}

CParaNdisTX::~CParaNdisTX()
{

    TPassiveSpinLocker LockedContext(m_Lock);
    CNBL *NBL = nullptr;

    NBL = m_SendQueue.Dequeue();

    while (NBL)
    {
        NBL->~CNBL();
        NBL = m_SendQueue.Dequeue();
    }

    DPrintf(1, "Pools state %d-> NB: %d, NBL: %d", m_queueIndex, m_nbPool.GetCount(), m_nblPool.GetCount());
    if (m_StateMachineRegistered)
    {
        m_Context->m_StateMachine.UnregisterFlow(m_StateMachine);
        m_StateMachineRegistered = false;
    }

    FreeExtraPages();
}

bool CParaNdisTX::Create(PPARANDIS_ADAPTER Context, UINT DeviceQueueIndex)
{
    m_Context = Context;
    m_queueIndex = (u16)DeviceQueueIndex;

    Context->m_StateMachine.RegisterFlow(m_StateMachine);
    m_StateMachineRegistered = true;

    m_nbPool.Create(Context->MiniportHandle);
    m_nblPool.Create(Context->MiniportHandle);

    CreatePath();

    if (!AllocateExtraPages())
    {
        FreeExtraPages();
        return false;
    }

    return m_VirtQueue.Create(DeviceQueueIndex,
                              &m_Context->IODevice,
                              m_Context->MiniportHandle,
                              m_Context->maxFreeTxDescriptors,
                              m_Context->nVirtioHeaderSize,
                              m_Context) &&
           m_SendQueue.Create(Context,
                              IsPowerOfTwo(m_Context->maxFreeTxDescriptors) ? 8 * m_Context->maxFreeTxDescriptors
                                                                            : PARANDIS_TX_LOCK_FREE_QUEUE_DEFAULT_SIZE);
}

void CParaNdisTX::CompleteOutstandingNBLChain(PNET_BUFFER_LIST NBL, ULONG Flags)
{
    ULONG NBLNum = ParaNdis_CountNBLs(NBL);

    DPrintf(3, "completing %d nbls", NBLNum);
    ParaNdis_CompleteNBLChain(m_Context->MiniportHandle, NBL, Flags);

    m_StateMachine.UnregisterOutstandingItems(NBLNum);
}

void CParaNdisTX::CompleteOutstandingInternalNBL(PNET_BUFFER_LIST NBL, BOOLEAN UnregisterOutstanding /*= TRUE*/)
{
    ULONG NBLNum = ParaNdis_CountNBLs(NBL);

    CGuestAnnouncePackets::NblCompletionCallback(NBL);

    if (UnregisterOutstanding)
    {
        m_StateMachine.UnregisterOutstandingItems(NBLNum);
    }
}

// called in DISPATCH (under read lock of RSS parameters)
// if called with internal (announce) NBL, the NBL is always single
void CParaNdisTX::Send(PNET_BUFFER_LIST NBL)
{
    PNET_BUFFER_LIST nextNBL = nullptr;
    NDIS_STATUS RejectionStatus = NDIS_STATUS_FAILURE;
    BOOLEAN CallCompletion = CallCompletionForNBL(m_Context, NBL);
    CRawCNBLList chain;
    ULONG count = ParaNdis_CountNBLs(NBL);

    ParaNdis_DebugHistory(NBL, eHistoryLogOperation::hopSend, NULL, count, 0, m_queueIndex);

    if (!m_StateMachine.RegisterOutstandingItems(count, &RejectionStatus))
    {
        if (CallCompletion)
        {
            ParaNdis_CompleteNBLChainWithStatus(m_Context->MiniportHandle, NBL, RejectionStatus);
        }
        else
        {
            CompleteOutstandingInternalNBL(NBL, FALSE);
        }
        return;
    }

    TraceNoPrefix(3, "%s:Q#%d NBL count = %d\n", __FUNCTION__, m_queueIndex, count);
    for (auto currNBL = NBL; currNBL != nullptr; currNBL = nextNBL)
    {
        nextNBL = NET_BUFFER_LIST_NEXT_NBL(currNBL);
        NET_BUFFER_LIST_NEXT_NBL(currNBL) = nullptr;

        auto NBLHolder = new (&m_nblPool) CNBL(currNBL, m_Context, *this, &m_nblPool, &m_nbPool);

        if (NBLHolder == nullptr)
        {
            currNBL->Status = NDIS_STATUS_RESOURCES;
            if (CallCompletion)
            {
                CompleteOutstandingNBLChain(currNBL);
            }
            else
            {
                CompleteOutstandingInternalNBL(NBL);
            }
            DPrintf(0, "ERROR: Failed to allocate CNBL instance");
            // regular NBL chain and first NBL failed to allocate memory
            // we'll refer next one as the first NBL in chain
            // not relevant for internal NBL, as it is always single
            if (currNBL == NBL)
            {
                NBL = nextNBL;
                CallCompletion = true;
            }
            continue;
        }

        if (NBLHolder->Prepare() && ParaNdis_IsTxRxPossible(m_Context))
        {
            if (currNBL == NBL)
            {
                LPCSTR title = nextNBL ? "NBL-Head" : "NBL-Single";
                NBLHolder->AddHistory(__FUNCTION__, "Created", title, currNBL);
                ParaNdis_DebugHistory(NBLHolder, eHistoryLogOperation::hopSendNBLRequest, currNBL, 1, 0, m_queueIndex);
            }
            else
            {
                NBLHolder->AddHistory(__FUNCTION__, "Created", "NBL-Chained", currNBL);
                ParaNdis_DebugHistory(NBLHolder, eHistoryLogOperation::hopSendNBLRequest, currNBL, 0, 0, m_queueIndex);
            }
            NBLHolder->SetInChain(NBL);
            chain.PushBack(NBLHolder);
        }
        else
        {
            NBLHolder->SetStatus(ParaNdis_ExactSendFailureStatus(m_Context));
            NBLHolder->Release();
            if (currNBL == NBL)
            {
                NBL = nextNBL;
                CallCompletion = true;
            }
        }
    }

    chain.ForEachDetached([](CNBL *nbl) { nbl->StartMapping(); });
}

void CParaNdisTX::NBLMappingDone(CNBL *NBLHolder)
{
    NETKVM_ASSERT(KeGetCurrentIrql() == DISPATCH_LEVEL);

    if (NBLHolder->MappingSucceeded() && m_VirtQueue.Alive())
    {
        UpdateTimestamp(m_AuditState.LastSendTime);
        m_SendQueue.Enqueue(NBLHolder);

        LONG isDpcThere = m_DpcWaiting;
        NBLHolder->AddHistory(__FUNCTION__, "", "", NULL, "Skip processing", isDpcThere);

        if (isDpcThere == 0)
        {
            DoPendingTasks(NBLHolder);
        }
    }
    else
    {
        DPrintf(0, "ERROR: one or more NBs failed to be mapped!");
        NBLHolder->SetStatus(NDIS_STATUS_FAILURE);
        m_Context->Statistics.ifOutErrors += NBLHolder->NumberOfBuffers();
        NBLHolder->UnsetInChain();
        NBLHolder->Release();
    }
}

bool CParaNdisTX::AllocateExtraPages()
{
    for (ULONG j = 0; j < MAX_PACKET_PAGES; j++)
    {
        auto Page = new (m_Context->MiniportHandle) CNdisSharedMemory();
        if (Page == nullptr)
        {
            return false;
        }
        Page->Initialize(m_Context->MiniportHandle);
        if (Page->Allocate(PAGE_SIZE))
        {
            m_ExtraPages.Push(Page);
        }
        else
        {
            DPrintf(0, "failed to allocate pages, [%d] pages have been allocated", j);
            return false;
        }
    }
    return true;
}

void CParaNdisTX::FreeExtraPages()
{
    m_ExtraPages.ForEachDetached([this](CNdisSharedMemory *e) {
        CNdisSharedMemory::Destroy(e, m_Context->MiniportHandle);
    });
}

bool CParaNdisTX::BorrowPages(CExtendedNBStorage *extraNBStorage, ULONG NeedPages)
{
    if (m_ExtraPages.GetCount() < NeedPages)
    {
        return false;
    }
    for (ULONG i = 0; i < NeedPages; i++)
    {
        auto pPage = m_ExtraPages.Pop();
        extraNBStorage->m_UsedPages[i] = pPage;
    }
    extraNBStorage->m_UsedPagesCount = NeedPages;
    return true;
}

void CParaNdisTX::ReturnPages(CExtendedNBStorage *extraNBStorage)
{
    if (extraNBStorage)
    {
        for (ULONG i = 0; i < extraNBStorage->m_UsedPagesCount; i++)
        {
            m_ExtraPages.Push(extraNBStorage->m_UsedPages[i]);
        }
        extraNBStorage->m_UsedPagesCount = 0;
    }
}

CNB *CNBL::PopMappedNB()
{
    m_MappedBuffersDetached.AddRef();
    return m_Buffers.Pop();
}
void CNBL::PushMappedNB(CNB *NB)
{
    m_MappedBuffersDetached.Release();
    m_Buffers.Push(NB);
}

void CNBL::NBComplete()
{
    // do we need 'detached' at all?
    // 2 simultaneous users may never see both
    // currentDone = m_BuffersNumber and detached = 0
    m_MappedBuffersDetached.Release();
    if (m_BuffersDone.AddRef() == (LONG)m_BuffersNumber)
    {
        AddHistory(__FUNCTION__, "", "", NULL, "Buffers", m_BuffersNumber);
        ParaNdis_DebugHistory(this,
                              eHistoryLogOperation::hopSendComplete,
                              m_NBL,
                              !m_Chain.m_FirstInChain,
                              0, // status
                              GetCurrentRefCounterUnsafe());
        CompleteInChain();
        m_AllNBCompleted = true;
    }
}

bool CNBL::IsSendDone()
{
    // m_Chain.m_FirstInChain = NULL for first NBL in chain, also for single NBL
    //    for single NBL m_NumChained = m_NumCompleted = 0
    //    for head of chain - m_NumChained >= 1
    // m_Chain.m_FirstInChain != NULL for any additional NBLs
    if (!m_AllNBCompleted)
    {
        return false;
    }
    CNBL *head = m_Chain.m_FirstInChain;
    if (head == nullptr)
    {
        return m_Chain.m_NumChained == m_Chain.m_NumCompleted;
    }
    return head->IsSendDone();
}

PNET_BUFFER_LIST CNBL::DetachInternalObject()
{
    m_Buffers.ForEach([this](CNB *NB) { NB->ReleaseResources(); });

    // do it for both LsoV1 and LsoV2
    if (IsLSO())
    {
        m_LsoInfo.LsoV1TransmitComplete.TcpPayload = m_TransferSize;
    }

    // Flush changes made in LSO structures
    NET_BUFFER_LIST_INFO(m_NBL, TcpLargeSendNetBufferListInfo) = m_LsoInfo.Value;

    auto Res = m_NBL;
    ParaNdis_DebugNBLOut(m_LogIndex, Res);
    m_NBL = nullptr;
    return Res;
}

NBLHistory::NBLHistory(LPCSTR Func,
                       LPCSTR Title,
                       LPCSTR ParameterMeaning,
                       PVOID Parameter,
                       LPCSTR ValueMeaning,
                       ULONG Value)
{
    m_Function = Func;
    m_Title = Title;
    m_ParameterMeaning = ParameterMeaning;
    m_Parameter = Parameter;
    m_ValueMeaning = ValueMeaning;
    m_Value = Value;
    m_CurrentCPU = KeGetCurrentProcessorNumber();
    UpdateTimestamp(m_Timestamp);
}

#if NBL_MAINTAIN_HISTORY
void CNBL::AddHistory(LPCSTR Func,
                      LPCSTR Title,
                      LPCSTR ParameterMeaning,
                      PVOID Parameter,
                      LPCSTR ValueMeaning,
                      ULONG Value)
{
    auto entry = new (m_Context->MiniportHandle) NBLHistory(Func,
                                                            Title,
                                                            ParameterMeaning,
                                                            Parameter,
                                                            ValueMeaning,
                                                            Value);
    if (entry)
    {
        m_History.PushBack(entry);
    }
}
#endif

static bool TestNBLChain(PNET_BUFFER_LIST RawNBL, LPCSTR Message)
{
    ULONG expected = 0;
    auto next = RawNBL;
    do
    {
        ULONG got = (ULONG)(ULONG_PTR)next->Scratch;
        if (got != expected)
        {
            if (Message)
            {
                TraceNoPrefix(0, "%s: Mismatch in chain: exp %d != %d\n", Message, expected, got);
            }
            return false;
        }
        expected++;
        next = NET_BUFFER_LIST_NEXT_NBL(next);
    } while (next);
    return true;
}

static void RepairNBLChain(PNET_BUFFER_LIST RawNBL)
{
    PNET_BUFFER_LIST *tail = &NET_BUFFER_LIST_NEXT_NBL(RawNBL);
    PNET_BUFFER_LIST prev, head, current;
    prev = head = current = NET_BUFFER_LIST_NEXT_NBL(RawNBL);
    ULONG expected = 1;
    *tail = nullptr;
    while (current)
    {
        ULONG got = (ULONG)(ULONG_PTR)current->Scratch;
        if (got != expected)
        {
            // go to next one
            prev = current;
            current = NET_BUFFER_LIST_NEXT_NBL(current);
            continue;
        }

        *tail = current;
        tail = &NET_BUFFER_LIST_NEXT_NBL(current);

        if (current == head)
        {
            // removed current which is a head
            head = NET_BUFFER_LIST_NEXT_NBL(current);
        }
        else
        {
            // removed current which is not a head
            NET_BUFFER_LIST_NEXT_NBL(prev) = NET_BUFFER_LIST_NEXT_NBL(current);
        }

        current = prev = head;
        *tail = nullptr;
        expected++;
    };
    if (head)
    {
        *tail = head;
    }
}

PNET_BUFFER_LIST CParaNdisTX::ProcessWaitingList(CRawCNBLList &completed)
{
    // TODO: keep NBLs order and check NBL chain when extracting NBLs

    PNET_BUFFER_LIST CompletedNBLs = nullptr;
    PNET_BUFFER_LIST *tail = &CompletedNBLs;

    // locked part under waiting list lock
    {
        TDPCSpinLocker LockedContext(m_WaitingListLock);

        completed.ForEachDetached([&](CNBL *NBL) { m_WaitingList.PushBack(NBL); });

        m_WaitingList.ForEachDetachedIf([](CNBL *NBL) { return NBL->IsSendDone(); },
                                        [&](CNBL *NBL) { completed.PushBack(NBL); });
    }
    // end of locked part under waiting list lock

    completed.ForEachDetached([&](CNBL *NBL) {
        if (!NBL->HasNBL())
        {
            ParaNdis_DebugHistory(NBL,
                                  eHistoryLogOperation::hopSendDone,
                                  NULL,
                                  0,
                                  m_queueIndex,
                                  NBL->GetCurrentRefCounterUnsafe());
            NBL->Release();
            return;
        }
        NBL->SetStatus(NDIS_STATUS_SUCCESS);
        auto RawNBL = NBL->DetachInternalObject();
        ParaNdis_DebugHistory(NBL,
                              eHistoryLogOperation::hopSendDone,
                              RawNBL,
                              0,
                              m_queueIndex,
                              NBL->GetCurrentRefCounterUnsafe());
        if (CallCompletionForNBL(m_Context, RawNBL))
        {
            static bool bTestNBLs = REPAIR_CHAIN;
            if (bTestNBLs && NET_BUFFER_LIST_NEXT_NBL(RawNBL))
            {
                // this is a chain
                RawNBL->Scratch = NULL;
                if (!TestNBLChain(RawNBL, nullptr))
                {
                    RepairNBLChain(RawNBL);
                    TestNBLChain(RawNBL, "Second pass");
                }
            }
            *tail = RawNBL;
            do
            {
                tail = &NET_BUFFER_LIST_NEXT_NBL(*tail);

            } while (*tail);
        }
        else
        {
            CompleteOutstandingInternalNBL(RawNBL);
        }
        NBL->Release();
    });

    return CompletedNBLs;
}

// call under TX lock
void CParaNdisTX::DropAllNBls(CRawCNBLList &Completed, NDIS_STATUS Status)
{
    while (HaveMappedNBLs())
    {
        CNBL *nbl = PopMappedNBL();
        nbl->SetStatus(Status);
        Completed.Push(nbl);
        while (nbl->HaveMappedBuffers())
        {
            CNB *nb = nbl->PopMappedNB();
            CNB::Destroy(nb);
            nbl->NBComplete();
        }
    }
}

void CParaNdisTX::Notify(SMNotifications message)
{
    CRawCNBList nbToFree;
    CRawCNBLList completedNBLs;

    __super::Notify(message);

    if (message != SMNotifications::SupriseRemoved)
    {
        // probably similar processing we'll do in case of reset
        return;
    }

    CDpcIrqlRaiser OnDpc;

    // pause follows after we return all the NBLs
    DoWithTXLock([&]() {
        //
        m_VirtQueue.ProcessTXCompletions(nbToFree, true);
        DropAllNBls(completedNBLs, NDIS_STATUS_SEND_ABORTED);
    });

    PostProcessPendingTask(nbToFree, completedNBLs);
}

// with lockless queue the implementation of NBL cancellation
// becomes very complicated: we do not have enumeration of
// queued NBLs, there is no synchronization with Enqueue() etc
// this functionality is optional, so leave it for now
void CParaNdisTX::CancelNBLs(PVOID CancelId)
{
    UNREFERENCED_PARAMETER(CancelId);
    DPrintf(0, "not supported");
}

// called with TX lock held
bool CParaNdisTX::RestartQueue()
{
    auto res = ParaNdis_SynchronizeWithInterrupt(m_Context, m_messageIndex, RestartQueueSynchronously, this) ? true
                                                                                                             : false;
    return res;
}

// called with TX lock held
// returns queue restart status
bool CParaNdisTX::SendMapped(bool IsInterrupt, CRawCNBLList &toWaitingList)
{
    bool SentOutSomeBuffers = false;
    bool bRestartStatus = false;
    bool HaveBuffers = true;

    if (ParaNdis_IsTxRxPossible(m_Context))
    {

        while (HaveBuffers && HaveMappedNBLs())
        {
            auto NBLHolder = PeekMappedNBL();

            if (NBLHolder->HaveMappedBuffers())
            {
                auto NBHolder = NBLHolder->PopMappedNB();
                auto result = m_VirtQueue.SubmitPacket(*NBHolder);

                switch (result)
                {
                    case SubmitTxPacketResult::SUBMIT_NO_PLACE_IN_QUEUE:
                        NBLHolder->PushMappedNB(NBHolder);
                        HaveBuffers = false;
                        // break the loop, allow to kick and free some buffers
                        break;

                    case SubmitTxPacketResult::SUBMIT_FAILURE:
                        __fallthrough;
                    case SubmitTxPacketResult::SUBMIT_SUCCESS:
                        __fallthrough;
                    case SubmitTxPacketResult::SUBMIT_PACKET_TOO_LARGE:
                        // if this NBL finished?
                        if (!NBLHolder->HaveMappedBuffers())
                        {
                            /* We use PeekMappedNBL method to get the current NBL
                             * that should be processed from the queue, when we finish
                             * sending all it's NBs, we should pop it from the queue.
                             */
                            PopMappedNBL();
                            toWaitingList.Push(NBLHolder);
                        }
                        else
                        {
                            // no, keep it in the queue
                        }

                        if (result == SubmitTxPacketResult::SUBMIT_SUCCESS)
                        {
                            SentOutSomeBuffers = true;
                        }
                        else
                        {
                            CNB::Destroy(NBHolder);
                            NBLHolder->NBComplete();
                        }
                        break;
                    default:
                        NETKVM_ASSERT(false);
                        break;
                }
            }
            else
            {
                NETKVM_ASSERT(false);
            }
        }
    }
    else
    {
        // link loss or removal
        DropAllNBls(toWaitingList, ParaNdis_ExactSendFailureStatus(m_Context));
    }

    if (IsInterrupt)
    {
        bRestartStatus = RestartQueue();
    }

    if (SentOutSomeBuffers || !HaveBuffers)
    {
        m_VirtQueue.Kick();
    }

    return bRestartStatus;
}

void CParaNdisTX::PostProcessPendingTask(CRawCNBList &nbToFree, CRawCNBLList &completedNBLs)
{
    PNET_BUFFER_LIST pNBLReturnNow = nullptr;
    if (!nbToFree.IsEmpty() || !completedNBLs.IsEmpty())
    {
        nbToFree.ForEachDetached([](CNB *NB) {
            CNBL *NBL = NB->GetParentNBL();
            CNB::Destroy(NB);
            NBL->NBComplete();
        });

        pNBLReturnNow = ProcessWaitingList(completedNBLs);
        if (pNBLReturnNow)
        {
            CompleteOutstandingNBLChain(pNBLReturnNow, NDIS_SEND_COMPLETE_FLAGS_DISPATCH_LEVEL);
        }
    }
}

// should be called on DISPATCH with no locks held
// or on PASSIVE
// returns flags:
// 1 - waiting list is not empty
// 2 - send queue is not empty
ULONG CParaNdisTX::IsStuck()
{
    ULONG ret = 0;
    DoWithTXLock([&]() {
        TDPCSpinLocker LockedContext(m_WaitingListLock);
        ret |= m_WaitingList.IsEmpty() ? 0 : 1;
        ret |= m_SendQueue.IsEmpty() ? 0 : 2;
    });
    return ret;
}

void CParaNdisTX::CheckStuckPackets(ULONG GraceTimeMillies)
{
    // skip if less than N sec ago TX was serviced
    ULONGLONG &timeStamp = m_AuditState.LastAudit;
    ULONGLONG diff, threshold = (ULONGLONG)GraceTimeMillies * 1000L * 10L;

    UpdateTimestamp(timeStamp);
    diff = timeStamp - m_AuditState.LastTxProcess;
    if (diff < threshold)
    {
        return;
    }
    DPrintf(3, "TXQ#%d", m_queueIndex);

    ULONG flags = IsStuck();
    if (flags)
    {
        DPrintf(0, "STUCK condition=%d detected TXQ#%d", flags, m_queueIndex);
        m_AuditState.Stucks++;
        DoPendingTasks(NULL);
        flags = IsStuck();
        DPrintf(0, "On recovery: condition=%d TXQ#%d", flags, m_queueIndex);
        m_AuditState.Recovered += flags == 0;
    }
}

bool CParaNdisTX::DoPendingTasks(CNBL *nblHolder)
{
    bool bRestartQueueStatus = false;
    bool bFromDpc = nblHolder == nullptr;
    CRawCNBList nbToFree;
    CRawCNBLList completedNBLs;

    if (bFromDpc)
    {
        m_DpcWaiting.AddRef();
    }

    DoWithTXLock([&]() {
        m_VirtQueue.ProcessTXCompletions(nbToFree);

        if (bFromDpc)
        {
            m_DpcWaiting.Release();
        }

        if (bFromDpc || 0 == (LONG)m_DpcWaiting)
        {
            bRestartQueueStatus = SendMapped(bFromDpc, completedNBLs);
            if (bRestartQueueStatus)
            {
                // we can enter here only when we called from DPC
                // if we can't enable interrupts on queue right now,
                // we can retrieve completed packets and try again
                m_VirtQueue.ProcessTXCompletions(nbToFree);
                bRestartQueueStatus = SendMapped(true, completedNBLs);
            }
            UpdateTimestamp(m_AuditState.LastTxProcess);
        }
        else
        {
            // the call initiated by Send(), we can give up
            // and let pending DPC do the job instead of wait
        }
    });

    PostProcessPendingTask(nbToFree, completedNBLs);

    return bRestartQueueStatus;
}

void CNB::MappingDone(PSCATTER_GATHER_LIST SGL)
{
    m_SGL = SGL;
    m_ParentNBL->RegisterMappedNB(this);
}

CNB::~CNB()
{
    NETKVM_ASSERT(KeGetCurrentIrql() == DISPATCH_LEVEL);

    if (m_SGL != nullptr)
    {
        NdisMFreeNetBufferSGList(m_Context->DmaHandle, m_SGL, m_NB);
    }
    if (m_ExtraNBStorage)
    {
        // for unknown case it was not freed before
        ReturnPages();
    }
}

void CNB::ReleaseResources()
{
    if (m_SGL != nullptr)
    {
        NdisMFreeNetBufferSGList(m_Context->DmaHandle, m_SGL, m_NB);
        m_SGL = nullptr;
    }
}

bool CNB::ScheduleBuildSGListForTx()
{
    NETKVM_ASSERT(KeGetCurrentIrql() == DISPATCH_LEVEL);

    return NdisMAllocateNetBufferSGList(m_Context->DmaHandle, m_NB, this, NDIS_SG_LIST_WRITE_TO_DEVICE, nullptr, 0) ==
           NDIS_STATUS_SUCCESS;
}

void CNB::PopulateIPLength(IPHeader *IpHeader, USHORT IpLength) const
{
    if ((IpHeader->v4.ip_verlen & 0xF0) == 0x40)
    {
        if (!IpHeader->v4.ip_length)
        {
            IpHeader->v4.ip_length = swap_short(IpLength);
        }
    }
    else if ((IpHeader->v6.ip6_ver_tc & 0xF0) == 0x60)
    {
        if (!IpHeader->v6.ip6_payload_len)
        {
            IpHeader->v6.ip6_payload_len = swap_short(IpLength - IPV6_HEADER_MIN_SIZE);
        }
    }
    else
    {
        DPrintf(0, "ERROR: Bad version of IP header!");
    }
}

void CNB::SetupLSO(virtio_net_hdr *VirtioHeader, PVOID IpHeader, ULONG EthPayloadLength) const
{
    PopulateIPLength(reinterpret_cast<IPHeader *>(IpHeader), static_cast<USHORT>(EthPayloadLength));

    tTcpIpPacketParsingResult packetReview;
    packetReview = ParaNdis_CheckSumVerifyFlat(reinterpret_cast<IPv4Header *>(IpHeader),
                                               EthPayloadLength,
                                               tPacketOffloadRequest::pcrIpChecksum | tPacketOffloadRequest::pcrFixIPChecksum | tPacketOffloadRequest::pcrTcpChecksum |
                                                                                                                                                   tPacketOffloadRequest::pcrFixPHChecksum,
                                               FALSE,
                                               __FUNCTION__);

    if (packetReview.xxpCheckSum == ppResult::ppresPCSOK || packetReview.fixedXxpCS)
    {
        auto IpHeaderOffset = m_Context->Offload.ipHeaderOffset;
        auto VHeader = static_cast<virtio_net_hdr *>(VirtioHeader);
        auto PriorityHdrLen = (m_ParentNBL->TCI() != 0) ? ETH_PRIORITY_HEADER_SIZE : 0;

        VHeader->flags = VIRTIO_NET_HDR_F_NEEDS_CSUM;
        VHeader->gso_type = packetReview.ipStatus == ppResult::ppresIPV4 ? VIRTIO_NET_HDR_GSO_TCPV4
                                                                         : VIRTIO_NET_HDR_GSO_TCPV6;
        VHeader->hdr_len = (USHORT)(packetReview.XxpIpHeaderSize + IpHeaderOffset + PriorityHdrLen);
        VHeader->gso_size = (USHORT)m_ParentNBL->MSS();
        VHeader->csum_start = (USHORT)(m_ParentNBL->TCPHeaderOffset() + PriorityHdrLen);
        VHeader->csum_offset = TCP_CHECKSUM_OFFSET;
    }
}

void CNB::SetupUSO(virtio_net_hdr *VirtioHeader, PVOID IpHeader, ULONG EthPayloadLength) const
{
    PopulateIPLength(reinterpret_cast<IPHeader *>(IpHeader), static_cast<USHORT>(EthPayloadLength));

    tTcpIpPacketParsingResult packetReview;
    packetReview = ParaNdis_CheckSumVerifyFlat(reinterpret_cast<IPv4Header *>(IpHeader),
                                               EthPayloadLength,
                                               tPacketOffloadRequest::pcrIpChecksum | tPacketOffloadRequest::pcrFixIPChecksum | tPacketOffloadRequest::pcrUdpChecksum |
                                                                                                                                                   tPacketOffloadRequest::pcrFixPHChecksum,
                                               FALSE,
                                               __FUNCTION__);

    if (packetReview.xxpCheckSum == ppResult::ppresPCSOK || packetReview.fixedXxpCS)
    {
        auto IpHeaderOffset = m_Context->Offload.ipHeaderOffset;
        auto VHeader = static_cast<virtio_net_hdr *>(VirtioHeader);
        auto PriorityHdrLen = (m_ParentNBL->TCI() != 0) ? ETH_PRIORITY_HEADER_SIZE : 0;

        VHeader->flags = VIRTIO_NET_HDR_F_NEEDS_CSUM;
        VHeader->gso_type = VIRTIO_NET_HDR_GSO_UDP_L4;
        VHeader->hdr_len = (USHORT)(packetReview.XxpIpHeaderSize + IpHeaderOffset + PriorityHdrLen);
        VHeader->gso_size = (USHORT)m_ParentNBL->UsoMSS();
        VHeader->csum_start = (USHORT)(m_ParentNBL->UsoHeaderOffset() + PriorityHdrLen);
        VHeader->csum_offset = UDP_CHECKSUM_OFFSET;
        DPrintf(0, "mss %d, hdr %d, total %d", VHeader->gso_size, VHeader->hdr_len, GetDataLength());
    }
}

USHORT CNB::QueryL4HeaderOffset(PVOID PacketData, ULONG IpHeaderOffset) const
{
    USHORT Res;
    auto ppr = ParaNdis_ReviewIPPacket(RtlOffsetToPointer(PacketData, IpHeaderOffset),
                                       GetDataLength(),
                                       FALSE,
                                       __FUNCTION__);
    if (ppr.ipStatus != ppResult::ppresNotIP)
    {
        Res = static_cast<USHORT>(IpHeaderOffset + ppr.ipHeaderSize);
    }
    else
    {
        DPrintf(0, "ERROR: NOT an IP packet - expected troubles!");
        Res = 0;
    }
    return Res;
}

void CNB::SetupCSO(virtio_net_hdr *VirtioHeader, ULONG L4HeaderOffset) const
{
    u16 PriorityHdrLen = m_ParentNBL->TCI() ? ETH_PRIORITY_HEADER_SIZE : 0;

    VirtioHeader->flags = VIRTIO_NET_HDR_F_NEEDS_CSUM;
    VirtioHeader->csum_start = static_cast<u16>(L4HeaderOffset) + PriorityHdrLen;
    VirtioHeader->csum_offset = m_ParentNBL->IsTcpCSO() ? TCP_CHECKSUM_OFFSET : UDP_CHECKSUM_OFFSET;
}

void CNB::DoIPHdrCSO(PVOID IpHeader, ULONG EthPayloadLength) const
{
    ParaNdis_CheckSumVerifyFlat(IpHeader,
                                EthPayloadLength,
                                tPacketOffloadRequest::pcrIpChecksum | tPacketOffloadRequest::pcrFixIPChecksum,
                                FALSE,
                                __FUNCTION__);
}

NBMappingStatus CNB::FillDescriptorSGList(CTXDescriptor &Descriptor, ULONG ParsedHeadersLength)
{
    if (!Descriptor.SetupHeaders(ParsedHeadersLength))
    {
        return NBMappingStatus::FAILURE;
    }
    if (Descriptor.HasRoom(m_SGL->NumberOfElements))
    {
        return MapDataToVirtioSGL(Descriptor, ParsedHeadersLength + NET_BUFFER_DATA_OFFSET(m_NB));
    }
    else
    {
        auto res = AllocateAndFillCopySGL(ParsedHeadersLength);
        if (res != NBMappingStatus::SUCCESS)
        {
            return res;
        }
        return MapCopyDataToVirtioSGL(Descriptor);
    }
}

NBMappingStatus CNB::MapDataToVirtioSGL(CTXDescriptor &Descriptor, ULONG Offset) const
{
    for (ULONG i = 0; i < m_SGL->NumberOfElements; i++)
    {
        if (Offset < m_SGL->Elements[i].Length)
        {
            PHYSICAL_ADDRESS PA;
            PA.QuadPart = m_SGL->Elements[i].Address.QuadPart + Offset;

            if (!Descriptor.AddDataChunk(PA, m_SGL->Elements[i].Length - Offset))
            {
                return NBMappingStatus::FAILURE;
            }

            Offset = 0;
        }
        else
        {
            Offset -= m_SGL->Elements[i].Length;
        }
    }

    return NBMappingStatus::SUCCESS;
}

bool CNB::CopyHeaders(PVOID Destination, ULONG MaxSize, ULONG &HeadersLength, ULONG &L4HeaderOffset) const
{
    HeadersLength = 0;
    L4HeaderOffset = 0;

    if (m_ParentNBL->IsLSO() || m_ParentNBL->IsTcpCSO())
    {
        L4HeaderOffset = m_ParentNBL->TCPHeaderOffset();
        HeadersLength = L4HeaderOffset + sizeof(TCPHeader);
        Copy(Destination, HeadersLength);
    }
    else if (m_ParentNBL->IsUSO())
    {
        L4HeaderOffset = m_ParentNBL->UsoHeaderOffset();
        HeadersLength = L4HeaderOffset + sizeof(UDPHeader);
        HeadersLength = Copy(Destination, HeadersLength);
    }
    else if (m_ParentNBL->IsUdpCSO())
    {
        HeadersLength = Copy(Destination, MaxSize);
        L4HeaderOffset = QueryL4HeaderOffset(Destination, m_Context->Offload.ipHeaderOffset);
    }
    else if (m_ParentNBL->IsIPHdrCSO())
    {
        HeadersLength = Copy(Destination, MaxSize);
        L4HeaderOffset = QueryL4HeaderOffset(Destination, m_Context->Offload.ipHeaderOffset);
    }
    else
    {
        HeadersLength = ETH_HEADER_SIZE;
        Copy(Destination, HeadersLength);
    }

    return (HeadersLength <= MaxSize);
}

void CNB::BuildPriorityHeader(PETH_HEADER EthHeader, PVLAN_HEADER VlanHeader) const
{
    VlanHeader->TCI = RtlUshortByteSwap(m_ParentNBL->TCI());

    if (VlanHeader->TCI != 0)
    {
        VlanHeader->EthType = EthHeader->EthType;
        EthHeader->EthType = RtlUshortByteSwap(PRIO_HEADER_ETH_TYPE);
    }
}

void CNB::PrepareOffloads(virtio_net_hdr *VirtioHeader,
                          PVOID IpHeader,
                          ULONG EthPayloadLength,
                          ULONG L4HeaderOffset) const
{
    *VirtioHeader = {};

    if (m_ParentNBL->IsLSO())
    {
        SetupLSO(VirtioHeader, IpHeader, EthPayloadLength);
    }
    else if (m_ParentNBL->IsUSO())
    {
        SetupUSO(VirtioHeader, IpHeader, EthPayloadLength);
    }
    else if (m_ParentNBL->IsTcpCSO() || m_ParentNBL->IsUdpCSO())
    {
        SetupCSO(VirtioHeader, L4HeaderOffset);
    }

    if (m_ParentNBL->IsIPHdrCSO())
    {
        DoIPHdrCSO(IpHeader, EthPayloadLength);
    }
}

void CNB::Report(int level, bool Success)
{
    DPrintf(level,
            "%s packet of %d:%d bytes:frag (NBL of %d)",
            Success ? "OK" : "Failed",
            GetDataLength(),
            GetSGLLength(),
            m_ParentNBL->NumberOfBuffers());
    if (!Success)
    {
        m_Context->extraStatistics.droppedTxPackets++;
    }
}

NBMappingStatus CNB::BindToDescriptor(CTXDescriptor &Descriptor)
{
    if (m_SGL == nullptr)
    {
        return NBMappingStatus::FAILURE;
    }

    Descriptor.SetNB(this);

    auto &HeadersArea = Descriptor.HeadersAreaAccessor();
    auto EthHeaders = HeadersArea.EthHeadersAreaVA();
    ULONG HeadersLength;
    ULONG L4HeaderOffset;

    if (!CopyHeaders(EthHeaders, HeadersArea.MaxEthHeadersSize(), HeadersLength, L4HeaderOffset))
    {
        return NBMappingStatus::FAILURE;
    }

    BuildPriorityHeader(HeadersArea.EthHeader(), HeadersArea.VlanHeader());
    PrepareOffloads(HeadersArea.VirtioHeader(),
                    HeadersArea.IPHeaders(),
                    GetDataLength() - m_Context->Offload.ipHeaderOffset,
                    L4HeaderOffset);

    return FillDescriptorSGList(Descriptor, HeadersLength);
}

ULONG CNB::Copy(PVOID Dst, ULONG Length) const
{
    ULONG CurrOffset = NET_BUFFER_CURRENT_MDL_OFFSET(m_NB);
    ULONG Copied = 0;

    Length = min(Length, NET_BUFFER_DATA_LENGTH(m_NB));

    for (PMDL CurrMDL = NET_BUFFER_CURRENT_MDL(m_NB); CurrMDL != nullptr && Copied < Length; CurrMDL = CurrMDL->Next)
    {
        ULONG CurrLen;
        PVOID CurrAddr = nullptr;

#if NDIS_SUPPORT_NDIS620
        NdisQueryMdl(CurrMDL, &CurrAddr, &CurrLen, MM_PAGE_PRIORITY(LowPagePriority | MdlMappingNoExecute));
#else
        NdisQueryMdl(CurrMDL, &CurrAddr, &CurrLen, MM_PAGE_PRIORITY(LowPagePriority));
#endif

        if (CurrAddr == nullptr)
        {
            break;
        }

        CurrLen = min(CurrLen - CurrOffset, Length - Copied);

        NdisMoveMemory(RtlOffsetToPointer(Dst, Copied), RtlOffsetToPointer(CurrAddr, CurrOffset), CurrLen);

        Copied += CurrLen;
        CurrOffset = 0;
    }

    return Copied;
}

/*
 * Copies data from an MDL chain to a destination buffer.
 *
 * @param Dst Pointer to the destination buffer where the data will be copied to.
 * @param Length The maximum number of bytes to copy.
 * @param Source Reference to a pointer to the current MDL in the chain. This MDL pointer
 *               will be updated as the function traverses the MDL chain.
 * @param Offset Reference to the current offset within the current MDL. This offset will
 *               be updated as data is copied from the MDL.
 * @return The total number of bytes copied to the destination buffer.
 *
 * Note: The function updates 'Source' to point to the next MDL in the chain after the
 *       current MDL has been processed. If the end of the current MDL is reached before
 *       copying 'Length' bytes, 'Source' will point to the next MDL for further copying.
 *       'Offset' is updated to reflect the position within the current MDL after the copy
 *       operation. If the end of the MDL is reached, 'Offset' will be set to zero to indicate
 *       the start of the next MDL in the chain.
 */
ULONG CNB::CopyFromMdlChain(PVOID Dst, ULONG Length, PMDL &Source, ULONG &Offset)
{
    /*Skip over MDLs until we reach the one that contains the offset we're interested in*/
    while (Source && MmGetMdlByteCount(Source) < Offset)
    {
        Offset -= MmGetMdlByteCount(Source);
        Source = Source->Next;
    }

    ULONG Copied = 0;

    while (Source && Copied < Length)
    {
        ULONG CurrLen;
        PVOID CurrAddr = nullptr;

        NdisQueryMdl(Source, &CurrAddr, &CurrLen, MM_PAGE_PRIORITY(LowPagePriority | MdlMappingNoExecute));

        if (CurrAddr == nullptr)
        {
            break;
        }

        ULONG toCopyNow = min(CurrLen - Offset, Length - Copied);
        NdisMoveMemory(RtlOffsetToPointer(Dst, Copied), RtlOffsetToPointer(CurrAddr, Offset), toCopyNow);

        if (Length - Copied < CurrLen - Offset)
        {
            /*We have not reached the end of the current MDL and have copied the required data*/
            Copied += toCopyNow;
            Offset += toCopyNow;
            break;
        }
        else
        {
            /*We have reached the end of the current MDL*/
            Copied += toCopyNow;
            Offset = 0;
            Source = Source->Next;
        }
    }

    return Copied;
}

void CNB::ReturnPages()
{
    if (m_ExtraNBStorage)
    {
        m_ParentNBL->GetParentTXPath()->ReturnPages(m_ExtraNBStorage);
        NdisFreeMemory(m_ExtraNBStorage, sizeof(CExtendedNBStorage), 0);
    }
    m_ExtraNBStorage = nullptr;
}

NBMappingStatus CNB::AllocateAndFillCopySGL(ULONG ParsedHeadersLength)
{
    // Calculate the number of pages needed for data, excluding protocol header length.
    ULONG DataLength = GetDataLength() - ParsedHeadersLength;
    ULONG Pages = (DataLength + PAGE_SIZE - 1) / PAGE_SIZE;

    m_ExtraNBStorage = (CExtendedNBStorage *)ParaNdis_AllocateMemory(m_Context, sizeof(CExtendedNBStorage));

    if (m_ExtraNBStorage == nullptr)
    {
        DPrintf(0, "ExtendedNBStorage allocation failed");
        return NBMappingStatus::FAILURE;
    }

    RtlZeroMemory(m_ExtraNBStorage, sizeof(CExtendedNBStorage));
    if (!m_ParentNBL->GetParentTXPath()->BorrowPages(m_ExtraNBStorage, Pages))
    {
        NdisFreeMemory(m_ExtraNBStorage, sizeof(CExtendedNBStorage), 0);
        m_ExtraNBStorage = NULL;
        return NBMappingStatus::NO_RESOURCE;
    }

    PMDL mdl = NET_BUFFER_CURRENT_MDL(m_NB);
    ULONG DataOffset = ParsedHeadersLength + NET_BUFFER_CURRENT_MDL_OFFSET(m_NB);
    for (ULONG i = 0; i < Pages; i++)
    {
        PVOID pVirtualAddress = m_ExtraNBStorage->m_UsedPages[i]->GetVA();
        ULONG toCopyNow = (i == Pages - 1) ? DataLength : PAGE_SIZE;
        ULONG Copied = CopyFromMdlChain(pVirtualAddress, toCopyNow, mdl, DataOffset);
        if (Copied != toCopyNow)
        {
            DPrintf(0, "copy failed! expected %lu, copied %lu bytes", toCopyNow, Copied);
            return NBMappingStatus::FAILURE;
        }
        m_ExtraNBStorage->m_Elements[i].Address = m_ExtraNBStorage->m_UsedPages[i]->GetPA();
        m_ExtraNBStorage->m_Elements[i].Length = toCopyNow;
        DataLength -= toCopyNow;
    }
    return NBMappingStatus::SUCCESS;
}

NBMappingStatus CNB::MapCopyDataToVirtioSGL(CTXDescriptor &Descriptor) const
{
    for (ULONG i = 0; i < m_ExtraNBStorage->m_UsedPagesCount; i++)
    {
        if (!Descriptor.AddDataChunk(m_ExtraNBStorage->m_Elements[i].Address, m_ExtraNBStorage->m_Elements[i].Length))
        {
            return NBMappingStatus::FAILURE;
        }
    }
    m_Context->extraStatistics.copiedTxPackets++;
    return NBMappingStatus::SUCCESS;
}