From 750f40fada90ce22e8aceb4bf2775be8478bf63b Mon Sep 17 00:00:00 2001
From: Simon Rozman <simon@rozman.si>
Date: Mon, 22 May 2017 08:23:18 +0200
Subject: [PATCH] GetTokenInformation<>() template introduced

---
 include/WinStd/Win.h | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/include/WinStd/Win.h b/include/WinStd/Win.h
index 98bd0444..bdef1c3e 100644
--- a/include/WinStd/Win.h
+++ b/include/WinStd/Win.h
@@ -278,6 +278,13 @@ template<class _Elem, class _Traits, class _Ax> inline BOOL LookupAccountSidA(_I
 ///
 template<class _Elem, class _Traits, class _Ax> inline BOOL LookupAccountSidW(_In_opt_ LPCWSTR lpSystemName, _In_ PSID lpSid, _Out_opt_ std::basic_string<_Elem, _Traits, _Ax> *sName, _Out_opt_ std::basic_string<_Elem, _Traits, _Ax> *sReferencedDomainName, _Out_ PSID_NAME_USE peUse);
 
+///
+/// Retrieves a specified type of information about an access token. The calling process must have appropriate access rights to obtain the information.
+///
+/// \sa [GetTokenInformation function](https://msdn.microsoft.com/en-us/library/windows/desktop/aa446671.aspx)
+///
+template<class _Ty> inline BOOL GetTokenInformation(_In_ HANDLE TokenHandle, _In_ TOKEN_INFORMATION_CLASS TokenInformationClass, _Out_ std::unique_ptr<_Ty> &TokenInformation);
+
 /// @}
 
 #pragma once
@@ -1504,3 +1511,31 @@ inline BOOL LookupAccountSidW(_In_opt_ LPCWSTR lpSystemName, _In_ PSID lpSid, _O
 
     return FALSE;
 }
+
+
+template<class _Ty>
+inline BOOL GetTokenInformation(_In_ HANDLE TokenHandle, _In_ TOKEN_INFORMATION_CLASS TokenInformationClass, _Out_ std::unique_ptr<_Ty> &TokenInformation)
+{
+    BYTE szStackBuffer[WINSTD_STACK_BUFFER_BYTES/sizeof(BYTE)];
+    DWORD dwSize;
+
+    if (GetTokenInformation(TokenHandle, TokenInformationClass, szStackBuffer, sizeof(szStackBuffer), &dwSize)) {
+        // The stack buffer was big enough to retrieve complete data. Alloc and copy.
+        TokenInformation.reset((_Ty*)(new BYTE[dwSize / sizeof(BYTE)]));
+        if (!TokenInformation) {
+            SetLastError(ERROR_OUTOFMEMORY);
+            return FALSE;
+        }
+        memcpy(TokenInformation.get(), szStackBuffer, dwSize);
+        return TRUE;
+    } else if (GetLastError() == ERROR_INSUFFICIENT_BUFFER) {
+        // The stack buffer was too small to retrieve complete data. Alloc and retry.
+        TokenInformation.reset((_Ty*)(new BYTE[dwSize / sizeof(BYTE)]));
+        if (!TokenInformation) {
+            SetLastError(ERROR_OUTOFMEMORY);
+            return FALSE;
+        }
+        return GetTokenInformation(TokenHandle, TokenInformationClass, TokenInformation.get(), dwSize, &dwSize);
+    } else
+        return FALSE;
+}