[ros-diffs] [tkreuzer] 61109: [NTOSKRNL] - Implement SepAdjustPrivileges, which does both the counting of (changed) privileges as well as applying them, when requested. Use it in NtAdjustPrivilegesToken twice ...

Author: tkreuzer
Date: Tue Nov 26 23:28:37 2013
New Revision: 61109

URL: http://svn.reactos.org/svn/reactos?rev=61109&view=rev
Log:
[NTOSKRNL]
- Implement SepAdjustPrivileges, which does both the counting of (changed) privileges as
well as applying them, when requested. Use it in NtAdjustPrivilegesToken twice instead of
duplicating the code there.
- Fix return value of NtAdjustPrivilegesToken by making sure to properly count the found
privileges and check against the provided ones
- Lock the Token, while messing with the privileges
- Add support for SE_PRIVILEGE_REMOVED
- Proplery (re)calculate Token flags after changing privileges
- Improve failure pathes by using a common cleanup label
- Don't free the allocations atatched to the token in SepCreateToken on failure, since
ObDereferenceObject will already do that.
- Make priviliges constants instead of initializing them.

Modified:
    trunk/reactos/ntoskrnl/include/internal/se.h
    trunk/reactos/ntoskrnl/se/priv.c
    trunk/reactos/ntoskrnl/se/token.c

Modified: trunk/reactos/ntoskrnl/include/internal/se.h
URL:
http://svn.reactos.org/svn/reactos/trunk/reactos/ntoskrnl/include/internal/…
==============================================================================
--- trunk/reactos/ntoskrnl/include/internal/se.h	[iso-8859-1] (original)
+++ trunk/reactos/ntoskrnl/include/internal/se.h	[iso-8859-1] Tue Nov 26 23:28:37 2013
@@ -145,32 +145,40 @@
 extern PSID SeAnonymousLogonSid;
 
 /* Privileges */
-extern LUID SeCreateTokenPrivilege;
-extern LUID SeAssignPrimaryTokenPrivilege;
-extern LUID SeLockMemoryPrivilege;
-extern LUID SeIncreaseQuotaPrivilege;
-extern LUID SeUnsolicitedInputPrivilege;
-extern LUID SeTcbPrivilege;
-extern LUID SeSecurityPrivilege;
-extern LUID SeTakeOwnershipPrivilege;
-extern LUID SeLoadDriverPrivilege;
-extern LUID SeCreatePagefilePrivilege;
-extern LUID SeIncreaseBasePriorityPrivilege;
-extern LUID SeSystemProfilePrivilege;
-extern LUID SeSystemtimePrivilege;
-extern LUID SeProfileSingleProcessPrivilege;
-extern LUID SeCreatePermanentPrivilege;
-extern LUID SeBackupPrivilege;
-extern LUID SeRestorePrivilege;
-extern LUID SeShutdownPrivilege;
-extern LUID SeDebugPrivilege;
-extern LUID SeAuditPrivilege;
-extern LUID SeSystemEnvironmentPrivilege;
-extern LUID SeChangeNotifyPrivilege;
-extern LUID SeRemoteShutdownPrivilege;
-extern LUID SeUndockPrivilege;
-extern LUID SeSyncAgentPrivilege;
-extern LUID SeEnableDelegationPrivilege;
+extern const LUID SeCreateTokenPrivilege;
+extern const LUID SeAssignPrimaryTokenPrivilege;
+extern const LUID SeLockMemoryPrivilege;
+extern const LUID SeIncreaseQuotaPrivilege;
+extern const LUID SeUnsolicitedInputPrivilege;
+extern const LUID SeTcbPrivilege;
+extern const LUID SeSecurityPrivilege;
+extern const LUID SeTakeOwnershipPrivilege;
+extern const LUID SeLoadDriverPrivilege;
+extern const LUID SeSystemProfilePrivilege;
+extern const LUID SeSystemtimePrivilege;
+extern const LUID SeProfileSingleProcessPrivilege;
+extern const LUID SeIncreaseBasePriorityPrivilege;
+extern const LUID SeCreatePagefilePrivilege;
+extern const LUID SeCreatePermanentPrivilege;
+extern const LUID SeBackupPrivilege;
+extern const LUID SeRestorePrivilege;
+extern const LUID SeShutdownPrivilege;
+extern const LUID SeDebugPrivilege;
+extern const LUID SeAuditPrivilege;
+extern const LUID SeSystemEnvironmentPrivilege;
+extern const LUID SeChangeNotifyPrivilege;
+extern const LUID SeRemoteShutdownPrivilege;
+extern const LUID SeUndockPrivilege;
+extern const LUID SeSyncAgentPrivilege;
+extern const LUID SeEnableDelegationPrivilege;
+extern const LUID SeManageVolumePrivilege;
+extern const LUID SeImpersonatePrivilege;
+extern const LUID SeCreateGlobalPrivilege;
+extern const LUID SeTrustedCredmanPrivilege;
+extern const LUID SeRelabelPrivilege;
+extern const LUID SeIncreaseWorkingSetPrivilege;
+extern const LUID SeTimeZonePrivilege;
+extern const LUID SeCreateSymbolicLinkPrivilege;
 
 /* DACLs */
 extern PACL SePublicDefaultUnrestrictedDacl;

Modified: trunk/reactos/ntoskrnl/se/priv.c
URL:
http://svn.reactos.org/svn/reactos/trunk/reactos/ntoskrnl/se/priv.c?rev=611…
==============================================================================
--- trunk/reactos/ntoskrnl/se/priv.c	[iso-8859-1] (original)
+++ trunk/reactos/ntoskrnl/se/priv.c	[iso-8859-1] Tue Nov 26 23:28:37 2013
@@ -19,32 +19,42 @@
 
 /* GLOBALS ********************************************************************/
 
-LUID SeCreateTokenPrivilege;
-LUID SeAssignPrimaryTokenPrivilege;
-LUID SeLockMemoryPrivilege;
-LUID SeIncreaseQuotaPrivilege;
-LUID SeUnsolicitedInputPrivilege;
-LUID SeTcbPrivilege;
-LUID SeSecurityPrivilege;
-LUID SeTakeOwnershipPrivilege;
-LUID SeLoadDriverPrivilege;
-LUID SeCreatePagefilePrivilege;
-LUID SeIncreaseBasePriorityPrivilege;
-LUID SeSystemProfilePrivilege;
-LUID SeSystemtimePrivilege;
-LUID SeProfileSingleProcessPrivilege;
-LUID SeCreatePermanentPrivilege;
-LUID SeBackupPrivilege;
-LUID SeRestorePrivilege;
-LUID SeShutdownPrivilege;
-LUID SeDebugPrivilege;
-LUID SeAuditPrivilege;
-LUID SeSystemEnvironmentPrivilege;
-LUID SeChangeNotifyPrivilege;
-LUID SeRemoteShutdownPrivilege;
-LUID SeUndockPrivilege;
-LUID SeSyncAgentPrivilege;
-LUID SeEnableDelegationPrivilege;
+#define CONST_LUID(x1, x2) {x1, x2}
+const LUID SeCreateTokenPrivilege = CONST_LUID(SE_CREATE_TOKEN_PRIVILEGE, 0);
+const LUID SeAssignPrimaryTokenPrivilege = CONST_LUID(SE_ASSIGNPRIMARYTOKEN_PRIVILEGE,
0);
+const LUID SeLockMemoryPrivilege = CONST_LUID(SE_LOCK_MEMORY_PRIVILEGE, 0);
+const LUID SeIncreaseQuotaPrivilege = CONST_LUID(SE_INCREASE_QUOTA_PRIVILEGE, 0);
+const LUID SeUnsolicitedInputPrivilege = CONST_LUID(6, 0);
+const LUID SeTcbPrivilege = CONST_LUID(SE_TCB_PRIVILEGE, 0);
+const LUID SeSecurityPrivilege = CONST_LUID(SE_SECURITY_PRIVILEGE, 0);
+const LUID SeTakeOwnershipPrivilege = CONST_LUID(SE_TAKE_OWNERSHIP_PRIVILEGE, 0);
+const LUID SeLoadDriverPrivilege = CONST_LUID(SE_LOAD_DRIVER_PRIVILEGE, 0);
+const LUID SeSystemProfilePrivilege = CONST_LUID(SE_SYSTEM_PROFILE_PRIVILEGE, 0);
+const LUID SeSystemtimePrivilege = CONST_LUID(SE_SYSTEMTIME_PRIVILEGE, 0);
+const LUID SeProfileSingleProcessPrivilege = CONST_LUID(SE_PROF_SINGLE_PROCESS_PRIVILEGE,
0);
+const LUID SeIncreaseBasePriorityPrivilege = CONST_LUID(SE_INC_BASE_PRIORITY_PRIVILEGE,
0);
+const LUID SeCreatePagefilePrivilege = CONST_LUID(SE_CREATE_PAGEFILE_PRIVILEGE, 0);
+const LUID SeCreatePermanentPrivilege = CONST_LUID(SE_CREATE_PERMANENT_PRIVILEGE, 0);
+const LUID SeBackupPrivilege = CONST_LUID(SE_BACKUP_PRIVILEGE, 0);
+const LUID SeRestorePrivilege = CONST_LUID(SE_RESTORE_PRIVILEGE, 0);
+const LUID SeShutdownPrivilege = CONST_LUID(SE_SHUTDOWN_PRIVILEGE, 0);
+const LUID SeDebugPrivilege = CONST_LUID(SE_DEBUG_PRIVILEGE, 0);
+const LUID SeAuditPrivilege = CONST_LUID(SE_AUDIT_PRIVILEGE, 0);
+const LUID SeSystemEnvironmentPrivilege = CONST_LUID(SE_SYSTEM_ENVIRONMENT_PRIVILEGE,
0);
+const LUID SeChangeNotifyPrivilege = CONST_LUID(SE_CHANGE_NOTIFY_PRIVILEGE, 0);
+const LUID SeRemoteShutdownPrivilege = CONST_LUID(SE_REMOTE_SHUTDOWN_PRIVILEGE, 0);
+const LUID SeUndockPrivilege = CONST_LUID(SE_UNDOCK_PRIVILEGE, 0);
+const LUID SeSyncAgentPrivilege = CONST_LUID(SE_SYNC_AGENT_PRIVILEGE, 0);
+const LUID SeEnableDelegationPrivilege = CONST_LUID(SE_ENABLE_DELEGATION_PRIVILEGE, 0);
+const LUID SeManageVolumePrivilege = CONST_LUID(SE_MANAGE_VOLUME_PRIVILEGE, 0);
+const LUID SeImpersonatePrivilege = CONST_LUID(SE_IMPERSONATE_PRIVILEGE, 0);
+const LUID SeCreateGlobalPrivilege = CONST_LUID(SE_CREATE_GLOBAL_PRIVILEGE, 0);
+const LUID SeTrustedCredmanPrivilege = CONST_LUID(SE_TRUSTED_CREDMAN_ACCESS_PRIVILEGE,
0);
+const LUID SeRelabelPrivilege = CONST_LUID(SE_RELABEL_PRIVILEGE, 0);
+const LUID SeIncreaseWorkingSetPrivilege = CONST_LUID(SE_INC_WORKING_SET_PRIVILEGE, 0);
+const LUID SeTimeZonePrivilege = CONST_LUID(SE_TIME_ZONE_PRIVILEGE, 0);
+const LUID SeCreateSymbolicLinkPrivilege = CONST_LUID(SE_CREATE_SYMBOLIC_LINK_PRIVILEGE,
0);
+
 
 /* PRIVATE FUNCTIONS **********************************************************/
 
@@ -53,58 +63,7 @@
 NTAPI
 SepInitPrivileges(VOID)
 {
-    SeCreateTokenPrivilege.LowPart = SE_CREATE_TOKEN_PRIVILEGE;
-    SeCreateTokenPrivilege.HighPart = 0;
-    SeAssignPrimaryTokenPrivilege.LowPart = SE_ASSIGNPRIMARYTOKEN_PRIVILEGE;
-    SeAssignPrimaryTokenPrivilege.HighPart = 0;
-    SeLockMemoryPrivilege.LowPart = SE_LOCK_MEMORY_PRIVILEGE;
-    SeLockMemoryPrivilege.HighPart = 0;
-    SeIncreaseQuotaPrivilege.LowPart = SE_INCREASE_QUOTA_PRIVILEGE;
-    SeIncreaseQuotaPrivilege.HighPart = 0;
-    SeUnsolicitedInputPrivilege.LowPart = SE_UNSOLICITED_INPUT_PRIVILEGE;
-    SeUnsolicitedInputPrivilege.HighPart = 0;
-    SeTcbPrivilege.LowPart = SE_TCB_PRIVILEGE;
-    SeTcbPrivilege.HighPart = 0;
-    SeSecurityPrivilege.LowPart = SE_SECURITY_PRIVILEGE;
-    SeSecurityPrivilege.HighPart = 0;
-    SeTakeOwnershipPrivilege.LowPart = SE_TAKE_OWNERSHIP_PRIVILEGE;
-    SeTakeOwnershipPrivilege.HighPart = 0;
-    SeLoadDriverPrivilege.LowPart = SE_LOAD_DRIVER_PRIVILEGE;
-    SeLoadDriverPrivilege.HighPart = 0;
-    SeSystemProfilePrivilege.LowPart = SE_SYSTEM_PROFILE_PRIVILEGE;
-    SeSystemProfilePrivilege.HighPart = 0;
-    SeSystemtimePrivilege.LowPart = SE_SYSTEMTIME_PRIVILEGE;
-    SeSystemtimePrivilege.HighPart = 0;
-    SeProfileSingleProcessPrivilege.LowPart = SE_PROF_SINGLE_PROCESS_PRIVILEGE;
-    SeProfileSingleProcessPrivilege.HighPart = 0;
-    SeIncreaseBasePriorityPrivilege.LowPart = SE_INC_BASE_PRIORITY_PRIVILEGE;
-    SeIncreaseBasePriorityPrivilege.HighPart = 0;
-    SeCreatePagefilePrivilege.LowPart = SE_CREATE_PAGEFILE_PRIVILEGE;
-    SeCreatePagefilePrivilege.HighPart = 0;
-    SeCreatePermanentPrivilege.LowPart = SE_CREATE_PERMANENT_PRIVILEGE;
-    SeCreatePermanentPrivilege.HighPart = 0;
-    SeBackupPrivilege.LowPart = SE_BACKUP_PRIVILEGE;
-    SeBackupPrivilege.HighPart = 0;
-    SeRestorePrivilege.LowPart = SE_RESTORE_PRIVILEGE;
-    SeRestorePrivilege.HighPart = 0;
-    SeShutdownPrivilege.LowPart = SE_SHUTDOWN_PRIVILEGE;
-    SeShutdownPrivilege.HighPart = 0;
-    SeDebugPrivilege.LowPart = SE_DEBUG_PRIVILEGE;
-    SeDebugPrivilege.HighPart = 0;
-    SeAuditPrivilege.LowPart = SE_AUDIT_PRIVILEGE;
-    SeAuditPrivilege.HighPart = 0;
-    SeSystemEnvironmentPrivilege.LowPart = SE_SYSTEM_ENVIRONMENT_PRIVILEGE;
-    SeSystemEnvironmentPrivilege.HighPart = 0;
-    SeChangeNotifyPrivilege.LowPart = SE_CHANGE_NOTIFY_PRIVILEGE;
-    SeChangeNotifyPrivilege.HighPart = 0;
-    SeRemoteShutdownPrivilege.LowPart = SE_REMOTE_SHUTDOWN_PRIVILEGE;
-    SeRemoteShutdownPrivilege.HighPart = 0;
-    SeUndockPrivilege.LowPart = SE_UNDOCK_PRIVILEGE;
-    SeUndockPrivilege.HighPart = 0;
-    SeSyncAgentPrivilege.LowPart = SE_SYNC_AGENT_PRIVILEGE;
-    SeSyncAgentPrivilege.HighPart = 0;
-    SeEnableDelegationPrivilege.LowPart = SE_ENABLE_DELEGATION_PRIVILEGE;
-    SeEnableDelegationPrivilege.HighPart = 0;
+
 }
 
 
@@ -233,7 +192,7 @@
             return STATUS_INSUFFICIENT_RESOURCES;
         }
     }
-    
+
     /* copy the array to the buffer */
     _SEH2_TRY
     {

Modified: trunk/reactos/ntoskrnl/se/token.c
URL:
http://svn.reactos.org/svn/reactos/trunk/reactos/ntoskrnl/se/token.c?rev=61…
==============================================================================
--- trunk/reactos/ntoskrnl/se/token.c	[iso-8859-1] (original)
+++ trunk/reactos/ntoskrnl/se/token.c	[iso-8859-1] Tue Nov 26 23:28:37 2013
@@ -101,6 +101,95 @@
     return STATUS_SUCCESS;
 }
 
+static
+VOID
+SepUpdateSinglePrivilegeFlagToken(
+    _Inout_ PTOKEN Token,
+    _In_ ULONG Index)
+{
+    ULONG TokenFlag;
+    NT_ASSERT(Index < Token->PrivilegeCount);
+
+    /* The high part of all values we are interested in is 0 */
+    if (Token->Privileges[Index].Luid.HighPart != 0)
+    {
+        return;
+    }
+
+    /* Check for certain privileges to update flags */
+    if (Token->Privileges[Index].Luid.LowPart == SE_CHANGE_NOTIFY_PRIVILEGE)
+    {
+        TokenFlag = TOKEN_HAS_TRAVERSE_PRIVILEGE;
+    }
+    else if (Token->Privileges[Index].Luid.LowPart == SE_BACKUP_PRIVILEGE)
+    {
+        TokenFlag = TOKEN_HAS_BACKUP_PRIVILEGE;
+    }
+    else if (Token->Privileges[Index].Luid.LowPart == SE_RESTORE_PRIVILEGE)
+    {
+        TokenFlag = TOKEN_HAS_RESTORE_PRIVILEGE;
+    }
+    else if (Token->Privileges[Index].Luid.LowPart == SE_IMPERSONATE_PRIVILEGE)
+    {
+        TokenFlag = TOKEN_HAS_IMPERSONATE_PRIVILEGE;
+    }
+    else
+    {
+        /* Nothing to do */
+        return;
+    }
+
+    /* Check if the specified privilege is enabled */
+    if (Token->Privileges[Index].Attributes & SE_PRIVILEGE_ENABLED)
+    {
+        /* It is enabled, so set the flag */
+        Token->TokenFlags |= TokenFlag;
+    }
+    else
+    {
+        /* Is is disabled, so remove the flag */
+        Token->TokenFlags &= ~TokenFlag;
+    }
+}
+
+static
+VOID
+SepUpdatePrivilegeFlagsToken(
+    _Inout_ PTOKEN Token)
+{
+    ULONG i;
+
+    /* Loop all privileges */
+    for (i = 0; i < Token->PrivilegeCount; i++)
+    {
+        /* Updates the flags dor this privilege */
+        SepUpdateSinglePrivilegeFlagToken(Token, i);
+    }
+}
+
+static
+VOID
+SepRemovePrivilegeToken(
+    _Inout_ PTOKEN Token,
+    _In_ ULONG Index)
+{
+    ULONG MoveCount;
+    NT_ASSERT(Index < Token->PrivilegeCount);
+
+    /* Calculate the number of trailing privileges */
+    MoveCount = Token->PrivilegeCount - Index - 1;
+    if (MoveCount != 0)
+    {
+        /* Move them one location ahead */
+        RtlMoveMemory(&Token->Privileges[Index],
+                      &Token->Privileges[Index + 1],
+                      MoveCount * sizeof(LUID_AND_ATTRIBUTES));
+    }
+
+    /* Update privilege count */
+    Token->PrivilegeCount--;
+}
+
 VOID
 NTAPI
 SepFreeProxyData(PVOID ProxyData)
@@ -592,18 +681,6 @@
         }
     }
 
-    /* Loop all privileges */
-    for (i = 0; i < PrivilegeCount; i++)
-    {
-        /* For optimization, check for change notify and impersonate rights */
-        if (((RtlEqualLuid(&Privileges[i].Luid, &SeChangeNotifyPrivilege))
&&
-            (Privileges[i].Attributes & SE_PRIVILEGE_ENABLED)))
-        {
-            /* Remember token has traverse */
-            TokenFlags |= TOKEN_HAS_TRAVERSE_PRIVILEGE;
-        }
-    }
-
     Status = ZwAllocateLocallyUniqueId(&TokenId);
     if (!NT_SUCCESS(Status))
         return Status;
@@ -732,6 +809,9 @@
     if (!NT_SUCCESS(Status))
         goto done;
 
+    /* Update privilege flags */
+    SepUpdatePrivilegeFlagsToken(AccessToken);
+
     AccessToken->DefaultDacl = ExAllocatePoolWithTag(PagedPool,
                                                      DefaultDacl->AclSize,
                                                      TAG_TOKEN_ACL);
@@ -769,15 +849,7 @@
     {
         if (AccessToken)
         {
-            if (AccessToken->UserAndGroups)
-                ExFreePoolWithTag(AccessToken->UserAndGroups, TAG_TOKEN_USERS);
-
-            if (AccessToken->Privileges)
-                ExFreePoolWithTag(AccessToken->Privileges, TAG_TOKEN_PRIVILAGES);
-
-            if (AccessToken->DefaultDacl)
-                ExFreePoolWithTag(AccessToken->DefaultDacl, TAG_TOKEN_ACL);
-
+            /* Dereference the token, the delete procedure will clean up */
             ObDereferenceObject(AccessToken);
         }
     }
@@ -1995,16 +2067,136 @@
     return(STATUS_NOT_IMPLEMENTED);
 }
 
+
+static
+ULONG
+SepAdjustPrivileges(
+    _Inout_ PTOKEN Token,
+    _In_ BOOLEAN DisableAllPrivileges,
+    _In_opt_ PLUID_AND_ATTRIBUTES NewState,
+    _In_ ULONG NewStateCount,
+    _Out_opt_ PTOKEN_PRIVILEGES PreviousState,
+    _In_ BOOLEAN ApplyChanges,
+    _Out_ PULONG ChangedPrivileges)
+{
+    ULONG i, j, PrivilegeCount, ChangeCount, NewAttributes;
+
+    /* Count the found privileges and those that need to be changed */
+    PrivilegeCount = 0;
+    ChangeCount = 0;
+
+    /* Loop all privileges in the token */
+    for (i = 0; i < Token->PrivilegeCount; i++)
+    {
+        /* Shall all of them be disabled? */
+        if (DisableAllPrivileges)
+        {
+            /* The new attributes are the old ones, but disabled */
+            NewAttributes = Token->Privileges[i].Attributes &
~SE_PRIVILEGE_ENABLED;
+        }
+        else
+        {
+            /* Otherwise loop all provided privileges */
+            for (j = 0; j < NewStateCount; j++)
+            {
+                /* Check if this is the LUID we are looking for */
+                if (RtlEqualLuid(&Token->Privileges[i].Luid,
&NewState[j].Luid))
+                {
+                    DPRINT("Found privilege\n");
+
+                    /* Copy SE_PRIVILEGE_ENABLED | SE_PRIVILEGE_REMOVED */
+                    NewAttributes = NewState[j].Attributes;
+                    NewAttributes &= (SE_PRIVILEGE_ENABLED | SE_PRIVILEGE_REMOVED);
+                    NewAttributes |= Token->Privileges[i].Attributes &
~SE_PRIVILEGE_ENABLED;
+
+                    /* Stop looking */
+                    break;
+                }
+            }
+
+            /* Check if we didn't find the privilege */
+            if (j == NewStateCount)
+            {
+                /* Continue with the token's next privilege */
+                continue;
+            }
+        }
+
+        /* We found a privilege, count it */
+        PrivilegeCount++;
+
+        /* Does the privilege need to be changed? */
+        if (Token->Privileges[i].Attributes != NewAttributes)
+        {
+            /* Does the caller want the old privileges? */
+            if (PreviousState != NULL)
+            {
+                /* Copy the old privilege */
+                PreviousState->Privileges[ChangeCount] = Token->Privileges[i];
+            }
+
+            /* Does the caller want to apply the changes? */
+            if (ApplyChanges)
+            {
+                /* Shall we remove the privilege? */
+                if (NewAttributes & SE_PRIVILEGE_REMOVED)
+                {
+                    /* Set the token as disabled and update flags for it */
+                    Token->Privileges[i].Attributes &= ~SE_PRIVILEGE_ENABLED;
+                    SepUpdateSinglePrivilegeFlagToken(Token, i);
+
+                    /* Remove the privilege */
+                    SepRemovePrivilegeToken(Token, i);
+
+                    /* Fix the running index */
+                    i--;
+
+                    /* Continue with next */
+                    continue;
+                }
+
+                /* Set the new attributes and update flags */
+                Token->Privileges[i].Attributes = NewAttributes;
+                SepUpdateSinglePrivilegeFlagToken(Token, i);
+            }
+
+            /* Increment the change count */
+            ChangeCount++;
+        }
+    }
+
+    /* Set the number of saved privileges */
+    if (PreviousState != NULL)
+        PreviousState->PrivilegeCount = ChangeCount;
+
+    /* Return the number of changed privileges */
+    *ChangedPrivileges = ChangeCount;
+
+    /* Check if we missed some */
+    if (!DisableAllPrivileges && (PrivilegeCount < NewStateCount))
+    {
+        return STATUS_NOT_ALL_ASSIGNED;
+    }
+
+    return STATUS_SUCCESS;
+}
+
+
 /*
  * @implemented
  */
-NTSTATUS NTAPI
-NtAdjustPrivilegesToken(IN HANDLE TokenHandle,
-                        IN BOOLEAN DisableAllPrivileges,
-                        IN PTOKEN_PRIVILEGES NewState,
-                        IN ULONG BufferLength,
-                        OUT PTOKEN_PRIVILEGES PreviousState OPTIONAL,
-                        OUT PULONG ReturnLength OPTIONAL)
+_Must_inspect_result_
+__kernel_entry
+NTSTATUS
+NTAPI
+NtAdjustPrivilegesToken(
+    _In_ HANDLE TokenHandle,
+    _In_ BOOLEAN DisableAllPrivileges,
+    _In_opt_ PTOKEN_PRIVILEGES NewState,
+    _In_ ULONG BufferLength,
+    _Out_writes_bytes_to_opt_(BufferLength,*ReturnLength)
+        PTOKEN_PRIVILEGES PreviousState,
+    _When_(PreviousState!=NULL, _Out_) PULONG ReturnLength)
 {
     PLUID_AND_ATTRIBUTES CapturedPrivileges = NULL;
     KPROCESSOR_MODE PreviousMode;
@@ -2012,14 +2204,12 @@
     ULONG CapturedLength = 0;
     ULONG NewStateSize = 0;
     ULONG ChangeCount;
+    ULONG RequiredLength;
     PTOKEN Token;
-    ULONG i;
-    ULONG j;
     NTSTATUS Status;
-
     PAGED_CODE();
 
-    DPRINT ("NtAdjustPrivilegesToken() called\n");
+    DPRINT("NtAdjustPrivilegesToken() called\n");
 
     /* Fail, if we do not disable all privileges but NewState is NULL */
     if (DisableAllPrivileges == FALSE && NewState == NULL)
@@ -2033,29 +2223,20 @@
             /* Probe NewState */
             if (DisableAllPrivileges == FALSE)
             {
-                ProbeForRead(NewState,
-                             sizeof(TOKEN_PRIVILEGES),
-                             sizeof(ULONG));
+                /* First probe the header */
+                ProbeForRead(NewState, sizeof(TOKEN_PRIVILEGES), sizeof(ULONG));
 
                 CapturedCount = NewState->PrivilegeCount;
-                NewStateSize = (ULONG)sizeof(TOKEN_PRIVILEGES) +
-                               ((CapturedCount - ANYSIZE_ARRAY) *
(ULONG)sizeof(LUID_AND_ATTRIBUTES));
-
-                ProbeForRead(NewState,
-                             NewStateSize,
-                             sizeof(ULONG));
+                NewStateSize = FIELD_OFFSET(TOKEN_PRIVILEGES,
Privileges[CapturedCount]);
+
+                ProbeForRead(NewState, NewStateSize, sizeof(ULONG));
             }
 
             /* Probe PreviousState and ReturnLength */
             if (PreviousState != NULL)
             {
-                ProbeForWrite(PreviousState,
-                              BufferLength,
-                              sizeof(ULONG));
-
-                ProbeForWrite(ReturnLength,
-                              sizeof(ULONG),
-                              sizeof(ULONG));
+                ProbeForWrite(PreviousState, BufferLength, sizeof(ULONG));
+                ProbeForWrite(ReturnLength, sizeof(ULONG), sizeof(ULONG));
             }
         }
         _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
@@ -2067,10 +2248,12 @@
     }
     else
     {
+        /* This is kernel mode, we trust the caller */
         if (DisableAllPrivileges == FALSE)
             CapturedCount = NewState->PrivilegeCount;
     }
 
+    /* Do we need to capture the new state? */
     if (DisableAllPrivileges == FALSE)
     {
         _SEH2_TRY
@@ -2089,7 +2272,7 @@
         _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
         {
             /* Return the exception code */
-            _SEH2_YIELD(return _SEH2_GetExceptionCode());
+            Status = _SEH2_GetExceptionCode();
         }
         _SEH2_END;
 
@@ -2106,7 +2289,7 @@
                                        NULL);
     if (!NT_SUCCESS(Status))
     {
-        DPRINT1 ("Failed to reference token (Status %lx)\n", Status);
+        DPRINT1("Failed to reference token (Status %lx)\n", Status);
 
         /* Release the captured privileges */
         if (CapturedPrivileges != NULL)
@@ -2117,51 +2300,23 @@
         return Status;
     }
 
-    /* Count the privileges that need to be changed */
-    ChangeCount = 0;
-    for (i = 0; i < Token->PrivilegeCount; i++)
-    {
-        if (DisableAllPrivileges)
-        {
-            if (Token->Privileges[i].Attributes & SE_PRIVILEGE_ENABLED)
-            {
-                DPRINT("Privilege enabled\n");
-
-                ChangeCount++;
-            }
-        }
-        else
-        {
-            for (j = 0; j < CapturedCount; j++)
-            {
-                if (Token->Privileges[i].Luid.LowPart ==
CapturedPrivileges[j].Luid.LowPart &&
-                    Token->Privileges[i].Luid.HighPart ==
CapturedPrivileges[j].Luid.HighPart)
-                {
-                    DPRINT("Found privilege\n");
-
-                    if ((Token->Privileges[i].Attributes & SE_PRIVILEGE_ENABLED)
!=
-                        (CapturedPrivileges[j].Attributes & SE_PRIVILEGE_ENABLED))
-                    {
-                        DPRINT("Attributes differ\n");
-                        DPRINT("Current attributes %lx  New attributes %lx\n",
-                               Token->Privileges[i].Attributes,
-                               CapturedPrivileges[j].Attributes);
-
-                        ChangeCount++;
-                    }
-                }
-            }
-        }
-    }
-
-    /*
-     * Return the required buffer size and
-     * check if the available buffer is large enough
-     */
+    /* Lock the token */
+    ExEnterCriticalRegionAndAcquireResourceExclusive(Token->TokenLock);
+
+    /* Count the privileges that need to be changed, do not apply them yet */
+    Status = SepAdjustPrivileges(Token,
+                                 DisableAllPrivileges,
+                                 CapturedPrivileges,
+                                 CapturedCount,
+                                 NULL,
+                                 FALSE,
+                                 &ChangeCount);
+
+    /* Check if the caller asked for the previous state */
     if (PreviousState != NULL)
     {
-        ULONG RequiredLength = (ULONG)sizeof(TOKEN_PRIVILEGES) +
-                               ((ChangeCount - ANYSIZE_ARRAY) *
(ULONG)sizeof(LUID_AND_ATTRIBUTES));
+        /* Calculate the required length */
+        RequiredLength = FIELD_OFFSET(TOKEN_PRIVILEGES, Privileges[ChangeCount]);
 
         /* Try to return the required buffer length */
         _SEH2_TRY
@@ -2170,125 +2325,44 @@
         }
         _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
         {
-            /* Dereference the token */
-            ObDereferenceObject(Token);
-
-            /* Release the captured privileges */
-            if (CapturedPrivileges != NULL)
-                SeReleaseLuidAndAttributesArray(CapturedPrivileges,
-                                                PreviousMode,
-                                                TRUE);
-
-            /* Return the exception code */
-            _SEH2_YIELD(return _SEH2_GetExceptionCode());
+            /* Do cleanup and return the exception code */
+            Status = _SEH2_GetExceptionCode();
+            goto Cleanup;
         }
         _SEH2_END;
 
         /* Fail, if the buffer length is smaller than the required length */
         if (BufferLength < RequiredLength)
         {
-            /* Dereference the token */
-            ObDereferenceObject(Token);
-
-            /* Release the captured privileges */
-            if (CapturedPrivileges != NULL)
-                SeReleaseLuidAndAttributesArray(CapturedPrivileges,
-                                                PreviousMode,
-                                                TRUE);
-
-            return STATUS_BUFFER_TOO_SMALL;
-        }
-    }
-
-    /* Change the privilege attributes */
-    ChangeCount = 0;
+            Status = STATUS_BUFFER_TOO_SMALL;
+            goto Cleanup;
+        }
+    }
+
+    /* Now enter SEH, since we might return the old privileges */
     _SEH2_TRY
     {
-        for (i = 0; i < Token->PrivilegeCount; i++)
-        {
-            if (DisableAllPrivileges == TRUE)
-            {
-                if (Token->Privileges[i].Attributes & SE_PRIVILEGE_ENABLED)
-                {
-                    DPRINT("Privilege enabled\n");
-
-                    /* Save the current privilege */
-                    if (PreviousState != NULL)
-                    {
-                        PreviousState->Privileges[ChangeCount].Luid =
Token->Privileges[i].Luid;
-                        PreviousState->Privileges[ChangeCount].Attributes =
Token->Privileges[i].Attributes;
-                    }
-
-                    /* Disable the current privlege */
-                    Token->Privileges[i].Attributes &= ~SE_PRIVILEGE_ENABLED;
-
-                    ChangeCount++;
-                }
-            }
-            else
-            {
-                for (j = 0; j < CapturedCount; j++)
-                {
-                    if (Token->Privileges[i].Luid.LowPart ==
CapturedPrivileges[j].Luid.LowPart &&
-                        Token->Privileges[i].Luid.HighPart ==
CapturedPrivileges[j].Luid.HighPart)
-                    {
-                        DPRINT("Found privilege\n");
-
-                        /* Check whether the attributes differ */
-                        if ((Token->Privileges[i].Attributes &
SE_PRIVILEGE_ENABLED) !=
-                            (CapturedPrivileges[j].Attributes &
SE_PRIVILEGE_ENABLED))
-                        {
-                            DPRINT("Attributes differ\n");
-                            DPRINT("Current attributes %lx  New attributes
%lx\n",
-                                   Token->Privileges[i].Attributes,
-                                   CapturedPrivileges[j].Attributes);
-
-                            /* Save the current privilege */
-                            if (PreviousState != NULL)
-                            {
-                                PreviousState->Privileges[ChangeCount].Luid =
Token->Privileges[i].Luid;
-                                PreviousState->Privileges[ChangeCount].Attributes =
Token->Privileges[i].Attributes;
-                            }
-
-                            /* Update the current privlege */
-                            Token->Privileges[i].Attributes &=
~SE_PRIVILEGE_ENABLED;
-                            Token->Privileges[i].Attributes |=
-                            (CapturedPrivileges[j].Attributes &
SE_PRIVILEGE_ENABLED);
-                            DPRINT("New attributes %lx\n",
-                                   Token->Privileges[i].Attributes);
-
-                            ChangeCount++;
-                        }
-                    }
-                }
-            }
-        }
-
-        /* Set the number of saved privileges */
-        if (PreviousState != NULL)
-            PreviousState->PrivilegeCount = ChangeCount;
+        /* This time apply the changes */
+        Status = SepAdjustPrivileges(Token,
+                                     DisableAllPrivileges,
+                                     CapturedPrivileges,
+                                     CapturedCount,
+                                     PreviousState,
+                                     TRUE,
+                                     &ChangeCount);
     }
     _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
     {
-        /* Dereference the token */
-        ObDereferenceObject(Token);
-
-        /* Release the captured privileges */
-        if (CapturedPrivileges != NULL)
-            SeReleaseLuidAndAttributesArray(CapturedPrivileges,
-                                            PreviousMode,
-                                            TRUE);
-
-        /* Return the exception code */
-        _SEH2_YIELD(return _SEH2_GetExceptionCode());
+        /* Do cleanup and return the exception code */
+        Status = _SEH2_GetExceptionCode();
+        goto Cleanup;
     }
     _SEH2_END;
 
-    /* Set the status */
-    Status = (ChangeCount < CapturedCount) ? STATUS_NOT_ALL_ASSIGNED :
STATUS_SUCCESS;
-
-    /* Dereference the token */
-    ObDereferenceObject (Token);
+Cleanup:
+    /* Unlock and dereference the token */
+    ExReleaseResourceAndLeaveCriticalRegion(Token->TokenLock);
+    ObDereferenceObject(Token);
 
     /* Release the captured privileges */
     if (CapturedPrivileges != NULL)
@@ -2297,7 +2371,6 @@
                                         TRUE);
 
     DPRINT ("NtAdjustPrivilegesToken() done\n");
-
     return Status;
 }