diff --git a/src/core/unittest/RecvBufferTest.cpp b/src/core/unittest/RecvBufferTest.cpp index 0bcbdb256f..8f53fa6f76 100644 --- a/src/core/unittest/RecvBufferTest.cpp +++ b/src/core/unittest/RecvBufferTest.cpp @@ -55,6 +55,9 @@ struct RecvBuffer { PreallocChunk = (QUIC_RECV_CHUNK*)CXPLAT_ALLOC_NONPAGED( sizeof(QUIC_RECV_CHUNK) + AllocBufferLength, QUIC_POOL_RECVBUF); // Use the recv buffer pool tag as this memory is moved to the recv buffer. + if (PreallocChunk == nullptr) { + return QUIC_STATUS_OUT_OF_MEMORY; + } QuicRecvChunkInitialize(PreallocChunk, AllocBufferLength, (uint8_t*)(PreallocChunk + 1), FALSE); } printf("Initializing: [mode=%u,vlen=%u,alen=%u]\n", RecvMode, VirtualBufferLength, AllocBufferLength); @@ -75,11 +78,25 @@ struct RecvBuffer { CXPLAT_LIST_ENTRY ChunkList; CxPlatListInitializeHead(&ChunkList); AppOwnedBuffer = (uint8_t *)CXPLAT_ALLOC_NONPAGED(VirtualBufferLength, QUIC_POOL_TEST); + if (AppOwnedBuffer == nullptr) { + return QUIC_STATUS_OUT_OF_MEMORY; + } auto* Chunk = (QUIC_RECV_CHUNK *)CxPlatPoolAlloc(&AppBufferChunkPool); + if (Chunk == nullptr) { + CXPLAT_FREE(AppOwnedBuffer, QUIC_POOL_TEST); + AppOwnedBuffer = nullptr; + return QUIC_STATUS_OUT_OF_MEMORY; + } QuicRecvChunkInitialize(Chunk, AllocBufferLength, AppOwnedBuffer, TRUE); CxPlatListInsertHead(&ChunkList, &Chunk->Link); if (VirtualBufferLength > AllocBufferLength) { auto* Chunk2 = (QUIC_RECV_CHUNK *)CxPlatPoolAlloc(&AppBufferChunkPool); + if (Chunk2 == nullptr) { + CxPlatPoolFree(Chunk); + CXPLAT_FREE(AppOwnedBuffer, QUIC_POOL_TEST); + AppOwnedBuffer = nullptr; + return QUIC_STATUS_OUT_OF_MEMORY; + } QuicRecvChunkInitialize(Chunk2, VirtualBufferLength - AllocBufferLength, AppOwnedBuffer + AllocBufferLength, TRUE); CxPlatListInsertTail(&ChunkList, &Chunk2->Link); } diff --git a/src/perf/lib/Tcp.cpp b/src/perf/lib/Tcp.cpp index 180e6cdb23..d502e56d1a 100644 --- a/src/perf/lib/Tcp.cpp +++ b/src/perf/lib/Tcp.cpp @@ -704,6 +704,10 @@ bool TcpConnection::InitializeTls() { const uint32_t LocalTPLength = 2; uint8_t* LocalTP = (uint8_t*)CXPLAT_ALLOC_NONPAGED(CxPlatTlsTPHeaderSize + LocalTPLength, QUIC_POOL_TLS_TRANSPARAMS); + if (LocalTP == nullptr) { + WriteOutput("LocalTP allocation failed\n"); + return false; + } CxPlatZeroMemory(LocalTP, LocalTPLength); CXPLAT_TLS_CONFIG Config; diff --git a/src/platform/unittest/TlsTest.cpp b/src/platform/unittest/TlsTest.cpp index 0efc3b9418..dce27cab89 100644 --- a/src/platform/unittest/TlsTest.cpp +++ b/src/platform/unittest/TlsTest.cpp @@ -587,6 +587,9 @@ struct TlsTest : public ::testing::TestWithParam if (Context->ReceivedSessionTicket.Buffer == nullptr) { Context->ReceivedSessionTicket.Buffer = // N.B - Add one so we don't ever allocate zero bytes. (uint8_t*)CXPLAT_ALLOC_NONPAGED(TicketLength+1, QUIC_POOL_CRYPTO_RESUMPTION_TICKET); + if (Context->ReceivedSessionTicket.Buffer == nullptr) { + return FALSE; + } Context->ReceivedSessionTicket.Length = TicketLength; if (TicketLength != 0) { CxPlatCopyMemory(