[reactos] 01/01: [RTL][NTDLL_APITEST] Implement RtlRemovePrivileges (#4614)

[Ratin Gao]

https://git.reactos.org/?p=reactos.git;a=commitdiff;h=badd97043fbce3bbfe60a…

commit badd97043fbce3bbfe60aa0470749a0cf4678b34
Author:     Ratin Gao <ratin(a)knsoft.org>
AuthorDate: Wed Oct 5 20:31:39 2022 +0800
Commit:     GitHub <noreply(a)github.com>
CommitDate: Wed Oct 5 14:31:39 2022 +0200

    [RTL][NTDLL_APITEST] Implement RtlRemovePrivileges (#4614)
    
    Vista+ API, compile-time guarded.
    Add tests for it.
---
 dll/ntdll/def/ntdll.spec                           |   2 +-
 modules/rostests/apitests/ntdll/CMakeLists.txt     |   1 +
 .../rostests/apitests/ntdll/RtlRemovePrivileges.c  | 111 +++++++++++++++++++++
 modules/rostests/apitests/ntdll/testlist.c         |   2 +
 sdk/include/ndk/rtlfuncs.h                         |  14 +++
 sdk/lib/rtl/priv.c                                 | 102 +++++++++++++++++++
 6 files changed, 231 insertions(+), 1 deletion(-)

diff --git a/dll/ntdll/def/ntdll.spec b/dll/ntdll/def/ntdll.spec
index 80d27637413..3cd201aa3d4 100644
--- a/dll/ntdll/def/ntdll.spec
+++ b/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
diff --git a/modules/rostests/apitests/ntdll/CMakeLists.txt
b/modules/rostests/apitests/ntdll/CMakeLists.txt
index e1076a3a584..c2ec474a5a0 100644
--- a/modules/rostests/apitests/ntdll/CMakeLists.txt
+++ b/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
diff --git a/modules/rostests/apitests/ntdll/RtlRemovePrivileges.c
b/modules/rostests/apitests/ntdll/RtlRemovePrivileges.c
new file mode 100644
index 00000000000..861dba1a8a8
--- /dev/null
+++ b/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(a)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) */
+}
diff --git a/modules/rostests/apitests/ntdll/testlist.c
b/modules/rostests/apitests/ntdll/testlist.c
index f5b265f0305..29b3ae4ba21 100644
--- a/modules/rostests/apitests/ntdll/testlist.c
+++ b/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 },
diff --git a/sdk/include/ndk/rtlfuncs.h b/sdk/include/ndk/rtlfuncs.h
index c6bdcfcdd19..05dc37686ba 100644
--- a/sdk/include/ndk/rtlfuncs.h
+++ b/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
diff --git a/sdk/lib/rtl/priv.c b/sdk/lib/rtl/priv.c
index 164bd80f75f..0070140103f 100644
--- a/sdk/lib/rtl/priv.c
+++ b/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) */