From f8feaa243bf787e51378348ecd0f399613476189 Mon Sep 17 00:00:00 2001
From: Lev Stipakov <lev@openvpn.net>
Date: Mon, 25 Sep 2023 15:09:19 +0200
Subject: [PATCH] Improve pool memory management

 - Handle case when IoAllocateMdl() and ExAllocatePool2()
 return NULL

 - Set max pool size (in-flight packets) to 100'000.

This might help with https://github.com/OpenVPN/ovpn-dco-win/issues/47

By playing with Quick Edit in console window (which makes execution
stop) where iPerf3 was running I was able to get pool size growing
to unreasonable values (over 50k). My understanding is that normally it
should be <10k, but let's be on a safe side and cap it to 100k - in
kernel dump it was over 1mil and system ran out of memory.

Bump to 0.9.4.

Signed-off-by: Lev Stipakov <lev@openvpn.net>
---
 PropertySheet.props |  2 +-
 bufferpool.cpp      | 32 ++++++++++++++++++++++++++------
 2 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/PropertySheet.props b/PropertySheet.props
index 9d8be4b..a12508d 100644
--- a/PropertySheet.props
+++ b/PropertySheet.props
@@ -4,7 +4,7 @@
   <PropertyGroup Label="UserMacros">
     <OVPN_DCO_VERSION_MAJOR>0</OVPN_DCO_VERSION_MAJOR>
     <OVPN_DCO_VERSION_MINOR>9</OVPN_DCO_VERSION_MINOR>
-    <OVPN_DCO_VERSION_PATCH>3</OVPN_DCO_VERSION_PATCH>
+    <OVPN_DCO_VERSION_PATCH>4</OVPN_DCO_VERSION_PATCH>
   </PropertyGroup>
   <PropertyGroup />
   <ItemDefinitionGroup>
diff --git a/bufferpool.cpp b/bufferpool.cpp
index e8b4357..926c783 100644
--- a/bufferpool.cpp
+++ b/bufferpool.cpp
@@ -29,6 +29,9 @@
 
 #define OVPN_BUFFER_HEADROOM 26 // we prepend TCP packet size (2 bytes) and crypto overhead (24 bytes)
 
+// good enough limit for in-flight packets
+constexpr auto MAX_POOL_SIZE = 100'000;
+
 struct OVPN_BUFFER_POOL_IMPL
 {
     LIST_ENTRY ListHead;
@@ -145,10 +148,19 @@ OvpnBufferPoolGet(OVPN_BUFFER_POOL handle, POOL_ENTRY** entry) {
     if (slist_entry) {
         *entry = CONTAINING_RECORD(slist_entry, POOL_ENTRY, PoolListEntry);
     } else {
+        if (pool->PoolSize > MAX_POOL_SIZE)
+        {
+            *entry = NULL;
+            LOG_ERROR("Pool size is too large", TraceLoggingValue(pool->Tag, "tag"), TraceLoggingValue(pool->PoolSize, "size"));
+            return;
+        }
         *entry = (POOL_ENTRY*)ExAllocatePool2(POOL_FLAG_NON_PAGED, pool->ItemSize, 'ovpn');
-        InterlockedIncrement(&pool->PoolSize);
-        if ((pool->PoolSize % 64) == 0) {
-            LOG_INFO("Pool size", TraceLoggingValue(pool->Tag, "tag"), TraceLoggingValue(pool->PoolSize, "size"));
+        if (*entry)
+        {
+            InterlockedIncrement(&pool->PoolSize);
+            if ((pool->PoolSize % 256) == 0) {
+                LOG_INFO("Pool size", TraceLoggingValue(pool->Tag, "tag"), TraceLoggingValue(pool->PoolSize, "size"));
+            }
         }
     }
 }
@@ -161,10 +173,17 @@ OvpnTxBufferPoolGet(OVPN_TX_BUFFER_POOL handle, OVPN_TX_BUFFER** buffer)
     if (*buffer == NULL)
         return STATUS_INSUFFICIENT_RESOURCES;
 
+    (*buffer)->Pool = handle;
+
     (*buffer)->Mdl = IoAllocateMdl(*buffer, ((OVPN_BUFFER_POOL_IMPL*)handle)->ItemSize, FALSE, FALSE, NULL);
-    MmBuildMdlForNonPagedPool((*buffer)->Mdl);
+    if (((*buffer)->Mdl) == NULL)
+    {
+        OvpnTxBufferPoolPut(*buffer);
+        *buffer = NULL;
+        return STATUS_INSUFFICIENT_RESOURCES;
+    }
 
-    (*buffer)->Pool = handle;
+    MmBuildMdlForNonPagedPool((*buffer)->Mdl);
 
     (*buffer)->Data = (*buffer)->Head + OVPN_BUFFER_HEADROOM;
     (*buffer)->Tail = (*buffer)->Data;
@@ -206,7 +225,8 @@ _Use_decl_annotations_
 VOID
 OvpnTxBufferPoolPut(OVPN_TX_BUFFER* buffer)
 {
-    IoFreeMdl(buffer->Mdl);
+    if (buffer->Mdl)
+        IoFreeMdl(buffer->Mdl);
 
     OvpnBufferPoolPut(buffer);
 }