Synchronize up to trunk's revision r57784.
[reactos.git] / ntoskrnl / se / token.c
index 703a99a..7a20cab 100644 (file)
@@ -19,7 +19,7 @@
 
 /* GLOBALS ********************************************************************/
 
-POBJECT_TYPE SepTokenObjectType = NULL;
+POBJECT_TYPE SeTokenObjectType = NULL;
 ERESOURCE SepTokenLock;
 
 TOKEN_SOURCE SeSystemTokenSource = {"*SYSTEM*", {0}};
@@ -234,13 +234,13 @@ SepDuplicateToken(PTOKEN Token,
     ULONG uLength;
     ULONG i;
     PVOID EndMem;
-    PTOKEN AccessToken;
+    PTOKEN AccessToken = NULL;
     NTSTATUS Status;
 
     PAGED_CODE();
 
     Status = ObCreateObject(PreviousMode,
-                            SepTokenObjectType,
+                            SeTokenObjectType,
                             ObjectAttributes,
                             PreviousMode,
                             NULL,
@@ -290,10 +290,14 @@ SepDuplicateToken(PTOKEN Token,
     for (i = 0; i < Token->UserAndGroupCount; i++)
         uLength += RtlLengthSid(Token->UserAndGroups[i].Sid);
 
-    AccessToken->UserAndGroups =
-    (PSID_AND_ATTRIBUTES)ExAllocatePoolWithTag(PagedPool,
-                                               uLength,
-                                               'uKOT');
+    AccessToken->UserAndGroups = ExAllocatePoolWithTag(PagedPool,
+                                                       uLength,
+                                                       TAG_TOKEN_USERS);
+    if (AccessToken->UserAndGroups == NULL)
+    {
+        Status = STATUS_INSUFFICIENT_RESOURCES;
+        goto done;
+    }
 
     EndMem = &AccessToken->UserAndGroups[AccessToken->UserAndGroupCount];
 
@@ -304,48 +308,69 @@ SepDuplicateToken(PTOKEN Token,
                                           EndMem,
                                           &EndMem,
                                           &uLength);
-    if (NT_SUCCESS(Status))
+    if (!NT_SUCCESS(Status))
+        goto done;
+
+    Status = SepFindPrimaryGroupAndDefaultOwner(AccessToken,
+                                                Token->PrimaryGroup,
+                                                0);
+    if (!NT_SUCCESS(Status))
+        goto done;
+
+    AccessToken->PrivilegeCount = Token->PrivilegeCount;
+
+    uLength = AccessToken->PrivilegeCount * sizeof(LUID_AND_ATTRIBUTES);
+    AccessToken->Privileges = ExAllocatePoolWithTag(PagedPool,
+                                                    uLength,
+                                                    TAG_TOKEN_PRIVILAGES);
+    if (AccessToken->Privileges == NULL)
     {
-        Status = SepFindPrimaryGroupAndDefaultOwner(
-                                                    AccessToken,
-                                                    Token->PrimaryGroup,
-                                                    0);
+        Status = STATUS_INSUFFICIENT_RESOURCES;
+        goto done;
     }
 
-    if (NT_SUCCESS(Status))
+    for (i = 0; i < AccessToken->PrivilegeCount; i++)
     {
-        AccessToken->PrivilegeCount = Token->PrivilegeCount;
-
-        uLength = AccessToken->PrivilegeCount * sizeof(LUID_AND_ATTRIBUTES);
-        AccessToken->Privileges =
-        (PLUID_AND_ATTRIBUTES)ExAllocatePoolWithTag(PagedPool,
-                                                    uLength,
-                                                    'pKOT');
+        RtlCopyLuid(&AccessToken->Privileges[i].Luid,
+                    &Token->Privileges[i].Luid);
+        AccessToken->Privileges[i].Attributes =
+        Token->Privileges[i].Attributes;
+    }
 
-        for (i = 0; i < AccessToken->PrivilegeCount; i++)
+    if (Token->DefaultDacl)
+    {
+        AccessToken->DefaultDacl = ExAllocatePoolWithTag(PagedPool,
+                                                         Token->DefaultDacl->AclSize,
+                                                         TAG_TOKEN_ACL);
+        if (AccessToken->DefaultDacl == NULL)
         {
-            RtlCopyLuid(&AccessToken->Privileges[i].Luid,
-                        &Token->Privileges[i].Luid);
-            AccessToken->Privileges[i].Attributes =
-            Token->Privileges[i].Attributes;
+            Status = STATUS_INSUFFICIENT_RESOURCES;
+            goto done;
         }
 
-        if (Token->DefaultDacl)
-        {
-            AccessToken->DefaultDacl =
-            (PACL) ExAllocatePoolWithTag(PagedPool,
-                                         Token->DefaultDacl->AclSize,
-                                         'kDOT');
-            memcpy(AccessToken->DefaultDacl,
-                   Token->DefaultDacl,
-                   Token->DefaultDacl->AclSize);
-        }
+        memcpy(AccessToken->DefaultDacl,
+               Token->DefaultDacl,
+               Token->DefaultDacl->AclSize);
     }
 
-    if (NT_SUCCESS(Status))
+    *NewAccessToken = AccessToken;
+
+done:
+    if (!NT_SUCCESS(Status))
     {
-        *NewAccessToken = AccessToken;
-        return(STATUS_SUCCESS);
+        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);
+
+            ObDereferenceObject(AccessToken);
+        }
     }
 
     return Status;
@@ -460,13 +485,13 @@ SepDeleteToken(PVOID ObjectBody)
     PTOKEN AccessToken = (PTOKEN)ObjectBody;
 
     if (AccessToken->UserAndGroups)
-        ExFreePool(AccessToken->UserAndGroups);
+        ExFreePoolWithTag(AccessToken->UserAndGroups, TAG_TOKEN_USERS);
 
     if (AccessToken->Privileges)
-        ExFreePool(AccessToken->Privileges);
+        ExFreePoolWithTag(AccessToken->Privileges, TAG_TOKEN_PRIVILAGES);
 
     if (AccessToken->DefaultDacl)
-        ExFreePool(AccessToken->DefaultDacl);
+        ExFreePoolWithTag(AccessToken->DefaultDacl, TAG_TOKEN_ACL);
 }
 
 
@@ -494,7 +519,7 @@ SepInitializeTokenImplementation(VOID)
     ObjectTypeInitializer.ValidAccessMask = TOKEN_ALL_ACCESS;
     ObjectTypeInitializer.UseDefaultObject = TRUE;
     ObjectTypeInitializer.DeleteProcedure = SepDeleteToken;
-    ObCreateObjectType(&Name, &ObjectTypeInitializer, NULL, &SepTokenObjectType);
+    ObCreateObjectType(&Name, &ObjectTypeInitializer, NULL, &SeTokenObjectType);
 }
 
 VOID
@@ -588,7 +613,7 @@ SepCreateToken(OUT PHANDLE TokenHandle,
         return Status;
 
     Status = ObCreateObject(PreviousMode,
-                            SepTokenObjectType,
+                            SeTokenObjectType,
                             ObjectAttributes,
                             PreviousMode,
                             NULL,
@@ -636,10 +661,14 @@ SepCreateToken(OUT PHANDLE TokenHandle,
     for (i = 0; i < GroupCount; i++)
         uLength += RtlLengthSid(Groups[i].Sid);
 
-    AccessToken->UserAndGroups =
-    (PSID_AND_ATTRIBUTES)ExAllocatePoolWithTag(PagedPool,
-                                               uLength,
-                                               'uKOT');
+    AccessToken->UserAndGroups = ExAllocatePoolWithTag(PagedPool,
+                                                       uLength,
+                                                       TAG_TOKEN_USERS);
+    if (AccessToken->UserAndGroups == NULL)
+    {
+        Status = STATUS_INSUFFICIENT_RESOURCES;
+        goto done;
+    }
 
     EndMem = &AccessToken->UserAndGroups[AccessToken->UserAndGroupCount];
 
@@ -650,66 +679,72 @@ SepCreateToken(OUT PHANDLE TokenHandle,
                                           EndMem,
                                           &EndMem,
                                           &uLength);
-    if (NT_SUCCESS(Status))
-    {
-        Status = RtlCopySidAndAttributesArray(GroupCount,
-                                              Groups,
-                                              uLength,
-                                              &AccessToken->UserAndGroups[1],
-                                              EndMem,
-                                              &EndMem,
-                                              &uLength);
-    }
+    if (!NT_SUCCESS(Status))
+        goto done;
 
-    if (NT_SUCCESS(Status))
+    Status = RtlCopySidAndAttributesArray(GroupCount,
+                                          Groups,
+                                          uLength,
+                                          &AccessToken->UserAndGroups[1],
+                                          EndMem,
+                                          &EndMem,
+                                          &uLength);
+    if (!NT_SUCCESS(Status))
+        goto done;
+
+    Status = SepFindPrimaryGroupAndDefaultOwner(AccessToken,
+                                                PrimaryGroup,
+                                                Owner);
+    if (!NT_SUCCESS(Status))
+        goto done;
+
+    uLength = PrivilegeCount * sizeof(LUID_AND_ATTRIBUTES);
+    AccessToken->Privileges = ExAllocatePoolWithTag(PagedPool,
+                                                    uLength,
+                                                    TAG_TOKEN_PRIVILAGES);
+    if (AccessToken->Privileges == NULL)
     {
-        Status = SepFindPrimaryGroupAndDefaultOwner(
-                                                    AccessToken,
-                                                    PrimaryGroup,
-                                                    Owner);
+        Status = STATUS_INSUFFICIENT_RESOURCES;
+        goto done;
     }
 
-    if (NT_SUCCESS(Status))
+    if (PreviousMode != KernelMode)
     {
-        uLength = PrivilegeCount * sizeof(LUID_AND_ATTRIBUTES);
-        AccessToken->Privileges =
-        (PLUID_AND_ATTRIBUTES)ExAllocatePoolWithTag(PagedPool,
-                                                    uLength,
-                                                    'pKOT');
-
-        if (PreviousMode != KernelMode)
-        {
-            _SEH2_TRY
-            {
-                RtlCopyMemory(AccessToken->Privileges,
-                              Privileges,
-                              PrivilegeCount * sizeof(LUID_AND_ATTRIBUTES));
-            }
-            _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
-            {
-                Status = _SEH2_GetExceptionCode();
-            }
-            _SEH2_END;
-        }
-        else
+        _SEH2_TRY
         {
             RtlCopyMemory(AccessToken->Privileges,
                           Privileges,
                           PrivilegeCount * sizeof(LUID_AND_ATTRIBUTES));
         }
+        _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
+        {
+            Status = _SEH2_GetExceptionCode();
+        }
+        _SEH2_END;
+    }
+    else
+    {
+        RtlCopyMemory(AccessToken->Privileges,
+                      Privileges,
+                      PrivilegeCount * sizeof(LUID_AND_ATTRIBUTES));
     }
 
-    if (NT_SUCCESS(Status))
+    if (!NT_SUCCESS(Status))
+        goto done;
+
+    AccessToken->DefaultDacl = ExAllocatePoolWithTag(PagedPool,
+                                                     DefaultDacl->AclSize,
+                                                     TAG_TOKEN_ACL);
+    if (AccessToken->DefaultDacl == NULL)
     {
-        AccessToken->DefaultDacl =
-        (PACL) ExAllocatePoolWithTag(PagedPool,
-                                     DefaultDacl->AclSize,
-                                     'kDOT');
-        memcpy(AccessToken->DefaultDacl,
-               DefaultDacl,
-               DefaultDacl->AclSize);
+        Status = STATUS_INSUFFICIENT_RESOURCES;
+        goto done;
     }
 
+    RtlCopyMemory(AccessToken->DefaultDacl,
+                  DefaultDacl,
+                  DefaultDacl->AclSize);
+
     if (!SystemToken)
     {
         Status = ObInsertObject((PVOID)AccessToken,
@@ -729,6 +764,24 @@ SepCreateToken(OUT PHANDLE TokenHandle,
         *TokenHandle = (HANDLE)AccessToken;
     }
 
+done:
+    if (!NT_SUCCESS(Status))
+    {
+        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);
+
+            ObDereferenceObject(AccessToken);
+        }
+    }
+
     return Status;
 }
 
@@ -1035,7 +1088,7 @@ NtQueryInformationToken(IN HANDLE TokenHandle,
 
     Status = ObReferenceObjectByHandle(TokenHandle,
                                        (TokenInformationClass == TokenSource) ? TOKEN_QUERY_SOURCE : TOKEN_QUERY,
-                                       SepTokenObjectType,
+                                       SeTokenObjectType,
                                        PreviousMode,
                                        (PVOID*)&Token,
                                        NULL);
@@ -1602,7 +1655,7 @@ NtSetInformationToken(IN HANDLE TokenHandle,
 
     Status = ObReferenceObjectByHandle(TokenHandle,
                                        NeededAccess,
-                                       SepTokenObjectType,
+                                       SeTokenObjectType,
                                        PreviousMode,
                                        (PVOID*)&Token,
                                        NULL);
@@ -1720,7 +1773,7 @@ NtSetInformationToken(IN HANDLE TokenHandle,
                             /* Free the previous dacl if present */
                             if(Token->DefaultDacl != NULL)
                             {
-                                ExFreePool(Token->DefaultDacl);
+                                ExFreePoolWithTag(Token->DefaultDacl, TAG_TOKEN_ACL);
                             }
 
                             /* Set the new dacl */
@@ -1732,7 +1785,7 @@ NtSetInformationToken(IN HANDLE TokenHandle,
                         /* Clear and free the default dacl if present */
                         if (Token->DefaultDacl != NULL)
                         {
-                            ExFreePool(Token->DefaultDacl);
+                            ExFreePoolWithTag(Token->DefaultDacl, TAG_TOKEN_ACL);
                             Token->DefaultDacl = NULL;
                         }
                     }
@@ -1845,7 +1898,7 @@ NtDuplicateToken(IN HANDLE ExistingTokenHandle,
 
     Status = ObReferenceObjectByHandle(ExistingTokenHandle,
                                        TOKEN_DUPLICATE,
-                                       SepTokenObjectType,
+                                       SeTokenObjectType,
                                        PreviousMode,
                                        (PVOID*)&Token,
                                        &HandleInformation);
@@ -1953,177 +2006,295 @@ NtAdjustPrivilegesToken(IN HANDLE TokenHandle,
                         OUT PTOKEN_PRIVILEGES PreviousState OPTIONAL,
                         OUT PULONG ReturnLength OPTIONAL)
 {
-    //  PLUID_AND_ATTRIBUTES Privileges;
+    PLUID_AND_ATTRIBUTES CapturedPrivileges = NULL;
     KPROCESSOR_MODE PreviousMode;
-    ULONG PrivilegeCount;
+    ULONG CapturedCount = 0;
+    ULONG CapturedLength = 0;
+    ULONG NewStateSize = 0;
+    ULONG ChangeCount;
     PTOKEN Token;
-    //  ULONG Length;
     ULONG i;
     ULONG j;
-    ULONG k;
-    ULONG Count;
-#if 0
-    ULONG a;
-    ULONG b;
-    ULONG c;
-#endif
     NTSTATUS Status;
 
     PAGED_CODE();
 
     DPRINT ("NtAdjustPrivilegesToken() called\n");
 
-    //  PrivilegeCount = NewState->PrivilegeCount;
+    /* Fail, if we do not disable all privileges but NewState is NULL */
+    if (DisableAllPrivileges == FALSE && NewState == NULL)
+        return STATUS_INVALID_PARAMETER;
+
     PreviousMode = KeGetPreviousMode ();
-    //  SeCaptureLuidAndAttributesArray(NewState->Privileges,
-    //                                  PrivilegeCount,
-    //                                  PreviousMode,
-    //                                  NULL,
-    //                                  0,
-    //                                  NonPagedPool,
-    //                                  1,
-    //                                  &Privileges,
-    //                                  &Length);
-
-    Status = ObReferenceObjectByHandle (TokenHandle,
-                                        TOKEN_ADJUST_PRIVILEGES | (PreviousState != NULL ? TOKEN_QUERY : 0),
-                                        SepTokenObjectType,
-                                        PreviousMode,
-                                        (PVOID*)&Token,
-                                        NULL);
-    if (!NT_SUCCESS(Status))
+    if (PreviousMode != KernelMode)
     {
-        DPRINT1 ("Failed to reference token (Status %lx)\n", Status);
-        //      SeReleaseLuidAndAttributesArray(Privileges,
-        //                                      PreviousMode,
-        //                                      0);
-        return Status;
+        _SEH2_TRY
+        {
+            /* Probe NewState */
+            if (DisableAllPrivileges == FALSE)
+            {
+                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));
+            }
+
+            /* Probe PreviousState and ReturnLength */
+            if (PreviousState != NULL)
+            {
+                ProbeForWrite(PreviousState,
+                              BufferLength,
+                              sizeof(ULONG));
+
+                ProbeForWrite(ReturnLength,
+                              sizeof(ULONG),
+                              sizeof(ULONG));
+            }
+        }
+        _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
+        {
+            /* Return the exception code */
+            _SEH2_YIELD(return _SEH2_GetExceptionCode());
+        }
+        _SEH2_END;
+    }
+    else
+    {
+        if (DisableAllPrivileges == FALSE)
+            CapturedCount = NewState->PrivilegeCount;
     }
 
+    if (DisableAllPrivileges == FALSE)
+    {
+        _SEH2_TRY
+        {
+            /* Capture the new state array of privileges */
+            Status = SeCaptureLuidAndAttributesArray(NewState->Privileges,
+                                                     CapturedCount,
+                                                     PreviousMode,
+                                                     NULL,
+                                                     0,
+                                                     PagedPool,
+                                                     TRUE,
+                                                     &CapturedPrivileges,
+                                                     &CapturedLength);
+        }
+        _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
+        {
+            /* Return the exception code */
+            _SEH2_YIELD(return _SEH2_GetExceptionCode());
+        }
+        _SEH2_END;
 
-#if 0
-    SepAdjustPrivileges(Token,
-                        0,
-                        PreviousMode,
-                        PrivilegeCount,
-                        Privileges,
-                        PreviousState,
-                        &a,
-                        &b,
-                        &c);
-#endif
+        if (!NT_SUCCESS(Status))
+            return Status;
+    }
 
-    PrivilegeCount = (BufferLength - FIELD_OFFSET(TOKEN_PRIVILEGES, Privileges)) /
-    sizeof(LUID_AND_ATTRIBUTES);
+    /* Reference the token */
+    Status = ObReferenceObjectByHandle(TokenHandle,
+                                       TOKEN_ADJUST_PRIVILEGES | (PreviousState != NULL ? TOKEN_QUERY : 0),
+                                       SeTokenObjectType,
+                                       PreviousMode,
+                                       (PVOID*)&Token,
+                                       NULL);
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1 ("Failed to reference token (Status %lx)\n", Status);
 
-    if (PreviousState != NULL)
-        PreviousState->PrivilegeCount = 0;
+        /* Release the captured privileges */
+        if (CapturedPrivileges != NULL)
+            SeReleaseLuidAndAttributesArray(CapturedPrivileges,
+                                            PreviousMode,
+                                            TRUE);
 
-    k = 0;
-    if (DisableAllPrivileges == TRUE)
+        return Status;
+    }
+
+    /* Count the privileges that need to be changed */
+    ChangeCount = 0;
+    for (i = 0; i < Token->PrivilegeCount; i++)
     {
-        for (i = 0; i < Token->PrivilegeCount; i++)
+        if (DisableAllPrivileges)
         {
-            if (Token->Privileges[i].Attributes != 0)
+            if (Token->Privileges[i].Attributes & SE_PRIVILEGE_ENABLED)
             {
-                DPRINT ("Attributes differ\n");
+                DPRINT("Privilege enabled\n");
 
-                /* Save current privilege */
-                if (PreviousState != NULL)
+                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)
                 {
-                    if (k < PrivilegeCount)
-                    {
-                        PreviousState->PrivilegeCount++;
-                        PreviousState->Privileges[k].Luid = Token->Privileges[i].Luid;
-                        PreviousState->Privileges[k].Attributes = Token->Privileges[i].Attributes;
-                    }
-                    else
+                    DPRINT("Found privilege\n");
+
+                    if ((Token->Privileges[i].Attributes & SE_PRIVILEGE_ENABLED) !=
+                        (CapturedPrivileges[j].Attributes & SE_PRIVILEGE_ENABLED))
                     {
-                        /*
-                         * FIXME: Should revert all the changes, calculate how
-                         * much space would be needed, set ResultLength
-                         * accordingly and fail.
-                         */
-                    }
+                        DPRINT("Attributes differ\n");
+                        DPRINT("Current attributes %lx  New attributes %lx\n",
+                               Token->Privileges[i].Attributes,
+                               CapturedPrivileges[j].Attributes);
 
-                    k++;
+                        ChangeCount++;
+                    }
                 }
-
-                /* Update current privlege */
-                Token->Privileges[i].Attributes &= ~SE_PRIVILEGE_ENABLED;
             }
         }
+    }
+
+    /*
+     * Return the required buffer size and
+     * check if the available buffer is large enough
+     */
+    if (PreviousState != NULL)
+    {
+        ULONG RequiredLength = (ULONG)sizeof(TOKEN_PRIVILEGES) +
+                               ((ChangeCount - ANYSIZE_ARRAY) * (ULONG)sizeof(LUID_AND_ATTRIBUTES));
+
+        /* Try to return the required buffer length */
+        _SEH2_TRY
+        {
+            *ReturnLength = RequiredLength;
+        }
+        _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
+        {
+            /* Dereference the token */
+            ObDereferenceObject(Token);
 
-        Status = STATUS_SUCCESS;
+            /* Release the captured privileges */
+            if (CapturedPrivileges != NULL)
+                SeReleaseLuidAndAttributesArray(CapturedPrivileges,
+                                                PreviousMode,
+                                                TRUE);
+
+            /* Return the exception code */
+            _SEH2_YIELD(return _SEH2_GetExceptionCode());
+        }
+        _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;
+        }
     }
-    else
+
+    /* Change the privilege attributes */
+    ChangeCount = 0;
+    _SEH2_TRY
     {
-        Count = 0;
         for (i = 0; i < Token->PrivilegeCount; i++)
         {
-            for (j = 0; j < NewState->PrivilegeCount; j++)
+            if (DisableAllPrivileges == TRUE)
             {
-                if (Token->Privileges[i].Luid.LowPart == NewState->Privileges[j].Luid.LowPart &&
-                    Token->Privileges[i].Luid.HighPart == NewState->Privileges[j].Luid.HighPart)
+                if (Token->Privileges[i].Attributes & SE_PRIVILEGE_ENABLED)
                 {
-                    DPRINT ("Found privilege\n");
+                    DPRINT("Privilege enabled\n");
 
-                    if ((Token->Privileges[i].Attributes & SE_PRIVILEGE_ENABLED) !=
-                        (NewState->Privileges[j].Attributes & SE_PRIVILEGE_ENABLED))
+                    /* 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 ("Attributes differ\n");
-                        DPRINT ("Current attributes %lx  desired attributes %lx\n",
-                                Token->Privileges[i].Attributes,
-                                NewState->Privileges[j].Attributes);
+                        DPRINT("Found privilege\n");
 
-                        /* Save current privilege */
-                        if (PreviousState != NULL)
+                        /* Check whether the attributes differ */
+                        if ((Token->Privileges[i].Attributes & SE_PRIVILEGE_ENABLED) !=
+                            (CapturedPrivileges[j].Attributes & SE_PRIVILEGE_ENABLED))
                         {
-                            if (k < PrivilegeCount)
-                            {
-                                PreviousState->PrivilegeCount++;
-                                PreviousState->Privileges[k].Luid = Token->Privileges[i].Luid;
-                                PreviousState->Privileges[k].Attributes = Token->Privileges[i].Attributes;
-                            }
-                            else
+                            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)
                             {
-                                /*
-                                 * FIXME: Should revert all the changes, calculate how
-                                 * much space would be needed, set ResultLength
-                                 * accordingly and fail.
-                                 */
+                                PreviousState->Privileges[ChangeCount].Luid = Token->Privileges[i].Luid;
+                                PreviousState->Privileges[ChangeCount].Attributes = Token->Privileges[i].Attributes;
                             }
 
-                            k++;
-                        }
+                            /* 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);
 
-                        /* Update current privlege */
-                        Token->Privileges[i].Attributes &= ~SE_PRIVILEGE_ENABLED;
-                        Token->Privileges[i].Attributes |=
-                        (NewState->Privileges[j].Attributes & SE_PRIVILEGE_ENABLED);
-                        DPRINT ("New attributes %lx\n",
-                                Token->Privileges[i].Attributes);
+                            ChangeCount++;
+                        }
                     }
-
-                    Count++;
                 }
             }
         }
 
-        Status = Count < NewState->PrivilegeCount ? STATUS_NOT_ALL_ASSIGNED : STATUS_SUCCESS;
+        /* Set the number of saved privileges */
+        if (PreviousState != NULL)
+            PreviousState->PrivilegeCount = ChangeCount;
     }
-
-    if (ReturnLength != NULL)
+    _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
     {
-        *ReturnLength = sizeof(TOKEN_PRIVILEGES) +
-        (sizeof(LUID_AND_ATTRIBUTES) * (k - 1));
+        /* 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());
     }
+    _SEH2_END;
 
+    /* Set the status */
+    Status = (ChangeCount < CapturedCount) ? STATUS_NOT_ALL_ASSIGNED : STATUS_SUCCESS;
+
+    /* Dereference the token */
     ObDereferenceObject (Token);
 
-    //  SeReleaseLuidAndAttributesArray(Privileges,
-    //                                  PreviousMode,
-    //                                  0);
+    /* Release the captured privileges */
+    if (CapturedPrivileges != NULL)
+        SeReleaseLuidAndAttributesArray(CapturedPrivileges,
+                                        PreviousMode,
+                                        TRUE);
 
     DPRINT ("NtAdjustPrivilegesToken() done\n");
 
@@ -2356,11 +2527,11 @@ NtOpenThreadTokenEx(IN HANDLE ThreadHandle,
     else
     {
         Status = ObOpenObjectByPointer(Token, HandleAttributes,
-                                       NULL, DesiredAccess, SepTokenObjectType,
+                                       NULL, DesiredAccess, SeTokenObjectType,
                                        PreviousMode, &hToken);
     }
 
-    if (Dacl) ExFreePool(Dacl);
+    if (Dacl) ExFreePoolWithTag(Dacl, TAG_TOKEN_ACL);
 
     if (OpenAsSelf)
     {
@@ -2443,7 +2614,7 @@ NtCompareTokens(IN HANDLE FirstTokenHandle,
 
     Status = ObReferenceObjectByHandle(FirstTokenHandle,
                                        TOKEN_QUERY,
-                                       SepTokenObjectType,
+                                       SeTokenObjectType,
                                        PreviousMode,
                                        (PVOID*)&FirstToken,
                                        NULL);
@@ -2452,7 +2623,7 @@ NtCompareTokens(IN HANDLE FirstTokenHandle,
 
     Status = ObReferenceObjectByHandle(SecondTokenHandle,
                                        TOKEN_QUERY,
-                                       SepTokenObjectType,
+                                       SeTokenObjectType,
                                        PreviousMode,
                                        (PVOID*)&SecondToken,
                                        NULL);