Add primary token check for PhImpersonateToken

codingman · Jun 26, 2021 · 0ea21e7 · 0ea21e7
1 parent bdf8f84
commit 0ea21e7
Showing 1 changed file with 56 additions and 30 deletions.
diff --git a/phlib/native.c b/phlib/native.c
@@ -9905,44 +9905,70 @@ NTSTATUS PhImpersonateToken(
     )
 {
     NTSTATUS status;
-    SECURITY_QUALITY_OF_SERVICE securityService;
-    OBJECT_ATTRIBUTES objectAttributes;
-    HANDLE tokenHandle;
-
-    InitializeObjectAttributes(
-        &objectAttributes,
-        NULL,
-        0,
-        NULL,
-        NULL
-        );
-
-    securityService.Length = sizeof(SECURITY_QUALITY_OF_SERVICE);
-    securityService.ImpersonationLevel = SecurityImpersonation;
-    securityService.ContextTrackingMode = SECURITY_DYNAMIC_TRACKING;
-    securityService.EffectiveOnly = FALSE;
-    objectAttributes.SecurityQualityOfService = &securityService;
+    TOKEN_TYPE tokenType;
+    ULONG returnLength;
 
-    status = NtDuplicateToken(
+    status = NtQueryInformationToken(
         TokenHandle,
-        TOKEN_IMPERSONATE | TOKEN_QUERY,
-        &objectAttributes,
-        FALSE,
-        TokenImpersonation,
-        &tokenHandle
+        TokenType,
+        &tokenType,
+        sizeof(TOKEN_TYPE),
+        &returnLength
         );
 
     if (!NT_SUCCESS(status))
         return status;
 
-    status = NtSetInformationThread(
-        ThreadHandle,
-        ThreadImpersonationToken,
-        &tokenHandle,
-        sizeof(HANDLE)
-        );
+    if (tokenType == TokenPrimary)
+    {
+        SECURITY_QUALITY_OF_SERVICE securityService;
+        OBJECT_ATTRIBUTES objectAttributes;
+        HANDLE tokenHandle;
+
+        InitializeObjectAttributes(
+            &objectAttributes,
+            NULL,
+            0,
+            NULL,
+            NULL
+            );
+
+        securityService.Length = sizeof(SECURITY_QUALITY_OF_SERVICE);
+        securityService.ImpersonationLevel = SecurityImpersonation;
+        securityService.ContextTrackingMode = SECURITY_DYNAMIC_TRACKING;
+        securityService.EffectiveOnly = FALSE;
+        objectAttributes.SecurityQualityOfService = &securityService;
+
+        status = NtDuplicateToken(
+            TokenHandle,
+            TOKEN_IMPERSONATE | TOKEN_QUERY,
+            &objectAttributes,
+            FALSE,
+            TokenImpersonation,
+            &tokenHandle
+            );
 
-    NtClose(tokenHandle);
+        if (!NT_SUCCESS(status))
+            return status;
+
+        status = NtSetInformationThread(
+            ThreadHandle,
+            ThreadImpersonationToken,
+            &tokenHandle,
+            sizeof(HANDLE)
+            );
+
+        NtClose(tokenHandle);
+    }
+    else
+    {
+        status = NtSetInformationThread(
+            ThreadHandle,
+            ThreadImpersonationToken,
+            &TokenHandle,
+            sizeof(HANDLE)
+            );
+    }
 
     return status;
 }