From 4425bd8db392a50eb1d0734bb3e7ab33927b4885 Mon Sep 17 00:00:00 2001
From: Serge Gautherie <32623169+SergeGautherie@users.noreply.github.com>
Date: Mon, 1 Jun 2020 13:17:29 +0200
Subject: [PATCH] [CSRSRV] CsrSetProcessSecurity(): Check 1st
 NtQueryInformationToken() result (#2862)

Also:
* Add 1 NtClose(hToken), in an error case.
* Do not call RtlFreeHeap(..., ..., NULL).

Follow-up to #2857.
---
 subsystems/win32/csrsrv/init.c | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/subsystems/win32/csrsrv/init.c b/subsystems/win32/csrsrv/init.c
index 49df131a862..69e9cfdea00 100644
--- a/subsystems/win32/csrsrv/init.c
+++ b/subsystems/win32/csrsrv/init.c
@@ -74,12 +74,18 @@ CsrSetProcessSecurity(VOID)
     if (!NT_SUCCESS(Status)) goto Quickie;
 
     /* Get the Token User Length */
-    NtQueryInformationToken(hToken, TokenUser, NULL, 0, &Length);
+    Status = NtQueryInformationToken(hToken, TokenUser, NULL, 0, &Length);
+    if (Status != STATUS_BUFFER_TOO_SMALL)
+    {
+        NtClose(hToken);
+        goto Quickie;
+    }
 
     /* Allocate space for it */
     TokenInfo = RtlAllocateHeap(CsrHeap, HEAP_ZERO_MEMORY, Length);
     if (!TokenInfo)
     {
+        NtClose(hToken);
         Status = STATUS_NO_MEMORY;
         goto Quickie;
     }
@@ -153,7 +159,7 @@ CsrSetProcessSecurity(VOID)
     /* Free the memory and return */
 Quickie:
     if (ProcSd) RtlFreeHeap(CsrHeap, 0, ProcSd);
-    RtlFreeHeap(CsrHeap, 0, TokenInfo);
+    if (TokenInfo) RtlFreeHeap(CsrHeap, 0, TokenInfo);
     return Status;
 }