convert DefaultSetInfoBufferCheck and DefaultQueryInfoBufferCheck to inlined functions
[reactos.git] / reactos / ntoskrnl / se / token.c
index b8ddffb..009fa98 100644 (file)
@@ -1,12 +1,11 @@
 /* $Id$
  *
- * COPYRIGHT:         See COPYING in the top level directory
- * PROJECT:           ReactOS kernel
- * PURPOSE:           Security manager
- * FILE:              kernel/se/token.c
- * PROGRAMER:         David Welch <welch@cwcom.net>
- * REVISION HISTORY:
- *                 26/07/98: Added stubs for security functions
+ * COPYRIGHT:       See COPYING in the top level directory
+ * PROJECT:         ReactOS kernel
+ * FILE:            ntoskrnl/se/token.c
+ * PURPOSE:         Security manager
+ *
+ * PROGRAMMERS:     David Welch <welch@cwcom.net>
  */
 
 /* INCLUDES *****************************************************************/
 /* GLOBALS *******************************************************************/
 
 POBJECT_TYPE SepTokenObjectType = NULL;
+ERESOURCE SepTokenLock;
 
 static GENERIC_MAPPING SepTokenMapping = {TOKEN_READ,
                                          TOKEN_WRITE,
                                          TOKEN_EXECUTE,
                                          TOKEN_ALL_ACCESS};
 
+static const INFORMATION_CLASS_INFO SeTokenInformationClass[] = {
+
+    /* Class 0 not used, blame M$! */
+    ICI_SQ_SAME( 0, 0, 0),
+
+    /* TokenUser */
+    ICI_SQ_SAME( sizeof(TOKEN_USER),                   sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE | ICIF_SET | ICIF_SET_SIZE_VARIABLE ),
+    /* TokenGroups */
+    ICI_SQ_SAME( sizeof(TOKEN_GROUPS),                 sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE | ICIF_SET | ICIF_SET_SIZE_VARIABLE ),
+    /* TokenPrivileges */
+    ICI_SQ_SAME( sizeof(TOKEN_PRIVILEGES),             sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE | ICIF_SET | ICIF_SET_SIZE_VARIABLE ),
+    /* TokenOwner */
+    ICI_SQ_SAME( sizeof(TOKEN_OWNER),                  sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE | ICIF_SET | ICIF_SET_SIZE_VARIABLE ),
+    /* TokenPrimaryGroup */
+    ICI_SQ_SAME( sizeof(TOKEN_PRIMARY_GROUP),          sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE | ICIF_SET | ICIF_SET_SIZE_VARIABLE ),
+    /* TokenDefaultDacl */
+    ICI_SQ_SAME( sizeof(TOKEN_DEFAULT_DACL),           sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE | ICIF_SET | ICIF_SET_SIZE_VARIABLE ),
+    /* TokenSource */
+    ICI_SQ_SAME( sizeof(TOKEN_SOURCE),                 sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE | ICIF_SET | ICIF_SET_SIZE_VARIABLE ),
+    /* TokenType */
+    ICI_SQ_SAME( sizeof(TOKEN_TYPE),                   sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE ),
+    /* TokenImpersonationLevel */
+    ICI_SQ_SAME( sizeof(SECURITY_IMPERSONATION_LEVEL), sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE ),
+    /* TokenStatistics */
+    ICI_SQ_SAME( sizeof(TOKEN_STATISTICS),             sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE | ICIF_SET | ICIF_SET_SIZE_VARIABLE ),
+    /* TokenRestrictedSids */
+    ICI_SQ_SAME( sizeof(TOKEN_GROUPS),                 sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE ),
+    /* TokenSessionId */
+    ICI_SQ_SAME( sizeof(ULONG),                        sizeof(ULONG), ICIF_QUERY | ICIF_SET ),
+    /* TokenGroupsAndPrivileges */
+    ICI_SQ_SAME( sizeof(TOKEN_GROUPS_AND_PRIVILEGES),  sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE ),
+    /* TokenSessionReference */
+    ICI_SQ_SAME( /* FIXME */0,                         sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE ),
+    /* TokenSandBoxInert */
+    ICI_SQ_SAME( sizeof(ULONG),                        sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE ),
+    /* TokenAuditPolicy */
+    ICI_SQ_SAME( /* FIXME */0,                         sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE ),
+    /* TokenOrigin */
+    ICI_SQ_SAME( sizeof(TOKEN_ORIGIN),                 sizeof(ULONG), ICIF_QUERY | ICIF_QUERY_SIZE_VARIABLE ),
+};
+
 /* FUNCTIONS *****************************************************************/
 
-VOID SepFreeProxyData(PVOID ProxyData)
+VOID
+NTAPI
+SepFreeProxyData(PVOID ProxyData)
 {
    UNIMPLEMENTED;
 }
 
-NTSTATUS SepCopyProxyData(PVOID* Dest, PVOID Src)
+NTSTATUS
+NTAPI
+SepCopyProxyData(PVOID* Dest, PVOID Src)
 {
    UNIMPLEMENTED;
    return(STATUS_NOT_IMPLEMENTED);
 }
 
-NTSTATUS SeExchangePrimaryToken(PEPROCESS Process,
-                               PACCESS_TOKEN NewTokenP,
-                               PACCESS_TOKEN* OldTokenP)
+NTSTATUS
+NTAPI
+SeExchangePrimaryToken(PEPROCESS Process,
+                       PACCESS_TOKEN NewTokenP,
+                       PACCESS_TOKEN* OldTokenP)
 {
-   PTOKEN OldToken;
-   PTOKEN NewToken = (PTOKEN)NewTokenP;
-   
-   if (NewToken->TokenType != TokenPrimary)
-     {
-       return(STATUS_UNSUCCESSFUL);
-     }
-   if (NewToken->TokenInUse != 0)
-     {
-       return(STATUS_UNSUCCESSFUL);
-     }
-   OldToken = Process->Token;
-   Process->Token = NewToken;
-   NewToken->TokenInUse = 1;
-   ObReferenceObjectByPointer(NewToken,
-                             TOKEN_ALL_ACCESS,
-                             SepTokenObjectType,
-                             KernelMode);
-   OldToken->TokenInUse = 0;
-   *OldTokenP = (PACCESS_TOKEN)OldToken;
-   return(STATUS_SUCCESS);
+    PTOKEN OldToken;
+    PTOKEN NewToken = (PTOKEN)NewTokenP;
+
+    PAGED_CODE();
+
+    if (NewToken->TokenType != TokenPrimary) return(STATUS_BAD_TOKEN_TYPE);
+    if (NewToken->TokenInUse) return(STATUS_TOKEN_ALREADY_IN_USE);
+
+    /* Mark new token in use */
+    NewToken->TokenInUse = 1;
+
+    /* Reference the New Token */
+    ObReferenceObject(NewToken);
+
+    /* Replace the old with the new */
+    OldToken = ObFastReplaceObject(&Process->Token, NewToken);
+
+    /* Mark the Old Token as free */
+    OldToken->TokenInUse = 0;
+
+    *OldTokenP = (PACCESS_TOKEN)OldToken;
+    return STATUS_SUCCESS;
+}
+
+VOID
+NTAPI
+SeDeassignPrimaryToken(PEPROCESS Process)
+{
+    PTOKEN OldToken;
+
+    /* Remove the Token */
+    OldToken = ObFastReplaceObject(&Process->Token, NULL);
+
+    /* Mark the Old Token as free */
+    OldToken->TokenInUse = 0;
 }
 
 static ULONG
@@ -72,6 +133,8 @@ RtlLengthSidAndAttributes(ULONG Count,
   ULONG i;
   ULONG uLength;
 
+  PAGED_CODE();
+
   uLength = Count * sizeof(SID_AND_ATTRIBUTES);
   for (i = 0; i < Count; i++)
     uLength += RtlLengthSid(Src[i].Sid);
@@ -81,6 +144,7 @@ RtlLengthSidAndAttributes(ULONG Count,
 
 
 NTSTATUS
+NTAPI
 SepFindPrimaryGroupAndDefaultOwner(PTOKEN Token,
                                   PSID PrimaryGroup,
                                   PSID DefaultOwner)
@@ -124,6 +188,7 @@ SepFindPrimaryGroupAndDefaultOwner(PTOKEN Token,
 
 
 NTSTATUS
+STDCALL
 SepDuplicateToken(PTOKEN Token,
                  POBJECT_ATTRIBUTES ObjectAttributes,
                  BOOLEAN EffectiveOnly,
@@ -132,13 +197,13 @@ SepDuplicateToken(PTOKEN Token,
                  KPROCESSOR_MODE PreviousMode,
                  PTOKEN* NewAccessToken)
 {
-  NTSTATUS Status;
   ULONG uLength;
   ULONG i;
-  
   PVOID EndMem;
-
   PTOKEN AccessToken;
+  NTSTATUS Status;
+
+  PAGED_CODE();
 
   Status = ObCreateObject(PreviousMode,
                          SepTokenObjectType,
@@ -169,6 +234,8 @@ SepDuplicateToken(PTOKEN Token,
       return(Status);
     }
 
+  AccessToken->TokenLock = &SepTokenLock;
+
   AccessToken->TokenInUse = 0;
   AccessToken->TokenType  = TokenType;
   AccessToken->ImpersonationLevel = Level;
@@ -187,8 +254,8 @@ SepDuplicateToken(PTOKEN Token,
   for (i = 0; i < Token->UserAndGroupCount; i++)
     uLength += RtlLengthSid(Token->UserAndGroups[i].Sid);
 
-  AccessToken->UserAndGroups = 
-    (PSID_AND_ATTRIBUTES)ExAllocatePoolWithTag(NonPagedPool,
+  AccessToken->UserAndGroups =
+    (PSID_AND_ATTRIBUTES)ExAllocatePoolWithTag(PagedPool,
                                               uLength,
                                               TAG('T', 'O', 'K', 'u'));
 
@@ -215,7 +282,7 @@ SepDuplicateToken(PTOKEN Token,
 
       uLength = AccessToken->PrivilegeCount * sizeof(LUID_AND_ATTRIBUTES);
       AccessToken->Privileges =
-       (PLUID_AND_ATTRIBUTES)ExAllocatePoolWithTag(NonPagedPool,
+       (PLUID_AND_ATTRIBUTES)ExAllocatePoolWithTag(PagedPool,
                                                    uLength,
                                                    TAG('T', 'O', 'K', 'p'));
 
@@ -223,14 +290,14 @@ SepDuplicateToken(PTOKEN Token,
        {
          RtlCopyLuid(&AccessToken->Privileges[i].Luid,
            &Token->Privileges[i].Luid);
-         AccessToken->Privileges[i].Attributes = 
+         AccessToken->Privileges[i].Attributes =
            Token->Privileges[i].Attributes;
        }
 
       if ( Token->DefaultDacl )
        {
          AccessToken->DefaultDacl =
-           (PACL) ExAllocatePoolWithTag(NonPagedPool,
+           (PACL) ExAllocatePoolWithTag(PagedPool,
                                         Token->DefaultDacl->AclSize,
                                         TAG('T', 'O', 'K', 'd'));
          memcpy(AccessToken->DefaultDacl,
@@ -253,39 +320,6 @@ SepDuplicateToken(PTOKEN Token,
   return(Status);
 }
 
-
-NTSTATUS
-SepInitializeNewProcess(struct _EPROCESS* NewProcess,
-                       struct _EPROCESS* ParentProcess)
-{
-  NTSTATUS Status;
-  PTOKEN pNewToken;
-  PTOKEN pParentToken;
-  
-  OBJECT_ATTRIBUTES ObjectAttributes;
-
-  pParentToken = (PACCESS_TOKEN) ParentProcess->Token;
-
-  InitializeObjectAttributes(&ObjectAttributes,
-                           NULL,
-                           0,
-                           NULL,
-                           NULL);
-
-  Status = SepDuplicateToken(pParentToken,
-                            &ObjectAttributes,
-                            FALSE,
-                            TokenPrimary,
-                            pParentToken->ImpersonationLevel,
-                            KernelMode,
-                            &pNewToken);
-  if ( ! NT_SUCCESS(Status) )
-    return Status;
-
-  NewProcess->Token = pNewToken;
-  return(STATUS_SUCCESS);
-}
-
 /*
  * @unimplemented
  */
@@ -309,7 +343,9 @@ SeCopyClientToken(PACCESS_TOKEN Token,
 {
    NTSTATUS Status;
    OBJECT_ATTRIBUTES ObjectAttributes;
-     
+
+   PAGED_CODE();
+
    InitializeObjectAttributes(&ObjectAttributes,
                              NULL,
                              0,
@@ -321,8 +357,8 @@ SeCopyClientToken(PACCESS_TOKEN Token,
                                TokenImpersonation,
                                Level,
                                PreviousMode,
-                           (PTOKEN*)&NewToken);
-   
+                               (PTOKEN*)NewToken);
+
    return(Status);
 }
 
@@ -342,7 +378,9 @@ SeCreateClientSecurity(IN struct _ETHREAD *Thread,
    PACCESS_TOKEN Token;
    ULONG g;
    PACCESS_TOKEN NewToken;
-   
+
+   PAGED_CODE();
+
    Token = PsReferenceEffectiveToken(Thread,
                                     &TokenType,
                                     &b,
@@ -381,7 +419,7 @@ SeCreateClientSecurity(IN struct _ETHREAD *Thread,
             ClientContext->DirectAccessEffectiveOnly = FALSE;
          }
      }
-   
+
    if (Qos->ContextTrackingMode == 0)
      {
        ClientContext->DirectlyAccessClientToken = FALSE;
@@ -487,7 +525,9 @@ SeImpersonateClient(IN PSECURITY_CLIENT_CONTEXT ClientContext,
                    IN PETHREAD ServerThread OPTIONAL)
 {
   UCHAR b;
-  
+
+  PAGED_CODE();
+
   if (ClientContext->DirectlyAccessClientToken == FALSE)
     {
       b = ClientContext->SecurityQos.EffectiveOnly;
@@ -524,33 +564,29 @@ SepDeleteToken(PVOID ObjectBody)
 }
 
 
-VOID INIT_FUNCTION
+VOID
+INIT_FUNCTION
+NTAPI
 SepInitializeTokenImplementation(VOID)
 {
-  SepTokenObjectType = ExAllocatePool(NonPagedPool, sizeof(OBJECT_TYPE));
-
-  SepTokenObjectType->Tag = TAG('T', 'O', 'K', 'T');
-  SepTokenObjectType->PeakObjects = 0;
-  SepTokenObjectType->PeakHandles = 0;
-  SepTokenObjectType->TotalObjects = 0;
-  SepTokenObjectType->TotalHandles = 0;
-  SepTokenObjectType->PagedPoolCharge = 0;
-  SepTokenObjectType->NonpagedPoolCharge = sizeof(TOKEN);
-  SepTokenObjectType->Mapping = &SepTokenMapping;
-  SepTokenObjectType->Dump = NULL;
-  SepTokenObjectType->Open = NULL;
-  SepTokenObjectType->Close = NULL;
-  SepTokenObjectType->Delete = SepDeleteToken;
-  SepTokenObjectType->Parse = NULL;
-  SepTokenObjectType->Security = NULL;
-  SepTokenObjectType->QueryName = NULL;
-  SepTokenObjectType->OkayToClose = NULL;
-  SepTokenObjectType->Create = NULL;
-  SepTokenObjectType->DuplicationNotify = NULL;
-
-  RtlCreateUnicodeString(&SepTokenObjectType->TypeName,
-                        L"Token");
-  ObpCreateTypeObject (SepTokenObjectType);
+    UNICODE_STRING Name;
+    OBJECT_TYPE_INITIALIZER ObjectTypeInitializer;
+    
+    ExInitializeResource(&SepTokenLock);
+    
+    DPRINT("Creating Token Object Type\n");
+  
+    /*  Initialize the Token type  */
+    RtlZeroMemory(&ObjectTypeInitializer, sizeof(ObjectTypeInitializer));
+    RtlInitUnicodeString(&Name, L"Token");
+    ObjectTypeInitializer.Length = sizeof(ObjectTypeInitializer);
+    ObjectTypeInitializer.DefaultPagedPoolCharge = sizeof(TOKEN);
+    ObjectTypeInitializer.GenericMapping = SepTokenMapping;
+    ObjectTypeInitializer.PoolType = PagedPool;
+    ObjectTypeInitializer.ValidAccessMask = TOKEN_ALL_ACCESS;
+    ObjectTypeInitializer.UseDefaultObject = TRUE;
+    ObjectTypeInitializer.DeleteProcedure = SepDeleteToken;
+    ObpCreateTypeObject(&ObjectTypeInitializer, &Name, &SepTokenObjectType);
 }
 
 
@@ -564,263 +600,455 @@ NtQueryInformationToken(IN HANDLE TokenHandle,
                        IN ULONG TokenInformationLength,
                        OUT PULONG ReturnLength)
 {
-  NTSTATUS Status, LengthStatus;
-  PVOID UnusedInfo;
-  PVOID EndMem;
+  union
+  {
+    PVOID Ptr;
+    ULONG Ulong;
+  } Unused;
   PTOKEN Token;
-  ULONG Length;
-  PTOKEN_GROUPS PtrTokenGroups;
-  PTOKEN_DEFAULT_DACL PtrDefaultDacl;
-  PTOKEN_STATISTICS PtrTokenStatistics;
+  ULONG RequiredLength;
+  KPROCESSOR_MODE PreviousMode;
+  NTSTATUS Status = STATUS_SUCCESS;
+
+  PAGED_CODE();
+
+  PreviousMode = ExGetPreviousMode();
+
+  /* Check buffers and class validity */
+  Status = DefaultQueryInfoBufferCheck(TokenInformationClass,
+                                       SeTokenInformationClass,
+                                       sizeof(SeTokenInformationClass) / sizeof(SeTokenInformationClass[0]),
+                                       TokenInformation,
+                                       TokenInformationLength,
+                                       ReturnLength,
+                                       PreviousMode);
+
+  if(!NT_SUCCESS(Status))
+  {
+    DPRINT("NtQueryInformationToken() failed, Status: 0x%x\n", Status);
+    return Status;
+  }
 
   Status = ObReferenceObjectByHandle(TokenHandle,
                                     (TokenInformationClass == TokenSource) ? TOKEN_QUERY_SOURCE : TOKEN_QUERY,
                                     SepTokenObjectType,
-                                    UserMode,
+                                    PreviousMode,
                                     (PVOID*)&Token,
                                     NULL);
-  if (!NT_SUCCESS(Status))
-    {
-      return(Status);
-    }
-
-  switch (TokenInformationClass)
+  if (NT_SUCCESS(Status))
+  {
+    switch (TokenInformationClass)
     {
       case TokenUser:
-       DPRINT("NtQueryInformationToken(TokenUser)\n");
-       Length = RtlLengthSidAndAttributes(1, Token->UserAndGroups);
-       if (TokenInformationLength < Length)
-         {
-           Status = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-           if (NT_SUCCESS(Status))
-             Status = STATUS_BUFFER_TOO_SMALL;
-         }
-       else
-         {
-           Status = RtlCopySidAndAttributesArray(1,
-                                                 Token->UserAndGroups,
-                                                 TokenInformationLength,
-                                                 TokenInformation,
-                                                 (char*)TokenInformation + 8,
-                                                 &UnusedInfo,
-                                                 &Length);
-           if (NT_SUCCESS(Status))
-             {
-               Length = TokenInformationLength - Length;
-               Status = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-             }
-         }
-       break;
-       
+      {
+        PTOKEN_USER tu = (PTOKEN_USER)TokenInformation;
+
+        DPRINT("NtQueryInformationToken(TokenUser)\n");
+        RequiredLength = sizeof(TOKEN_USER) +
+                         RtlLengthSid(Token->UserAndGroups[0].Sid);
+
+        _SEH_TRY
+        {
+          if(TokenInformationLength >= RequiredLength)
+          {
+            Status = RtlCopySidAndAttributesArray(1,
+                                                  &Token->UserAndGroups[0],
+                                                  RequiredLength - sizeof(TOKEN_USER),
+                                                  &tu->User,
+                                                  (PSID)(tu + 1),
+                                                  &Unused.Ptr,
+                                                  &Unused.Ulong);
+          }
+          else
+          {
+            Status = STATUS_BUFFER_TOO_SMALL;
+          }
+
+          if(ReturnLength != NULL)
+          {
+            *ReturnLength = RequiredLength;
+          }
+        }
+        _SEH_HANDLE
+        {
+          Status = _SEH_GetExceptionCode();
+        }
+        _SEH_END;
+
+        break;
+      }
+
       case TokenGroups:
-       DPRINT("NtQueryInformationToken(TokenGroups)\n");
-       Length = RtlLengthSidAndAttributes(Token->UserAndGroupCount - 1, &Token->UserAndGroups[1]) + sizeof(ULONG);
-       if (TokenInformationLength < Length)
-         {
-           Status = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-           if (NT_SUCCESS(Status))
-             Status = STATUS_BUFFER_TOO_SMALL;
-         }
-       else
-         {
-           EndMem = (char*)TokenInformation + Token->UserAndGroupCount * sizeof(SID_AND_ATTRIBUTES);
-           PtrTokenGroups = (PTOKEN_GROUPS)TokenInformation;
-           PtrTokenGroups->GroupCount = Token->UserAndGroupCount - 1;
-           Status = RtlCopySidAndAttributesArray(Token->UserAndGroupCount - 1,
-                                                 &Token->UserAndGroups[1],
-                                                 TokenInformationLength,
-                                                 PtrTokenGroups->Groups,
-                                                 EndMem,
-                                                 &UnusedInfo,
-                                                 &Length);
-           if (NT_SUCCESS(Status))
-             {
-               Length = TokenInformationLength - Length;
-               Status = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-             }
-         }
+      {
+        PTOKEN_GROUPS tg = (PTOKEN_GROUPS)TokenInformation;
+
+        DPRINT("NtQueryInformationToken(TokenGroups)\n");
+        RequiredLength = sizeof(tg->GroupCount) +
+                         RtlLengthSidAndAttributes(Token->UserAndGroupCount - 1, &Token->UserAndGroups[1]);
+
+        _SEH_TRY
+        {
+          if(TokenInformationLength >= RequiredLength)
+          {
+            ULONG SidLen = RequiredLength - sizeof(tg->GroupCount) -
+                           ((Token->UserAndGroupCount - 1) * sizeof(SID_AND_ATTRIBUTES));
+            PSID_AND_ATTRIBUTES Sid = (PSID_AND_ATTRIBUTES)((ULONG_PTR)TokenInformation + sizeof(tg->GroupCount) +
+                                                            ((Token->UserAndGroupCount - 1) * sizeof(SID_AND_ATTRIBUTES)));
+
+            tg->GroupCount = Token->UserAndGroupCount - 1;
+            Status = RtlCopySidAndAttributesArray(Token->UserAndGroupCount - 1,
+                                                  &Token->UserAndGroups[1],
+                                                  SidLen,
+                                                  &tg->Groups[0],
+                                                  (PSID)Sid,
+                                                  &Unused.Ptr,
+                                                  &Unused.Ulong);
+          }
+          else
+          {
+            Status = STATUS_BUFFER_TOO_SMALL;
+          }
+
+          if(ReturnLength != NULL)
+          {
+            *ReturnLength = RequiredLength;
+          }
+        }
+        _SEH_HANDLE
+        {
+          Status = _SEH_GetExceptionCode();
+        }
+        _SEH_END;
+
        break;
+      }
 
       case TokenPrivileges:
-       DPRINT("NtQueryInformationToken(TokenPrivileges)\n");
-       Length = sizeof(ULONG) + Token->PrivilegeCount * sizeof(LUID_AND_ATTRIBUTES);
-       if (TokenInformationLength < Length)
-         {
-           Status = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-           if (NT_SUCCESS(Status))
-             Status = STATUS_BUFFER_TOO_SMALL;
-         }
-       else
-         {
-           ULONG i;
-           TOKEN_PRIVILEGES* pPriv = (TOKEN_PRIVILEGES*)TokenInformation;
-
-           pPriv->PrivilegeCount = Token->PrivilegeCount;
-           for (i = 0; i < Token->PrivilegeCount; i++)
-             {
-               RtlCopyLuid(&pPriv->Privileges[i].Luid, &Token->Privileges[i].Luid);
-               pPriv->Privileges[i].Attributes = Token->Privileges[i].Attributes;
-             }
-           Status = STATUS_SUCCESS;
-         }
-       break;
+      {
+        PTOKEN_PRIVILEGES tp = (PTOKEN_PRIVILEGES)TokenInformation;
+
+        DPRINT("NtQueryInformationToken(TokenPrivileges)\n");
+        RequiredLength = sizeof(tp->PrivilegeCount) +
+                         (Token->PrivilegeCount * sizeof(LUID_AND_ATTRIBUTES));
+
+        _SEH_TRY
+        {
+          if(TokenInformationLength >= RequiredLength)
+          {
+            tp->PrivilegeCount = Token->PrivilegeCount;
+            RtlCopyLuidAndAttributesArray(Token->PrivilegeCount,
+                                          Token->Privileges,
+                                          &tp->Privileges[0]);
+          }
+          else
+          {
+            Status = STATUS_BUFFER_TOO_SMALL;
+          }
+
+          if(ReturnLength != NULL)
+          {
+            *ReturnLength = RequiredLength;
+          }
+        }
+        _SEH_HANDLE
+        {
+          Status = _SEH_GetExceptionCode();
+        }
+        _SEH_END;
+
+        break;
+      }
 
       case TokenOwner:
-       DPRINT("NtQueryInformationToken(TokenOwner)\n");
-       Length = RtlLengthSid(Token->UserAndGroups[Token->DefaultOwnerIndex].Sid) + sizeof(TOKEN_OWNER);
-       if (TokenInformationLength < Length)
-         {
-           Status = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-           if (NT_SUCCESS(Status))
-             Status = STATUS_BUFFER_TOO_SMALL;
-         }
-       else
-         {
-           ((PTOKEN_OWNER)TokenInformation)->Owner = 
-             (PSID)(((PTOKEN_OWNER)TokenInformation) + 1);
-           RtlCopySid(TokenInformationLength - sizeof(TOKEN_OWNER),
-                      ((PTOKEN_OWNER)TokenInformation)->Owner,
-                      Token->UserAndGroups[Token->DefaultOwnerIndex].Sid);
-           Status = STATUS_SUCCESS;
-         }
-       break;
+      {
+        ULONG SidLen;
+        PTOKEN_OWNER to = (PTOKEN_OWNER)TokenInformation;
+
+        DPRINT("NtQueryInformationToken(TokenOwner)\n");
+        SidLen = RtlLengthSid(Token->UserAndGroups[Token->DefaultOwnerIndex].Sid);
+        RequiredLength = sizeof(TOKEN_OWNER) + SidLen;
+
+        _SEH_TRY
+        {
+          if(TokenInformationLength >= RequiredLength)
+          {
+            to->Owner = (PSID)(to + 1);
+            Status = RtlCopySid(SidLen,
+                                to->Owner,
+                                Token->UserAndGroups[Token->DefaultOwnerIndex].Sid);
+          }
+          else
+          {
+            Status = STATUS_BUFFER_TOO_SMALL;
+          }
+
+          if(ReturnLength != NULL)
+          {
+            *ReturnLength = RequiredLength;
+          }
+        }
+        _SEH_HANDLE
+        {
+          Status = _SEH_GetExceptionCode();
+        }
+        _SEH_END;
+
+        break;
+      }
 
       case TokenPrimaryGroup:
-       DPRINT("NtQueryInformationToken(TokenPrimaryGroup),"
-              "Token->PrimaryGroup = 0x%08x\n", Token->PrimaryGroup);
-       Length = RtlLengthSid(Token->PrimaryGroup) + sizeof(TOKEN_PRIMARY_GROUP);
-       if (TokenInformationLength < Length)
-         {
-           Status = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-           if (NT_SUCCESS(Status))
-             Status = STATUS_BUFFER_TOO_SMALL;
-         }
-       else
-         {
-           ((PTOKEN_PRIMARY_GROUP)TokenInformation)->PrimaryGroup = 
-             (PSID)(((PTOKEN_PRIMARY_GROUP)TokenInformation) + 1);
-           RtlCopySid(TokenInformationLength - sizeof(TOKEN_PRIMARY_GROUP),
-                      ((PTOKEN_PRIMARY_GROUP)TokenInformation)->PrimaryGroup,
-                      Token->PrimaryGroup);
-           Status = STATUS_SUCCESS;
-         }
-       break;
+      {
+        ULONG SidLen;
+        PTOKEN_PRIMARY_GROUP tpg = (PTOKEN_PRIMARY_GROUP)TokenInformation;
+
+        DPRINT("NtQueryInformationToken(TokenPrimaryGroup)\n");
+        SidLen = RtlLengthSid(Token->PrimaryGroup);
+        RequiredLength = sizeof(TOKEN_PRIMARY_GROUP) + SidLen;
+
+        _SEH_TRY
+        {
+          if(TokenInformationLength >= RequiredLength)
+          {
+            tpg->PrimaryGroup = (PSID)(tpg + 1);
+            Status = RtlCopySid(SidLen,
+                                tpg->PrimaryGroup,
+                                Token->PrimaryGroup);
+          }
+          else
+          {
+            Status = STATUS_BUFFER_TOO_SMALL;
+          }
+
+          if(ReturnLength != NULL)
+          {
+            *ReturnLength = RequiredLength;
+          }
+        }
+        _SEH_HANDLE
+        {
+          Status = _SEH_GetExceptionCode();
+        }
+        _SEH_END;
+
+        break;
+      }
 
       case TokenDefaultDacl:
-       DPRINT("NtQueryInformationToken(TokenDefaultDacl)\n");
-       PtrDefaultDacl = (PTOKEN_DEFAULT_DACL) TokenInformation;
-       Length = (Token->DefaultDacl ? Token->DefaultDacl->AclSize : 0) + sizeof(TOKEN_DEFAULT_DACL);
-       if (TokenInformationLength < Length)
-         {
-           Status = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-           if (NT_SUCCESS(Status))
-             Status = STATUS_BUFFER_TOO_SMALL;
-         }
-       else if (!Token->DefaultDacl)
-         {
-           PtrDefaultDacl->DefaultDacl = 0;
-           Status = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-         }
-       else
-         {
-           PtrDefaultDacl->DefaultDacl = (PACL) (PtrDefaultDacl + 1);
-           memmove(PtrDefaultDacl->DefaultDacl,
-                   Token->DefaultDacl,
-                   Token->DefaultDacl->AclSize);
-           Status = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-         }
-       break;
+      {
+        PTOKEN_DEFAULT_DACL tdd = (PTOKEN_DEFAULT_DACL)TokenInformation;
+
+        DPRINT("NtQueryInformationToken(TokenDefaultDacl)\n");
+        RequiredLength = sizeof(TOKEN_DEFAULT_DACL);
+
+        if(Token->DefaultDacl != NULL)
+        {
+          RequiredLength += Token->DefaultDacl->AclSize;
+        }
+
+        _SEH_TRY
+        {
+          if(TokenInformationLength >= RequiredLength)
+          {
+            if(Token->DefaultDacl != NULL)
+            {
+              tdd->DefaultDacl = (PACL)(tdd + 1);
+              RtlCopyMemory(tdd->DefaultDacl,
+                            Token->DefaultDacl,
+                            Token->DefaultDacl->AclSize);
+            }
+            else
+            {
+              tdd->DefaultDacl = NULL;
+            }
+          }
+          else
+          {
+            Status = STATUS_BUFFER_TOO_SMALL;
+          }
+
+          if(ReturnLength != NULL)
+          {
+            *ReturnLength = RequiredLength;
+          }
+        }
+        _SEH_HANDLE
+        {
+          Status = _SEH_GetExceptionCode();
+        }
+        _SEH_END;
+
+        break;
+      }
 
       case TokenSource:
-       DPRINT("NtQueryInformationToken(TokenSource)\n");
-       if (TokenInformationLength < sizeof(TOKEN_SOURCE))
-         {
-           Length = sizeof(TOKEN_SOURCE);
-           Status = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-           if (NT_SUCCESS(Status))
-             Status = STATUS_BUFFER_TOO_SMALL;
-         }
-       else
-         {
-           Status = MmCopyToCaller(TokenInformation, &Token->TokenSource, sizeof(TOKEN_SOURCE));
-         }
-       break;
+      {
+        PTOKEN_SOURCE ts = (PTOKEN_SOURCE)TokenInformation;
+
+        DPRINT("NtQueryInformationToken(TokenSource)\n");
+        RequiredLength = sizeof(TOKEN_SOURCE);
+
+        _SEH_TRY
+        {
+          if(TokenInformationLength >= RequiredLength)
+          {
+            *ts = Token->TokenSource;
+          }
+          else
+          {
+            Status = STATUS_BUFFER_TOO_SMALL;
+          }
+
+          if(ReturnLength != NULL)
+          {
+            *ReturnLength = RequiredLength;
+          }
+        }
+        _SEH_HANDLE
+        {
+          Status = _SEH_GetExceptionCode();
+        }
+        _SEH_END;
+
+        break;
+      }
 
       case TokenType:
-       DPRINT("NtQueryInformationToken(TokenType)\n");
-       if (TokenInformationLength < sizeof(TOKEN_TYPE))
-         {
-           Length = sizeof(TOKEN_TYPE);
-           Status = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-           if (NT_SUCCESS(Status))
-             Status = STATUS_BUFFER_TOO_SMALL;
-         }
-       else
-         {
-           Status = MmCopyToCaller(TokenInformation, &Token->TokenType, sizeof(TOKEN_TYPE));
-         }
-       break;
+      {
+        PTOKEN_TYPE tt = (PTOKEN_TYPE)TokenInformation;
+
+        DPRINT("NtQueryInformationToken(TokenType)\n");
+        RequiredLength = sizeof(TOKEN_TYPE);
+
+        _SEH_TRY
+        {
+          if(TokenInformationLength >= RequiredLength)
+          {
+            *tt = Token->TokenType;
+          }
+          else
+          {
+            Status = STATUS_BUFFER_TOO_SMALL;
+          }
+
+          if(ReturnLength != NULL)
+          {
+            *ReturnLength = RequiredLength;
+          }
+        }
+        _SEH_HANDLE
+        {
+          Status = _SEH_GetExceptionCode();
+        }
+        _SEH_END;
+
+        break;
+      }
 
       case TokenImpersonationLevel:
-       DPRINT("NtQueryInformationToken(TokenImpersonationLevel)\n");
-       if (TokenInformationLength < sizeof(SECURITY_IMPERSONATION_LEVEL))
-         {
-           Length = sizeof(SECURITY_IMPERSONATION_LEVEL);
-           Status = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-           if (NT_SUCCESS(Status))
-             Status = STATUS_BUFFER_TOO_SMALL;
-         }
-       else
-         {
-           Status = MmCopyToCaller(TokenInformation, &Token->ImpersonationLevel, sizeof(SECURITY_IMPERSONATION_LEVEL));
-         }
-       break;
+      {
+        PSECURITY_IMPERSONATION_LEVEL sil = (PSECURITY_IMPERSONATION_LEVEL)TokenInformation;
+
+        DPRINT("NtQueryInformationToken(TokenImpersonationLevel)\n");
+        RequiredLength = sizeof(SECURITY_IMPERSONATION_LEVEL);
+
+        _SEH_TRY
+        {
+          if(TokenInformationLength >= RequiredLength)
+          {
+            *sil = Token->ImpersonationLevel;
+          }
+          else
+          {
+            Status = STATUS_BUFFER_TOO_SMALL;
+          }
+
+          if(ReturnLength != NULL)
+          {
+            *ReturnLength = RequiredLength;
+          }
+        }
+        _SEH_HANDLE
+        {
+          Status = _SEH_GetExceptionCode();
+        }
+        _SEH_END;
+
+        break;
+      }
 
       case TokenStatistics:
-       DPRINT("NtQueryInformationToken(TokenStatistics)\n");
-       if (TokenInformationLength < sizeof(TOKEN_STATISTICS))
-         {
-           Length = sizeof(TOKEN_STATISTICS);
-           Status = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-           if (NT_SUCCESS(Status))
-             Status = STATUS_BUFFER_TOO_SMALL;
-         }
-       else
-         {
-           PtrTokenStatistics = (PTOKEN_STATISTICS)TokenInformation;
-           PtrTokenStatistics->TokenId = Token->TokenId;
-           PtrTokenStatistics->AuthenticationId = Token->AuthenticationId;
-           PtrTokenStatistics->ExpirationTime = Token->ExpirationTime;
-           PtrTokenStatistics->TokenType = Token->TokenType;
-           PtrTokenStatistics->ImpersonationLevel = Token->ImpersonationLevel;
-           PtrTokenStatistics->DynamicCharged = Token->DynamicCharged;
-           PtrTokenStatistics->DynamicAvailable = Token->DynamicAvailable;
-           PtrTokenStatistics->GroupCount = Token->UserAndGroupCount - 1;
-           PtrTokenStatistics->PrivilegeCount = Token->PrivilegeCount;
-           PtrTokenStatistics->ModifiedId = Token->ModifiedId;
-
-           Status = STATUS_SUCCESS;
-         }
-       break;
+      {
+        PTOKEN_STATISTICS ts = (PTOKEN_STATISTICS)TokenInformation;
+
+        DPRINT("NtQueryInformationToken(TokenStatistics)\n");
+        RequiredLength = sizeof(TOKEN_STATISTICS);
+
+        _SEH_TRY
+        {
+          if(TokenInformationLength >= RequiredLength)
+          {
+            ts->TokenId = Token->TokenId;
+            ts->AuthenticationId = Token->AuthenticationId;
+            ts->ExpirationTime = Token->ExpirationTime;
+            ts->TokenType = Token->TokenType;
+            ts->ImpersonationLevel = Token->ImpersonationLevel;
+            ts->DynamicCharged = Token->DynamicCharged;
+            ts->DynamicAvailable = Token->DynamicAvailable;
+            ts->GroupCount = Token->UserAndGroupCount - 1;
+            ts->PrivilegeCount = Token->PrivilegeCount;
+            ts->ModifiedId = Token->ModifiedId;
+          }
+          else
+          {
+            Status = STATUS_BUFFER_TOO_SMALL;
+          }
+
+          if(ReturnLength != NULL)
+          {
+            *ReturnLength = RequiredLength;
+          }
+        }
+        _SEH_HANDLE
+        {
+          Status = _SEH_GetExceptionCode();
+        }
+        _SEH_END;
+
+        break;
+      }
 
       case TokenOrigin:
-       DPRINT("NtQueryInformationToken(TokenOrigin)\n");
-       if (TokenInformationLength < sizeof(TOKEN_ORIGIN))
-         {
-           Status = STATUS_BUFFER_TOO_SMALL;
-         }
-       else
-         {
-           Status = MmCopyToCaller(&((PTOKEN_ORIGIN)TokenInformation)->OriginatingLogonSession,
-                                   &Token->AuthenticationId, sizeof(LUID));
-         }
-       Length = sizeof(TOKEN_ORIGIN);
-       LengthStatus = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-       if (NT_SUCCESS(Status))
-         {
-           Status = LengthStatus;
-         }
-       break;
+      {
+        PTOKEN_ORIGIN to = (PTOKEN_ORIGIN)TokenInformation;
+
+        DPRINT("NtQueryInformationToken(TokenOrigin)\n");
+        RequiredLength = sizeof(TOKEN_ORIGIN);
+
+        _SEH_TRY
+        {
+          if(TokenInformationLength >= RequiredLength)
+          {
+            RtlCopyLuid(&to->OriginatingLogonSession,
+                        &Token->AuthenticationId);
+          }
+          else
+          {
+            Status = STATUS_BUFFER_TOO_SMALL;
+          }
+
+          if(ReturnLength != NULL)
+          {
+            *ReturnLength = RequiredLength;
+          }
+        }
+        _SEH_HANDLE
+        {
+          Status = _SEH_GetExceptionCode();
+        }
+        _SEH_END;
+
+        break;
+      }
 
       case TokenGroupsAndPrivileges:
        DPRINT1("NtQueryInformationToken(TokenGroupsAndPrivileges) not implemented\n");
@@ -828,9 +1056,49 @@ NtQueryInformationToken(IN HANDLE TokenHandle,
        break;
 
       case TokenRestrictedSids:
-       DPRINT1("NtQueryInformationToken(TokenRestrictedSids) not implemented\n");
-       Status = STATUS_NOT_IMPLEMENTED;
+      {
+        PTOKEN_GROUPS tg = (PTOKEN_GROUPS)TokenInformation;
+
+        DPRINT("NtQueryInformationToken(TokenRestrictedSids)\n");
+        RequiredLength = sizeof(tg->GroupCount) +
+                         RtlLengthSidAndAttributes(Token->RestrictedSidCount, Token->RestrictedSids);
+
+        _SEH_TRY
+        {
+          if(TokenInformationLength >= RequiredLength)
+          {
+            ULONG SidLen = RequiredLength - sizeof(tg->GroupCount) -
+                           (Token->RestrictedSidCount * sizeof(SID_AND_ATTRIBUTES));
+            PSID_AND_ATTRIBUTES Sid = (PSID_AND_ATTRIBUTES)((ULONG_PTR)TokenInformation + sizeof(tg->GroupCount) +
+                                                            (Token->RestrictedSidCount * sizeof(SID_AND_ATTRIBUTES)));
+
+            tg->GroupCount = Token->RestrictedSidCount;
+            Status = RtlCopySidAndAttributesArray(Token->RestrictedSidCount,
+                                                  Token->RestrictedSids,
+                                                  SidLen,
+                                                  &tg->Groups[0],
+                                                  (PSID)Sid,
+                                                  &Unused.Ptr,
+                                                  &Unused.Ulong);
+          }
+          else
+          {
+            Status = STATUS_BUFFER_TOO_SMALL;
+          }
+
+          if(ReturnLength != NULL)
+          {
+            *ReturnLength = RequiredLength;
+          }
+        }
+        _SEH_HANDLE
+        {
+          Status = _SEH_GetExceptionCode();
+        }
+        _SEH_END;
+
        break;
+      }
 
       case TokenSandBoxInert:
        DPRINT1("NtQueryInformationToken(TokenSandboxInert) not implemented\n");
@@ -838,27 +1106,44 @@ NtQueryInformationToken(IN HANDLE TokenHandle,
        break;
 
       case TokenSessionId:
-       DPRINT("NtQueryInformationToken(TokenSessionId)\n");
-       if (TokenInformationLength < sizeof(ULONG))
-         {
-           Length = sizeof(ULONG);
-           Status = MmCopyToCaller(ReturnLength, &Length, sizeof(ULONG));
-           if (NT_SUCCESS(Status))
-             Status = STATUS_BUFFER_TOO_SMALL;
-         }
-       else
-         {
-           Status = MmCopyToCaller(TokenInformation, &Token->SessionId, sizeof(ULONG));
-         }
-       break;
+      {
+        ULONG SessionId = 0;
+
+        DPRINT("NtQueryInformationToken(TokenSessionId)\n");
+
+        Status = SeQuerySessionIdToken(Token,
+                                       &SessionId);
+
+        if(NT_SUCCESS(Status))
+        {
+          _SEH_TRY
+          {
+            /* buffer size was already verified, no need to check here again */
+            *(PULONG)TokenInformation = SessionId;
+
+            if(ReturnLength != NULL)
+            {
+              *ReturnLength = sizeof(ULONG);
+            }
+          }
+          _SEH_HANDLE
+          {
+            Status = _SEH_GetExceptionCode();
+          }
+          _SEH_END;
+        }
+
+        break;
+      }
 
       default:
-       DPRINT1("NtQueryInformationToken(%d) invalid parameter\n");
-       Status = STATUS_INVALID_PARAMETER;
+       DPRINT1("NtQueryInformationToken(%d) invalid information class\n", TokenInformationClass);
+       Status = STATUS_INVALID_INFO_CLASS;
        break;
     }
 
-  ObDereferenceObject(Token);
+    ObDereferenceObject(Token);
+  }
 
   return(Status);
 }
@@ -879,7 +1164,7 @@ SeQueryInformationToken(
 }
 
 /*
- * @unimplemented
+ * @implemented
  */
 NTSTATUS
 STDCALL
@@ -888,14 +1173,14 @@ SeQuerySessionIdToken(
        IN PULONG pSessionId
        )
 {
-       UNIMPLEMENTED;
-       return STATUS_NOT_IMPLEMENTED;
+  *pSessionId = ((PTOKEN)Token)->SessionId;
+  return STATUS_SUCCESS;
 }
 
 /*
  * NtSetTokenInformation: Partly implemented.
  * Unimplemented:
- *  TokenOrigin, TokenDefaultDacl, TokenSessionId
+ *  TokenOrigin, TokenDefaultDacl
  */
 
 NTSTATUS STDCALL
@@ -904,121 +1189,229 @@ NtSetInformationToken(IN HANDLE TokenHandle,
                      OUT PVOID TokenInformation,
                      IN ULONG TokenInformationLength)
 {
-  NTSTATUS Status;
   PTOKEN Token;
-  TOKEN_OWNER TokenOwnerSet = { 0 };
-  TOKEN_PRIMARY_GROUP TokenPrimaryGroupSet = { 0 };
-  DWORD NeededAccess = 0;
+  KPROCESSOR_MODE PreviousMode;
+  ULONG NeededAccess = TOKEN_ADJUST_DEFAULT;
+  NTSTATUS Status = STATUS_SUCCESS;
 
-  switch (TokenInformationClass) 
-    {
-    case TokenOwner:
-    case TokenPrimaryGroup:
-      NeededAccess = TOKEN_ADJUST_DEFAULT;
-      break;
+  PAGED_CODE();
 
-    case TokenDefaultDacl:
-      if (TokenInformationLength < sizeof(TOKEN_DEFAULT_DACL))
-        return STATUS_BUFFER_TOO_SMALL;
-      NeededAccess = TOKEN_ADJUST_DEFAULT;
-      break;
+  PreviousMode = ExGetPreviousMode();
 
-    default:
-      DPRINT1("NtSetInformationToken: lying about success (stub) - %x\n", TokenInformationClass);
-      return STATUS_SUCCESS;  
+  Status = DefaultSetInfoBufferCheck(TokenInformationClass,
+                                     SeTokenInformationClass,
+                                     sizeof(SeTokenInformationClass) / sizeof(SeTokenInformationClass[0]),
+                                     TokenInformation,
+                                     TokenInformationLength,
+                                     PreviousMode);
 
-    }
+  if(!NT_SUCCESS(Status))
+  {
+    /* Invalid buffers */
+    DPRINT("NtSetInformationToken() failed, Status: 0x%x\n", Status);
+    return Status;
+  }
+
+  if(TokenInformationClass == TokenSessionId)
+  {
+    NeededAccess |= TOKEN_ADJUST_SESSIONID;
+  }
 
   Status = ObReferenceObjectByHandle(TokenHandle,
                                     NeededAccess,
                                     SepTokenObjectType,
-                                    UserMode,
+                                    PreviousMode,
                                     (PVOID*)&Token,
                                     NULL);
-  if (!NT_SUCCESS(Status))
+  if (NT_SUCCESS(Status))
+  {
+    switch (TokenInformationClass)
     {
-      return(Status);
-    }
+      case TokenOwner:
+      {
+        if(TokenInformationLength >= sizeof(TOKEN_OWNER))
+        {
+          PTOKEN_OWNER to = (PTOKEN_OWNER)TokenInformation;
+          PSID InputSid = NULL;
 
-  switch (TokenInformationClass)
-    {
-    case TokenOwner:
-      MmCopyFromCaller( &TokenOwnerSet, TokenInformation,
-                       min(sizeof(TokenOwnerSet),TokenInformationLength) );
-      RtlCopySid(TokenInformationLength - sizeof(TOKEN_OWNER),
-                Token->UserAndGroups[Token->DefaultOwnerIndex].Sid,
-                TokenOwnerSet.Owner);
-      Status = STATUS_SUCCESS;
-      DPRINT("NtSetInformationToken(TokenOwner)\n");
-      break;
-      
-    case TokenPrimaryGroup:
-      MmCopyFromCaller( &TokenPrimaryGroupSet, TokenInformation, 
-                       min(sizeof(TokenPrimaryGroupSet),
-                           TokenInformationLength) );
-      RtlCopySid(TokenInformationLength - sizeof(TOKEN_PRIMARY_GROUP),
-                Token->PrimaryGroup,
-                TokenPrimaryGroupSet.PrimaryGroup);
-      Status = STATUS_SUCCESS;
-      DPRINT("NtSetInformationToken(TokenPrimaryGroup),"
-            "Token->PrimaryGroup = 0x%08x\n", Token->PrimaryGroup);
-      break;
+          _SEH_TRY
+          {
+            InputSid = to->Owner;
+          }
+          _SEH_HANDLE
+          {
+            Status = _SEH_GetExceptionCode();
+          }
+          _SEH_END;
+
+          if(NT_SUCCESS(Status))
+          {
+            PSID CapturedSid;
+
+            Status = SepCaptureSid(InputSid,
+                                   PreviousMode,
+                                   PagedPool,
+                                   FALSE,
+                                   &CapturedSid);
+            if(NT_SUCCESS(Status))
+            {
+              RtlCopySid(RtlLengthSid(CapturedSid),
+                         Token->UserAndGroups[Token->DefaultOwnerIndex].Sid,
+                         CapturedSid);
+              SepReleaseSid(CapturedSid,
+                            PreviousMode,
+                            FALSE);
+            }
+          }
+        }
+        else
+        {
+          Status = STATUS_INFO_LENGTH_MISMATCH;
+        }
+        break;
+      }
 
-    case TokenDefaultDacl:
+      case TokenPrimaryGroup:
       {
-        TOKEN_DEFAULT_DACL TokenDefaultDacl = { 0 };
-        ACL OldAcl;
-        PACL NewAcl;
+        if(TokenInformationLength >= sizeof(TOKEN_PRIMARY_GROUP))
+        {
+          PTOKEN_PRIMARY_GROUP tpg = (PTOKEN_PRIMARY_GROUP)TokenInformation;
+          PSID InputSid = NULL;
 
-        Status = MmCopyFromCaller( &TokenDefaultDacl, TokenInformation, 
-                                   sizeof(TOKEN_DEFAULT_DACL) );
-        if (!NT_SUCCESS(Status))
+          _SEH_TRY
           {
-            Status = STATUS_INVALID_PARAMETER;
-            break;
+            InputSid = tpg->PrimaryGroup;
           }
-
-        Status = MmCopyFromCaller( &OldAcl, TokenDefaultDacl.DefaultDacl,
-                                   sizeof(ACL) );
-        if (!NT_SUCCESS(Status))
+          _SEH_HANDLE
           {
-            Status = STATUS_INVALID_PARAMETER;
-            break;
+            Status = _SEH_GetExceptionCode();
           }
+          _SEH_END;
 
-        NewAcl = ExAllocatePool(NonPagedPool, sizeof(ACL));
-        if (NewAcl == NULL)
+          if(NT_SUCCESS(Status))
           {
-            Status = STATUS_INSUFFICIENT_RESOURCES;
-            break;
+            PSID CapturedSid;
+
+            Status = SepCaptureSid(InputSid,
+                                   PreviousMode,
+                                   PagedPool,
+                                   FALSE,
+                                   &CapturedSid);
+            if(NT_SUCCESS(Status))
+            {
+              RtlCopySid(RtlLengthSid(CapturedSid),
+                         Token->PrimaryGroup,
+                         CapturedSid);
+              SepReleaseSid(CapturedSid,
+                            PreviousMode,
+                            FALSE);
+            }
           }
+        }
+        else
+        {
+          Status = STATUS_INFO_LENGTH_MISMATCH;
+        }
+        break;
+      }
+
+      case TokenDefaultDacl:
+      {
+        if(TokenInformationLength >= sizeof(TOKEN_DEFAULT_DACL))
+        {
+          PTOKEN_DEFAULT_DACL tdd = (PTOKEN_DEFAULT_DACL)TokenInformation;
+          PACL InputAcl = NULL;
 
-        Status = MmCopyFromCaller( NewAcl, TokenDefaultDacl.DefaultDacl,
-                                   OldAcl.AclSize );
-        if (!NT_SUCCESS(Status))
+          _SEH_TRY
           {
-            Status = STATUS_INVALID_PARAMETER;
-            ExFreePool(NewAcl);
-            break;
+            InputAcl = tdd->DefaultDacl;
           }
+          _SEH_HANDLE
+          {
+            Status = _SEH_GetExceptionCode();
+          }
+          _SEH_END;
 
-        if (Token->DefaultDacl)
+          if(NT_SUCCESS(Status))
           {
-            ExFreePool(Token->DefaultDacl);
+            if(InputAcl != NULL)
+            {
+              PACL CapturedAcl;
+
+              /* capture and copy the dacl */
+              Status = SepCaptureAcl(InputAcl,
+                                     PreviousMode,
+                                     PagedPool,
+                                     TRUE,
+                                     &CapturedAcl);
+              if(NT_SUCCESS(Status))
+              {
+                /* free the previous dacl if present */
+                if(Token->DefaultDacl != NULL)
+                {
+                  ExFreePool(Token->DefaultDacl);
+                }
+
+                /* set the new dacl */
+                Token->DefaultDacl = CapturedAcl;
+              }
+            }
+            else
+            {
+              /* clear and free the default dacl if present */
+              if(Token->DefaultDacl != NULL)
+              {
+                ExFreePool(Token->DefaultDacl);
+                Token->DefaultDacl = NULL;
+              }
+            }
           }
+        }
+        else
+        {
+          Status = STATUS_INFO_LENGTH_MISMATCH;
+        }
+        break;
+      }
+
+      case TokenSessionId:
+      {
+        ULONG SessionId = 0;
+
+        _SEH_TRY
+        {
+          /* buffer size was already verified, no need to check here again */
+          SessionId = *(PULONG)TokenInformation;
+        }
+        _SEH_HANDLE
+        {
+          Status = _SEH_GetExceptionCode();
+        }
+        _SEH_END;
 
-        Token->DefaultDacl = NewAcl;
+        if(NT_SUCCESS(Status))
+        {
+          if(!SeSinglePrivilegeCheck(SeTcbPrivilege,
+                                     PreviousMode))
+          {
+            Status = STATUS_PRIVILEGE_NOT_HELD;
+            break;
+          }
 
-        Status = STATUS_SUCCESS;
+          Token->SessionId = SessionId;
+        }
         break;
       }
 
-    default:
-      Status = STATUS_NOT_IMPLEMENTED;
-      break;
+      default:
+      {
+        Status = STATUS_NOT_IMPLEMENTED;
+        break;
+      }
     }
 
-  ObDereferenceObject(Token);
+    ObDereferenceObject(Token);
+  }
 
   return(Status);
 }
@@ -1034,64 +1427,104 @@ NtSetInformationToken(IN HANDLE TokenHandle,
  */
 NTSTATUS STDCALL
 NtDuplicateToken(IN HANDLE ExistingTokenHandle,
-                IN ACCESS_MASK DesiredAccess,
-       IN POBJECT_ATTRIBUTES ObjectAttributes OPTIONAL /*is it really optional?*/,
-       IN BOOLEAN EffectiveOnly,
-                IN TOKEN_TYPE TokenType,
-                OUT PHANDLE NewTokenHandle)
+                 IN ACCESS_MASK DesiredAccess,
+                 IN POBJECT_ATTRIBUTES ObjectAttributes  OPTIONAL,
+                 IN BOOLEAN EffectiveOnly,
+                 IN TOKEN_TYPE TokenType,
+                 OUT PHANDLE NewTokenHandle)
 {
   KPROCESSOR_MODE PreviousMode;
+  HANDLE hToken;
   PTOKEN Token;
   PTOKEN NewToken;
-  NTSTATUS Status;
+  PSECURITY_QUALITY_OF_SERVICE CapturedSecurityQualityOfService;
+  BOOLEAN QoSPresent;
+  NTSTATUS Status = STATUS_SUCCESS;
+
+  PAGED_CODE();
 
   PreviousMode = KeGetPreviousMode();
+
+  if(PreviousMode != KernelMode)
+  {
+    _SEH_TRY
+    {
+      ProbeForWriteHandle(NewTokenHandle);
+    }
+    _SEH_HANDLE
+    {
+      Status = _SEH_GetExceptionCode();
+    }
+    _SEH_END;
+
+    if(!NT_SUCCESS(Status))
+    {
+      return Status;
+    }
+  }
+
+  Status = SepCaptureSecurityQualityOfService(ObjectAttributes,
+                                              PreviousMode,
+                                              PagedPool,
+                                              FALSE,
+                                              &CapturedSecurityQualityOfService,
+                                              &QoSPresent);
+  if(!NT_SUCCESS(Status))
+  {
+    DPRINT1("NtDuplicateToken() failed to capture QoS! Status: 0x%x\n", Status);
+    return Status;
+  }
+
   Status = ObReferenceObjectByHandle(ExistingTokenHandle,
                                     TOKEN_DUPLICATE,
                                     SepTokenObjectType,
                                     PreviousMode,
                                     (PVOID*)&Token,
                                     NULL);
-  if (!NT_SUCCESS(Status))
+  if (NT_SUCCESS(Status))
+  {
+    Status = SepDuplicateToken(Token,
+                               ObjectAttributes,
+                               EffectiveOnly,
+                               TokenType,
+                               (QoSPresent ? CapturedSecurityQualityOfService->ImpersonationLevel : SecurityAnonymous),
+                              PreviousMode,
+                              &NewToken);
+
+    ObDereferenceObject(Token);
+
+    if (NT_SUCCESS(Status))
     {
-      DPRINT1("Failed to reference token (Status %lx)\n", Status);
-      return Status;
-    }
-
-  Status = SepDuplicateToken(Token,
-                            ObjectAttributes,
-                            EffectiveOnly,
-                            TokenType,
-              ObjectAttributes->SecurityQualityOfService ? 
-                  ((PSECURITY_QUALITY_OF_SERVICE)(ObjectAttributes->SecurityQualityOfService))->ImpersonationLevel : 
-                  0 /*SecurityAnonymous*/,
-                            PreviousMode,
-                            &NewToken);
+      Status = ObInsertObject((PVOID)NewToken,
+                         NULL,
+                         DesiredAccess,
+                         0,
+                         NULL,
+                         &hToken);
 
-  ObDereferenceObject(Token);
+      ObDereferenceObject(NewToken);
 
-  if (!NT_SUCCESS(Status))
-    {
-      DPRINT1("Failed to duplicate token (Status %lx)\n", Status);
-      return Status;
+      if (NT_SUCCESS(Status))
+      {
+        _SEH_TRY
+        {
+          *NewTokenHandle = hToken;
+        }
+        _SEH_HANDLE
+        {
+          Status = _SEH_GetExceptionCode();
+        }
+        _SEH_END;
+      }
     }
+  }
 
-  Status = ObInsertObject((PVOID)NewToken,
-                         NULL,
-                         DesiredAccess,
-                         0,
-                         NULL,
-                         NewTokenHandle);
-
-  ObDereferenceObject(NewToken);
-
-  if (!NT_SUCCESS(Status))
-    {
-      DPRINT1("Failed to create token handle (Status %lx)\n");
-      return Status;
-    }
+  /* free the captured structure */
+  SepReleaseSecurityQualityOfService(CapturedSecurityQualityOfService,
+                                     PreviousMode,
+                                     FALSE);
 
-  return STATUS_SUCCESS;
+  return Status;
 }
 
 
@@ -1124,15 +1557,17 @@ NtAdjustGroupsToken(IN HANDLE TokenHandle,
    ULONG a;
    ULONG b;
    ULONG c;
-   
+
+   PAGED_CODE();
+
    Status = ObReferenceObjectByHandle(TokenHandle,
                                      ?,
                                      SepTokenObjectType,
                                      UserMode,
                                      (PVOID*)&Token,
                                      NULL);
-   
-   
+
+
    SepAdjustGroups(Token,
                   0,
                   ResetToDefault,
@@ -1217,7 +1652,7 @@ NtAdjustPrivilegesToken (IN HANDLE TokenHandle,
 {
 //  PLUID_AND_ATTRIBUTES Privileges;
   KPROCESSOR_MODE PreviousMode;
-//  ULONG PrivilegeCount;
+  ULONG PrivilegeCount;
   PTOKEN Token;
 //  ULONG Length;
   ULONG i;
@@ -1231,6 +1666,8 @@ NtAdjustPrivilegesToken (IN HANDLE TokenHandle,
 #endif
   NTSTATUS Status;
 
+  PAGED_CODE();
+
   DPRINT ("NtAdjustPrivilegesToken() called\n");
 
 //  PrivilegeCount = NewState->PrivilegeCount;
@@ -1246,7 +1683,7 @@ NtAdjustPrivilegesToken (IN HANDLE TokenHandle,
 //                               &Length);
 
   Status = ObReferenceObjectByHandle (TokenHandle,
-                                     TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
+                                     TOKEN_ADJUST_PRIVILEGES | (PreviousState != NULL ? TOKEN_QUERY : 0),
                                      SepTokenObjectType,
                                      PreviousMode,
                                      (PVOID*)&Token,
@@ -1273,6 +1710,12 @@ NtAdjustPrivilegesToken (IN HANDLE TokenHandle,
                       &c);
 #endif
 
+  PrivilegeCount = (BufferLength - FIELD_OFFSET(TOKEN_PRIVILEGES, Privileges)) /
+                   sizeof(LUID_AND_ATTRIBUTES);
+
+  if (PreviousState != NULL)
+    PreviousState->PrivilegeCount = 0;
+
   k = 0;
   if (DisableAllPrivileges == TRUE)
     {
@@ -1283,11 +1726,22 @@ NtAdjustPrivilegesToken (IN HANDLE TokenHandle,
              DPRINT ("Attributes differ\n");
 
              /* Save current privilege */
-             if (PreviousState != NULL && k < PreviousState->PrivilegeCount)
+             if (PreviousState != NULL)
                {
-                 PreviousState->Privileges[k].Luid = Token->Privileges[i].Luid;
-                 PreviousState->Privileges[k].Attributes = Token->Privileges[i].Attributes;
-                 k++;
+                  if (k < PrivilegeCount)
+                    {
+                      PreviousState->PrivilegeCount++;
+                      PreviousState->Privileges[k].Luid = Token->Privileges[i].Luid;
+                      PreviousState->Privileges[k].Attributes = Token->Privileges[i].Attributes;
+                    }
+                  else
+                    {
+                      /* FIXME: Should revert all the changes, calculate how
+                       * much space would be needed, set ResultLength
+                       * accordingly and fail.
+                       */
+                    }
+                  k++;
                }
 
              /* Update current privlege */
@@ -1317,16 +1771,27 @@ NtAdjustPrivilegesToken (IN HANDLE TokenHandle,
                              NewState->Privileges[j].Attributes);
 
                      /* Save current privilege */
-                     if (PreviousState != NULL && k < PreviousState->PrivilegeCount)
+                     if (PreviousState != NULL)
                        {
-                         PreviousState->Privileges[k].Luid = Token->Privileges[i].Luid;
-                         PreviousState->Privileges[k].Attributes = Token->Privileges[i].Attributes;
-                         k++;
+                          if (k < PrivilegeCount)
+                            {
+                              PreviousState->PrivilegeCount++;
+                              PreviousState->Privileges[k].Luid = Token->Privileges[i].Luid;
+                              PreviousState->Privileges[k].Attributes = Token->Privileges[i].Attributes;
+                            }
+                          else
+                            {
+                              /* FIXME: Should revert all the changes, calculate how
+                               * much space would be needed, set ResultLength
+                               * accordingly and fail.
+                               */
+                            }
+                          k++;
                        }
 
                      /* Update current privlege */
                      Token->Privileges[i].Attributes &= ~SE_PRIVILEGE_ENABLED;
-                     Token->Privileges[i].Attributes |= 
+                     Token->Privileges[i].Attributes |=
                        (NewState->Privileges[j].Attributes & SE_PRIVILEGE_ENABLED);
                      DPRINT ("New attributes %lx\n",
                              Token->Privileges[i].Attributes);
@@ -1355,23 +1820,27 @@ NtAdjustPrivilegesToken (IN HANDLE TokenHandle,
   return Status;
 }
 
-
-NTSTATUS
-SepCreateSystemProcessToken(struct _EPROCESS* Process)
+PTOKEN
+STDCALL
+SepCreateSystemProcessToken(VOID)
 {
   NTSTATUS Status;
   ULONG uSize;
   ULONG i;
-
-  ULONG uLocalSystemLength = RtlLengthSid(SeLocalSystemSid);
-  ULONG uWorldLength       = RtlLengthSid(SeWorldSid);
-  ULONG uAuthUserLength    = RtlLengthSid(SeAuthenticatedUserSid);
-  ULONG uAdminsLength      = RtlLengthSid(SeAliasAdminsSid);
-
+  ULONG uLocalSystemLength;
+  ULONG uWorldLength;
+  ULONG uAuthUserLength;
+  ULONG uAdminsLength;
   PTOKEN AccessToken;
-
   PVOID SidArea;
 
+  PAGED_CODE();
+
+  uLocalSystemLength = RtlLengthSid(SeLocalSystemSid);
+  uWorldLength       = RtlLengthSid(SeWorldSid);
+  uAuthUserLength    = RtlLengthSid(SeAuthenticatedUserSid);
+  uAdminsLength      = RtlLengthSid(SeAliasAdminsSid);
+
  /*
   * Initialize the token
   */
@@ -1386,30 +1855,38 @@ SepCreateSystemProcessToken(struct _EPROCESS* Process)
                          (PVOID*)&AccessToken);
   if (!NT_SUCCESS(Status))
     {
-      return(Status);
+      return NULL;
     }
-
-  Status = NtAllocateLocallyUniqueId(&AccessToken->TokenId);
+  Status = ObInsertObject(AccessToken,
+                          NULL,
+                          TOKEN_ALL_ACCESS,
+                          0,
+                          NULL,
+                          NULL);
+
+  Status = ExpAllocateLocallyUniqueId(&AccessToken->TokenId);
   if (!NT_SUCCESS(Status))
     {
       ObDereferenceObject(AccessToken);
-      return(Status);
+      return NULL;
     }
 
-  Status = NtAllocateLocallyUniqueId(&AccessToken->ModifiedId);
+  Status = ExpAllocateLocallyUniqueId(&AccessToken->ModifiedId);
   if (!NT_SUCCESS(Status))
     {
       ObDereferenceObject(AccessToken);
-      return(Status);
+      return NULL;
     }
 
-  Status = NtAllocateLocallyUniqueId(&AccessToken->AuthenticationId);
+  Status = ExpAllocateLocallyUniqueId(&AccessToken->AuthenticationId);
   if (!NT_SUCCESS(Status))
     {
       ObDereferenceObject(AccessToken);
-      return Status;
+      return NULL;
     }
 
+  AccessToken->TokenLock = &SepTokenLock;
+
   AccessToken->TokenType = TokenPrimary;
   AccessToken->ImpersonationLevel = SecurityDelegation;
   AccessToken->TokenSource.SourceIdentifier.LowPart = 0;
@@ -1424,8 +1901,8 @@ SepCreateSystemProcessToken(struct _EPROCESS* Process)
   uSize += uAuthUserLength;
   uSize += uAdminsLength;
 
-  AccessToken->UserAndGroups = 
-    (PSID_AND_ATTRIBUTES)ExAllocatePoolWithTag(NonPagedPool,
+  AccessToken->UserAndGroups =
+    (PSID_AND_ATTRIBUTES)ExAllocatePoolWithTag(PagedPool,
                                               uSize,
                                               TAG('T', 'O', 'K', 'u'));
   SidArea = &AccessToken->UserAndGroups[AccessToken->UserAndGroupCount];
@@ -1457,7 +1934,7 @@ SepCreateSystemProcessToken(struct _EPROCESS* Process)
 
   uSize = AccessToken->PrivilegeCount * sizeof(LUID_AND_ATTRIBUTES);
   AccessToken->Privileges =
-       (PLUID_AND_ATTRIBUTES)ExAllocatePoolWithTag(NonPagedPool,
+       (PLUID_AND_ATTRIBUTES)ExAllocatePoolWithTag(PagedPool,
                                                    uSize,
                                                    TAG('T', 'O', 'K', 'p'));
 
@@ -1470,55 +1947,55 @@ SepCreateSystemProcessToken(struct _EPROCESS* Process)
 
   AccessToken->Privileges[i].Attributes = 0;
   AccessToken->Privileges[i++].Luid = SeTakeOwnershipPrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = SE_PRIVILEGE_ENABLED_BY_DEFAULT|SE_PRIVILEGE_ENABLED;
   AccessToken->Privileges[i++].Luid = SeCreatePagefilePrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = SE_PRIVILEGE_ENABLED_BY_DEFAULT|SE_PRIVILEGE_ENABLED;
   AccessToken->Privileges[i++].Luid = SeLockMemoryPrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = 0;
   AccessToken->Privileges[i++].Luid = SeAssignPrimaryTokenPrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = 0;
   AccessToken->Privileges[i++].Luid = SeIncreaseQuotaPrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = SE_PRIVILEGE_ENABLED_BY_DEFAULT|SE_PRIVILEGE_ENABLED;
   AccessToken->Privileges[i++].Luid = SeIncreaseBasePriorityPrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = SE_PRIVILEGE_ENABLED_BY_DEFAULT|SE_PRIVILEGE_ENABLED;
   AccessToken->Privileges[i++].Luid = SeCreatePermanentPrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = SE_PRIVILEGE_ENABLED_BY_DEFAULT|SE_PRIVILEGE_ENABLED;
   AccessToken->Privileges[i++].Luid = SeDebugPrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = SE_PRIVILEGE_ENABLED_BY_DEFAULT|SE_PRIVILEGE_ENABLED;
   AccessToken->Privileges[i++].Luid = SeAuditPrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = 0;
   AccessToken->Privileges[i++].Luid = SeSecurityPrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = 0;
   AccessToken->Privileges[i++].Luid = SeSystemEnvironmentPrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = SE_PRIVILEGE_ENABLED_BY_DEFAULT|SE_PRIVILEGE_ENABLED;
   AccessToken->Privileges[i++].Luid = SeChangeNotifyPrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = 0;
   AccessToken->Privileges[i++].Luid = SeBackupPrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = 0;
   AccessToken->Privileges[i++].Luid = SeRestorePrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = 0;
   AccessToken->Privileges[i++].Luid = SeShutdownPrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = 0;
   AccessToken->Privileges[i++].Luid = SeLoadDriverPrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = SE_PRIVILEGE_ENABLED_BY_DEFAULT|SE_PRIVILEGE_ENABLED;
   AccessToken->Privileges[i++].Luid = SeProfileSingleProcessPrivilege;
-  
+
   AccessToken->Privileges[i].Attributes = 0;
   AccessToken->Privileges[i++].Luid = SeSystemtimePrivilege;
 #if 0
@@ -1536,7 +2013,7 @@ SepCreateSystemProcessToken(struct _EPROCESS* Process)
   uSize += sizeof(ACE) + uAdminsLength;
   uSize = (uSize & (~3)) + 8;
   AccessToken->DefaultDacl =
-    (PACL) ExAllocatePoolWithTag(NonPagedPool,
+    (PACL) ExAllocatePoolWithTag(PagedPool,
                                 uSize,
                                 TAG('T', 'O', 'K', 'd'));
   Status = RtlCreateAcl(AccessToken->DefaultDacl, uSize, ACL_REVISION);
@@ -1553,18 +2030,17 @@ SepCreateSystemProcessToken(struct _EPROCESS* Process)
   if ( ! NT_SUCCESS(Status) )
     {
       ObDereferenceObject(AccessToken);
-      return Status;
+      return NULL;
     }
 
-  Process->Token = AccessToken;
-  return(STATUS_SUCCESS);
+  return AccessToken;
 }
 
 
 NTSTATUS STDCALL
-NtCreateToken(OUT PHANDLE UnsafeTokenHandle,
+NtCreateToken(OUT PHANDLE TokenHandle,
              IN ACCESS_MASK DesiredAccess,
-             IN POBJECT_ATTRIBUTES UnsafeObjectAttributes,
+             IN POBJECT_ATTRIBUTES ObjectAttributes,
              IN TOKEN_TYPE TokenType,
              IN PLUID AuthenticationId,
              IN PLARGE_INTEGER ExpirationTime,
@@ -1576,24 +2052,70 @@ NtCreateToken(OUT PHANDLE UnsafeTokenHandle,
              IN PTOKEN_DEFAULT_DACL TokenDefaultDacl,
              IN PTOKEN_SOURCE TokenSource)
 {
-  HANDLE TokenHandle;
+  HANDLE hToken;
   PTOKEN AccessToken;
-  NTSTATUS Status;
-  OBJECT_ATTRIBUTES SafeObjectAttributes;
-  POBJECT_ATTRIBUTES ObjectAttributes;
   LUID TokenId;
   LUID ModifiedId;
   PVOID EndMem;
   ULONG uLength;
   ULONG i;
+  ULONG nTokenPrivileges = 0;
+  LARGE_INTEGER LocalExpirationTime = {};
+  KPROCESSOR_MODE PreviousMode;
+  NTSTATUS Status = STATUS_SUCCESS;
 
-  Status = MmCopyFromCaller(&SafeObjectAttributes,
-                           UnsafeObjectAttributes,
-                           sizeof(OBJECT_ATTRIBUTES));
-  if (!NT_SUCCESS(Status))
-    return(Status);
+  PAGED_CODE();
+
+  PreviousMode = ExGetPreviousMode();
+
+  if(PreviousMode != KernelMode)
+  {
+    _SEH_TRY
+    {
+      ProbeForWriteHandle(TokenHandle);
+      ProbeForRead(AuthenticationId,
+                   sizeof(LUID),
+                   sizeof(ULONG));
+      LocalExpirationTime = ProbeForReadLargeInteger(ExpirationTime);
+      ProbeForRead(TokenUser,
+                   sizeof(TOKEN_USER),
+                   sizeof(ULONG));
+      ProbeForRead(TokenGroups,
+                   sizeof(TOKEN_GROUPS),
+                   sizeof(ULONG));
+      ProbeForRead(TokenPrivileges,
+                   sizeof(TOKEN_PRIVILEGES),
+                   sizeof(ULONG));
+      ProbeForRead(TokenOwner,
+                   sizeof(TOKEN_OWNER),
+                   sizeof(ULONG));
+      ProbeForRead(TokenPrimaryGroup,
+                   sizeof(TOKEN_PRIMARY_GROUP),
+                   sizeof(ULONG));
+      ProbeForRead(TokenDefaultDacl,
+                   sizeof(TOKEN_DEFAULT_DACL),
+                   sizeof(ULONG));
+      ProbeForRead(TokenSource,
+                   sizeof(TOKEN_SOURCE),
+                   sizeof(ULONG));
+      nTokenPrivileges = TokenPrivileges->PrivilegeCount;
+    }
+    _SEH_HANDLE
+    {
+      Status = _SEH_GetExceptionCode();
+    }
+    _SEH_END;
 
-  ObjectAttributes = &SafeObjectAttributes;
+    if(!NT_SUCCESS(Status))
+    {
+      return Status;
+    }
+  }
+  else
+  {
+    nTokenPrivileges = TokenPrivileges->PrivilegeCount;
+    LocalExpirationTime = *ExpirationTime;
+  }
 
   Status = ZwAllocateLocallyUniqueId(&TokenId);
   if (!NT_SUCCESS(Status))
@@ -1603,10 +2125,10 @@ NtCreateToken(OUT PHANDLE UnsafeTokenHandle,
   if (!NT_SUCCESS(Status))
     return(Status);
 
-  Status = ObCreateObject(ExGetPreviousMode(),
+  Status = ObCreateObject(PreviousMode,
                          SepTokenObjectType,
                          ObjectAttributes,
-                         ExGetPreviousMode(),
+                         PreviousMode,
                          NULL,
                          sizeof(TOKEN),
                          0,
@@ -1618,18 +2140,7 @@ NtCreateToken(OUT PHANDLE UnsafeTokenHandle,
       return(Status);
     }
 
-  Status = ObInsertObject ((PVOID)AccessToken,
-                          NULL,
-                          DesiredAccess,
-                          0,
-                          NULL,
-                          &TokenHandle);
-  if (!NT_SUCCESS(Status))
-    {
-      DPRINT1("ObInsertObject() failed (Status %lx)\n");
-      ObDereferenceObject (AccessToken);
-      return Status;
-    }
+  AccessToken->TokenLock = &SepTokenLock;
 
   RtlCopyLuid(&AccessToken->TokenSource.SourceIdentifier,
              &TokenSource->SourceIdentifier);
@@ -1662,8 +2173,8 @@ NtCreateToken(OUT PHANDLE UnsafeTokenHandle,
   for (i = 0; i < TokenGroups->GroupCount; i++)
     uLength += RtlLengthSid(TokenGroups->Groups[i].Sid);
 
-  AccessToken->UserAndGroups = 
-    (PSID_AND_ATTRIBUTES)ExAllocatePoolWithTag(NonPagedPool,
+  AccessToken->UserAndGroups =
+    (PSID_AND_ATTRIBUTES)ExAllocatePoolWithTag(PagedPool,
                                               uLength,
                                               TAG('T', 'O', 'K', 'u'));
 
@@ -1690,7 +2201,7 @@ NtCreateToken(OUT PHANDLE UnsafeTokenHandle,
   if (NT_SUCCESS(Status))
     {
       Status = SepFindPrimaryGroupAndDefaultOwner(
-       AccessToken, 
+       AccessToken,
        TokenPrimaryGroup->PrimaryGroup,
        TokenOwner->Owner);
     }
@@ -1699,24 +2210,36 @@ NtCreateToken(OUT PHANDLE UnsafeTokenHandle,
     {
       uLength = TokenPrivileges->PrivilegeCount * sizeof(LUID_AND_ATTRIBUTES);
       AccessToken->Privileges =
-       (PLUID_AND_ATTRIBUTES)ExAllocatePoolWithTag(NonPagedPool,
+       (PLUID_AND_ATTRIBUTES)ExAllocatePoolWithTag(PagedPool,
                                                    uLength,
                                                    TAG('T', 'O', 'K', 'p'));
 
-      for (i = 0; i < TokenPrivileges->PrivilegeCount; i++)
-       {
-         Status = MmCopyFromCaller(&AccessToken->Privileges[i],
-                                   &TokenPrivileges->Privileges[i],
-                                   sizeof(LUID_AND_ATTRIBUTES));
-         if (!NT_SUCCESS(Status))
-           break;
-       }
+      if (PreviousMode != KernelMode)
+        {
+          _SEH_TRY
+            {
+              RtlCopyMemory(AccessToken->Privileges,
+                            TokenPrivileges->Privileges,
+                            nTokenPrivileges * sizeof(LUID_AND_ATTRIBUTES));
+            }
+          _SEH_HANDLE
+            {
+              Status = _SEH_GetExceptionCode();
+            }
+          _SEH_END;
+        }
+      else
+        {
+          RtlCopyMemory(AccessToken->Privileges,
+                        TokenPrivileges->Privileges,
+                        nTokenPrivileges * sizeof(LUID_AND_ATTRIBUTES));
+        }
     }
 
   if (NT_SUCCESS(Status))
     {
       AccessToken->DefaultDacl =
-       (PACL) ExAllocatePoolWithTag(NonPagedPool,
+       (PACL) ExAllocatePoolWithTag(PagedPool,
                                     TokenDefaultDacl->DefaultDacl->AclSize,
                                     TAG('T', 'O', 'K', 'd'));
       memcpy(AccessToken->DefaultDacl,
@@ -1724,22 +2247,33 @@ NtCreateToken(OUT PHANDLE UnsafeTokenHandle,
             TokenDefaultDacl->DefaultDacl->AclSize);
     }
 
-  ObDereferenceObject(AccessToken);
-
-  if (NT_SUCCESS(Status))
+  Status = ObInsertObject ((PVOID)AccessToken,
+                          NULL,
+                          DesiredAccess,
+                          0,
+                          NULL,
+                          &hToken);
+  if (!NT_SUCCESS(Status))
     {
-      Status = MmCopyToCaller(UnsafeTokenHandle,
-                             &TokenHandle,
-                             sizeof(HANDLE));
+      DPRINT1("ObInsertObject() failed (Status %lx)\n", Status);
     }
 
-  if (!NT_SUCCESS(Status))
+  ObDereferenceObject(AccessToken);
+
+  if (NT_SUCCESS(Status))
     {
-      ZwClose(TokenHandle);
-      return(Status);
+      _SEH_TRY
+      {
+        *TokenHandle = hToken;
+      }
+      _SEH_HANDLE
+      {
+        Status = _SEH_GetExceptionCode();
+      }
+      _SEH_END;
     }
 
-  return(STATUS_SUCCESS);
+  return Status;
 }
 
 
@@ -1750,6 +2284,8 @@ NTSTATUS STDCALL
 SeQueryAuthenticationIdToken(IN PACCESS_TOKEN Token,
                             OUT PLUID LogonId)
 {
+  PAGED_CODE();
+
   *LogonId = ((PTOKEN)Token)->AuthenticationId;
 
   return STATUS_SUCCESS;
@@ -1763,6 +2299,8 @@ SECURITY_IMPERSONATION_LEVEL
 STDCALL
 SeTokenImpersonationLevel(IN PACCESS_TOKEN Token)
 {
+  PAGED_CODE();
+
   return ((PTOKEN)Token)->ImpersonationLevel;
 }
 
@@ -1773,6 +2311,8 @@ SeTokenImpersonationLevel(IN PACCESS_TOKEN Token)
 TOKEN_TYPE STDCALL
 SeTokenType(IN PACCESS_TOKEN Token)
 {
+  PAGED_CODE();
+
   return ((PTOKEN)Token)->TokenType;
 }
 
@@ -1829,6 +2369,7 @@ NtOpenThreadTokenEx(IN HANDLE ThreadHandle,
                     OUT PHANDLE TokenHandle)
 {
   PETHREAD Thread;
+  HANDLE hToken;
   PTOKEN Token, NewToken, PrimaryToken;
   BOOLEAN CopyOnOpen, EffectiveOnly;
   SECURITY_IMPERSONATION_LEVEL ImpersonationLevel;
@@ -1836,7 +2377,30 @@ NtOpenThreadTokenEx(IN HANDLE ThreadHandle,
   OBJECT_ATTRIBUTES ObjectAttributes;
   SECURITY_DESCRIPTOR SecurityDescriptor;
   PACL Dacl = NULL;
-  NTSTATUS Status;
+  KPROCESSOR_MODE PreviousMode;
+  NTSTATUS Status = STATUS_SUCCESS;
+
+  PAGED_CODE();
+
+  PreviousMode = ExGetPreviousMode();
+
+  if(PreviousMode != KernelMode)
+  {
+    _SEH_TRY
+    {
+      ProbeForWriteHandle(TokenHandle);
+    }
+    _SEH_HANDLE
+    {
+      Status = _SEH_GetExceptionCode();
+    }
+    _SEH_END;
+
+    if(!NT_SUCCESS(Status))
+    {
+      return Status;
+    }
+  }
 
   /*
    * At first open the thread token for information access and verify
@@ -1844,7 +2408,7 @@ NtOpenThreadTokenEx(IN HANDLE ThreadHandle,
    */
 
   Status = ObReferenceObjectByHandle(ThreadHandle, THREAD_QUERY_INFORMATION,
-                                     PsThreadType, UserMode, (PVOID*)&Thread,
+                                     PsThreadType, PreviousMode, (PVOID*)&Thread,
                                      NULL);
   if (!NT_SUCCESS(Status))
     {
@@ -1879,7 +2443,7 @@ NtOpenThreadTokenEx(IN HANDLE ThreadHandle,
   if (CopyOnOpen)
     {
       Status = ObReferenceObjectByHandle(ThreadHandle, THREAD_ALL_ACCESS,
-                                         PsThreadType, UserMode,
+                                         PsThreadType, PreviousMode,
                                          (PVOID*)&Thread, NULL);
       if (!NT_SUCCESS(Status))
         {
@@ -1890,7 +2454,7 @@ NtOpenThreadTokenEx(IN HANDLE ThreadHandle,
             }
           return Status;
         }
-   
+
       PrimaryToken = PsReferencePrimaryToken(Thread->ThreadsProcess);
       Status = SepCreateImpersonationTokenDacl(Token, PrimaryToken, &Dacl);
       ObfDereferenceObject(PrimaryToken);
@@ -1904,7 +2468,7 @@ NtOpenThreadTokenEx(IN HANDLE ThreadHandle,
             }
           return Status;
         }
-      
+
       RtlCreateSecurityDescriptor(&SecurityDescriptor,
                                   SECURITY_DESCRIPTOR_REVISION);
       RtlSetDaclSecurityDescriptor(&SecurityDescriptor, TRUE, Dacl,
@@ -1928,7 +2492,7 @@ NtOpenThreadTokenEx(IN HANDLE ThreadHandle,
         }
 
       Status = ObInsertObject(NewToken, NULL, DesiredAccess, 0, NULL,
-                              TokenHandle);
+                              &hToken);
 
       ObfDereferenceObject(NewToken);
     }
@@ -1936,7 +2500,7 @@ NtOpenThreadTokenEx(IN HANDLE ThreadHandle,
     {
       Status = ObOpenObjectByPointer(Token, HandleAttributes,
                                      NULL, DesiredAccess, SepTokenObjectType,
-                                     ExGetPreviousMode(), TokenHandle);
+                                     PreviousMode, &hToken);
     }
 
   ObfDereferenceObject(Token);
@@ -1946,6 +2510,19 @@ NtOpenThreadTokenEx(IN HANDLE ThreadHandle,
       PsRestoreImpersonation(PsGetCurrentThread(), &ImpersonationState);
     }
 
+  if(NT_SUCCESS(Status))
+  {
+    _SEH_TRY
+    {
+      *TokenHandle = hToken;
+    }
+    _SEH_HANDLE
+    {
+      Status = _SEH_GetExceptionCode();
+    }
+    _SEH_END;
+  }
+
   return Status;
 }