[PATCH v3 2/2] nsiproxy: Only set the endpoint count from udp_endpoint_enumerate_all when appropriate.

[Tim Clem]

From: Tim Clem

Fixes heap corruption downstream in NsiEnumerateObjectsAllParametersEx.

Fixes a regression from a3f737f6143.
---
 dlls/nsiproxy.sys/udp.c | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/dlls/nsiproxy.sys/udp.c b/dlls/nsiproxy.sys/udp.c
index 1d790484d8f..897c90c131b 100644
--- a/dlls/nsiproxy.sys/udp.c
+++ b/dlls/nsiproxy.sys/udp.c
@@ -231,13 +231,17 @@ static NTSTATUS udp_endpoint_enumerate_all( void *key_data, UINT key_size, void
             *count = reply->count;
         else if (ret == STATUS_BUFFER_TOO_SMALL)
         {
-            *count = reply->count;
-            if (want_data)
+            if (!want_data)
             {
-                free( endpoints );
-                return STATUS_BUFFER_OVERFLOW;
+                /* If we were given buffers, the outgoing count must never be
+                   greater than the incoming one. If we weren't, the count
+                   should be set to the actual count. */
+                *count = reply->count;
+                return STATUS_SUCCESS;
             }
-            else return STATUS_SUCCESS;
+
+            free( endpoints );
+            return STATUS_BUFFER_OVERFLOW;
         }
     }
     SERVER_END_REQ;