[RTL][NTDLL_APITEST] Implement RtlRemovePrivileges (#4614)

Vista+ API, compile-time guarded.
Add tests for it.

dll/ntdll/def/ntdll.spec

@@ -1111,7 +1111,7 @@
 @ stdcall -stub -version=0x600+ RtlReleaseSRWLockExclusive(ptr)
 @ stdcall -stub -version=0x600+ RtlReleaseSRWLockShared(ptr)
 @ stdcall RtlRemoteCall(ptr ptr ptr long ptr long long)
-@ stub -version=0x600+ RtlRemovePrivileges
+@ stdcall -version=0x600+ RtlRemovePrivileges(ptr ptr long)
 @ stdcall RtlRemoveVectoredContinueHandler(ptr)
 @ stdcall RtlRemoveVectoredExceptionHandler(ptr)
 @ stub -version=0x600+ RtlReportException

modules/rostests/apitests/ntdll/CMakeLists.txt

@@ -85,6 +85,7 @@ list(APPEND SOURCE
     RtlpEnsureBufferSize.c
     RtlQueryTimeZoneInfo.c
     RtlReAllocateHeap.c
+    RtlRemovePrivileges.c
     RtlUnicodeStringToAnsiString.c
     RtlUnicodeStringToCountedOemString.c
     RtlUnicodeToOemN.c

modules/rostests/apitests/ntdll/RtlRemovePrivileges.c

@@ -0,0 +1,111 @@
+/*
+ * PROJECT:         ReactOS api tests
+ * LICENSE:         See COPYING in the top level directory
+ * PURPOSE:         Test for RtlRemovePrivileges
+ * PROGRAMMER:      Ratin Gao <ratin@knsoft.org>
+ */
+
+#include "precomp.h"
+
+START_TEST(RtlRemovePrivileges)
+{
+#if (NTDDI_VERSION >= NTDDI_VISTA)
+    NTSTATUS Status;
+    HANDLE TokenHandle, TestTokenHandle;
+    ULONG ReturnLength;
+    UCHAR Buffer
+        [sizeof(TOKEN_PRIVILEGES) +
+         sizeof(LUID_AND_ATTRIBUTES) * (SE_MAX_WELL_KNOWN_PRIVILEGE - SE_MIN_WELL_KNOWN_PRIVILEGE)];
+    PTOKEN_PRIVILEGES Privileges;
+    ULONG PrivilegesToKeep[2];
+
+    /* Duplicate current process token to run this test */
+    Status = NtOpenProcessToken(GetCurrentProcess(), TOKEN_DUPLICATE, &TokenHandle);
+    if (!NT_SUCCESS(Status))
+    {
+        ok(0, "Failed to open current process token with TOKEN_DUPLICATE access (Status code %lx)!\n", Status);
+        return;
+    }
+
+    Status = NtDuplicateToken(TokenHandle, TOKEN_ALL_ACCESS, NULL, FALSE, TokenPrimary, &TestTokenHandle);
+    NtClose(TokenHandle);
+    if (!NT_SUCCESS(Status))
+    {
+        ok(0, "Failed to duplicate current process token (Status code %lx)!\n", Status);
+        return;
+    }
+
+    /* Retrieve token privileges, we need at least 3 privileges to run following tests */
+    Status = NtQueryInformationToken(TestTokenHandle, TokenPrivileges, Buffer, sizeof(Buffer), &ReturnLength);
+    if (!NT_SUCCESS(Status))
+    {
+        NtClose(TestTokenHandle);
+        ok(0, "Failed to retrieve token privileges (Status code %lx)!\n", Status);
+        return;
+    }
+    Privileges = (PTOKEN_PRIVILEGES)Buffer;
+    if (Privileges->PrivilegeCount < 3)
+    {
+        NtClose(TestTokenHandle);
+        ok(0, "No enough privileges to run the test (Number of privilege: %lu)!\n", Privileges->PrivilegeCount);
+        return;
+    }
+
+    /* Remove all privileges except 2nd and 3rd privileges, this should succeed */
+    PrivilegesToKeep[0] = Privileges->Privileges[1].Luid.LowPart;
+    PrivilegesToKeep[1] = Privileges->Privileges[2].Luid.LowPart;
+    Status = RtlRemovePrivileges(TestTokenHandle, PrivilegesToKeep, ARRAYSIZE(PrivilegesToKeep));
+
+    /* Do not use NT_SUCCESS, RtlRemovePrivileges may returns STATUS_NOT_ALL_ASSIGNED */
+    if (Status != STATUS_SUCCESS)
+    {
+        NtClose(TestTokenHandle);
+        ok_ntstatus(Status, STATUS_SUCCESS);
+        return;
+    }
+
+    /* Now, only two privileges we kept should be present */
+    Status = NtQueryInformationToken(TestTokenHandle, TokenPrivileges, Buffer, sizeof(Buffer), &ReturnLength);
+    if (!NT_SUCCESS(Status))
+    {
+        NtClose(TestTokenHandle);
+        ok(0, "Failed to retrieve token privileges (Status code %lx)!\n", Status);
+        return;
+    }
+    ok(Privileges->PrivilegeCount == ARRAYSIZE(PrivilegesToKeep),
+       "Number of privileges after RtlRemovePrivileges is %lu, expected %u\n", Privileges->PrivilegeCount,
+       ARRAYSIZE(PrivilegesToKeep));
+    ok(PrivilegesToKeep[0] + PrivilegesToKeep[1] ==
+           Privileges->Privileges[0].Luid.LowPart + Privileges->Privileges[1].Luid.LowPart,
+       "Incorrect privileges kept by RtlRemovePrivileges: %lu and %lu, expected %lu and %lu",
+       Privileges->Privileges[0].Luid.LowPart, Privileges->Privileges[1].Luid.LowPart, PrivilegesToKeep[0],
+       PrivilegesToKeep[1]);
+
+    /* Remove all privileges, this should succeed */
+    Status = RtlRemovePrivileges(TestTokenHandle, NULL, 0);
+
+    /* Do not use NT_SUCCESS, RtlRemovePrivileges may returns STATUS_NOT_ALL_ASSIGNED */
+    if (Status != STATUS_SUCCESS)
+    {
+        NtClose(TestTokenHandle);
+        ok_ntstatus(Status, STATUS_SUCCESS);
+        return;
+    }
+
+    /* Now, no privilege should be present */
+    Status = NtQueryInformationToken(TestTokenHandle, TokenPrivileges, Buffer, sizeof(Buffer), &ReturnLength);
+    if (!NT_SUCCESS(Status))
+    {
+        NtClose(TestTokenHandle);
+        ok(0, "Failed to retrieve token privileges (Status code %lx)!\n", Status);
+        return;
+    }
+    ok(Privileges->PrivilegeCount == 0, "There are %lu privileges still exist after RtlRemovePrivileges\n",
+       Privileges->PrivilegeCount);
+
+    NtClose(TestTokenHandle);
+    return;
+#else
+    skip("RtlRemovePrivileges available on NT6.0+ (NTDDI_VERSION >= NTDDI_VISTA)");
+#endif /* (NTDDI_VERSION >= NTDDI_VISTA) */
+}

modules/rostests/apitests/ntdll/testlist.c

@@ -81,6 +81,7 @@ extern void func_RtlpApplyLengthFunction(void);
 extern void func_RtlpEnsureBufferSize(void);
 extern void func_RtlQueryTimeZoneInformation(void);
 extern void func_RtlReAllocateHeap(void);
+extern void func_RtlRemovePrivileges(void);
 extern void func_RtlUnicodeStringToAnsiString(void);
 extern void func_RtlUnicodeStringToCountedOemString(void);
 extern void func_RtlUnicodeToOemN(void);
@@ -172,6 +173,7 @@ const struct test winetest_testlist[] =
     { "RtlpEnsureBufferSize",           func_RtlpEnsureBufferSize },
     { "RtlQueryTimeZoneInformation",    func_RtlQueryTimeZoneInformation },
     { "RtlReAllocateHeap",              func_RtlReAllocateHeap },
+    { "RtlRemovePrivileges",            func_RtlRemovePrivileges },
     { "RtlUnicodeStringToAnsiSize",     func_RtlxUnicodeStringToAnsiSize }, /* For some reason, starting test name with Rtlx hides it */
     { "RtlUnicodeStringToAnsiString",   func_RtlUnicodeStringToAnsiString },
     { "RtlUnicodeStringToCountedOemString", func_RtlUnicodeStringToCountedOemString },

sdk/include/ndk/rtlfuncs.h

@@ -1568,6 +1568,20 @@ RtlReleasePrivilege(
     _In_ PVOID ReturnedState
 );
 
+#if (NTDDI_VERSION >= NTDDI_VISTA)
+
+NTSYSAPI
+NTSTATUS
+NTAPI
+RtlRemovePrivileges(
+    _In_ HANDLE TokenHandle,
+    _In_reads_opt_(PrivilegeCount) _When_(PrivilegeCount != 0, _Notnull_)
+         PULONG PrivilegesToKeep,
+    _In_ ULONG PrivilegeCount
+);
+
+#endif /* (NTDDI_VERSION >= NTDDI_VISTA) */
+
 _IRQL_requires_max_(APC_LEVEL)
 NTSYSAPI
 NTSTATUS

sdk/lib/rtl/priv.c

@@ -486,3 +486,105 @@ RtlAdjustPrivilege(IN ULONG Privilege,
 
     return STATUS_SUCCESS;
 }
+
+#if (NTDDI_VERSION >= NTDDI_VISTA)
+
+/**
+ * @brief
+ * Removes all privileges in the specified access token.
+ *
+ * @param[in] TokenHandle
+ * A handle to the access token that contains the privileges to be removed.
+ *
+ * @param[in] PrivilegesToKeep
+ * A pointer to an array of privilege values (defined as SE_XXX_PRIVILEGE) that specify
+ * the privileges to keep in the token.
+ *
+ * @param[in] PrivilegeCount
+ * Specifies the number of entries in the PrivilegesToKeep array.
+ *
+ * @return
+ * Returns STATUS_SUCCESS if privileges removed successfully.
+ * STATUS_INVALID_PARAMETER is returned if input privilege value greater than
+ * SE_MAX_WELL_KNOWN_PRIVILEGE. STATUS_NOT_ALL_ASSIGNED is returned if The token does
+ * not have one or more of the privileges specified in the PrivilegesToKeep parameter,
+ * and no privileges were removed. A failure NTSTATUS code is returned otherwise.
+ */
+NTSTATUS
+NTAPI
+RtlRemovePrivileges(
+    _In_ HANDLE TokenHandle,
+    _In_reads_opt_(PrivilegeCount) _When_(PrivilegeCount != 0, _Notnull_)
+         PULONG PrivilegesToKeep,
+    _In_ ULONG PrivilegeCount)
+{
+    NTSTATUS Status;
+    UINT64 PrivilegesToKeepBitmap;
+    ULONG i, ReturnLength;
+    UCHAR Buffer[sizeof(TOKEN_PRIVILEGES) +
+                 sizeof(LUID_AND_ATTRIBUTES) * (SE_MAX_WELL_KNOWN_PRIVILEGE - SE_MIN_WELL_KNOWN_PRIVILEGE)];
+    PTOKEN_PRIVILEGES Privileges;
+
+    C_ASSERT(SE_MAX_WELL_KNOWN_PRIVILEGE < 64);
+
+    DPRINT("RtlRemovePrivileges(%p, %p, %u)\n", TokenHandle, PrivilegesToKeep, PrivilegeCount);
+
+    /* Save privileges that should be keep */
+    PrivilegesToKeepBitmap = 0;
+    if (PrivilegeCount)
+    {
+        for (i = 0; i < PrivilegeCount; i++)
+        {
+            if (PrivilegesToKeep[i] > SE_MAX_WELL_KNOWN_PRIVILEGE)
+            {
+                return STATUS_INVALID_PARAMETER;
+            }
+            PrivilegesToKeepBitmap |= (1ULL << PrivilegesToKeep[i]);
+        }
+    }
+
+    /* Get token privileges information */
+    Status = ZwQueryInformationToken(TokenHandle,
+                                     TokenPrivileges,
+                                     Buffer,
+                                     sizeof(Buffer),
+                                     &ReturnLength);
+    if (!NT_SUCCESS(Status))
+    {
+        return Status;
+    }
+
+    /* Remove all privileges that we don't need to keep */
+    Privileges = (PTOKEN_PRIVILEGES)Buffer;
+    for (i = 0; i < Privileges->PrivilegeCount; i++)
+    {
+        LARGE_INTEGER Privilege = *(LARGE_INTEGER*)&Privileges->Privileges[i].Luid;
+        ASSERT(Privilege.QuadPart <= SE_MAX_WELL_KNOWN_PRIVILEGE);
+        if (PrivilegesToKeepBitmap & (1ULL << Privilege.QuadPart))
+        {
+            PrivilegesToKeepBitmap &= ~(1ULL << Privilege.QuadPart);
+        }
+        else
+        {
+            Privileges->Privileges[i].Attributes = SE_PRIVILEGE_REMOVED;
+        }
+    }
+
+    if (PrivilegesToKeepBitmap)
+    {
+        Status = STATUS_NOT_ALL_ASSIGNED;
+    }
+    else
+    {
+        Status = ZwAdjustPrivilegesToken(TokenHandle,
+                                         FALSE,
+                                         (PTOKEN_PRIVILEGES)Buffer,
+                                         sizeof(Buffer),
+                                         NULL,
+                                         NULL);
+    }
+
+    return Status;
+}
+
+#endif /* (NTDDI_VERSION >= NTDDI_VISTA) */