1 /*
2  * PROJECT:         ReactOS API tests
3  * LICENSE:         GPL-2.0-or-later (https://spdx.org/licenses/GPL-2.0-or-later)
4  * PURPOSE:         Tests for the NtSetInformationToken API
5  * COPYRIGHT:       Copyright 2022 George Bișoc <george.bisoc@reactos.org>
6  */
7 
8 #include "precomp.h"
9 
10 static
11 HANDLE
12 OpenCurrentToken(VOID)
13 {
14     BOOL Success;
15     HANDLE Token;
16 
17     Success = OpenProcessToken(GetCurrentProcess(),
18                                TOKEN_READ | TOKEN_ADJUST_DEFAULT | TOKEN_ADJUST_SESSIONID,
19                                &Token);
20     if (!Success)
21     {
22         ok(0, "OpenProcessToken() has failed to get the process' token (error code: %lu)!\n", GetLastError());
23         return NULL;
24     }
25 
26     return Token;
27 }
28 
29 static
30 PTOKEN_DEFAULT_DACL
31 QueryOriginalDefaultDacl(
32     _In_ HANDLE Token,
33     _Out_ PULONG DaclLength)
34 {
35     NTSTATUS Status;
36     PTOKEN_DEFAULT_DACL Dacl;
37     ULONG BufferLength;
38 
39     *DaclLength = 0;
40 
41     Status = NtQueryInformationToken(Token,
42                                      TokenDefaultDacl,
43                                      NULL,
44                                      0,
45                                      &BufferLength);
46     if (!NT_SUCCESS(Status) && (Status != STATUS_BUFFER_TOO_SMALL))
47     {
48         ok(0, "Failed to query buffer length, STATUS_BUFFER_TOO_SMALL has to be expected (Status code %lx)!\n", Status);
49         return NULL;
50     }
51 
52     Dacl = RtlAllocateHeap(RtlGetProcessHeap(), 0, BufferLength);
53     if (!Dacl)
54     {
55         ok(0, "Failed to allocate from heap for token default DACL (required buffer length %lu)!\n", BufferLength);
56         return NULL;
57     }
58 
59     Status = NtQueryInformationToken(Token,
60                                      TokenDefaultDacl,
61                                      Dacl,
62                                      BufferLength,
63                                      &BufferLength);
64     if (!NT_SUCCESS(Status))
65     {
66         ok(0, "Failed to query default DACL (Status code %lx)!\n", Status);
67         RtlFreeHeap(RtlGetProcessHeap(), 0, Dacl);
68         return NULL;
69     }
70 
71     *DaclLength = BufferLength;
72     return Dacl;
73 }
74 
75 static
76 PACL
77 CreateNewDefaultDacl(
78     _Out_ PULONG DaclLength)
79 {
80     NTSTATUS Status;
81     PACL Dacl;
82     ULONG Length;
83     PSID LocalSystemSid;
84     static SID_IDENTIFIER_AUTHORITY NtAuthority = {SECURITY_NT_AUTHORITY};
85 
86     *DaclLength = 0;
87 
88     Status = RtlAllocateAndInitializeSid(&NtAuthority,
89                                          1,
90                                          SECURITY_LOCAL_SYSTEM_RID,
91                                          0, 0, 0, 0, 0, 0, 0,
92                                          &LocalSystemSid);
93     if (!NT_SUCCESS(Status))
94     {
95         ok(0, "Failed to allocate Local System SID (Status code %lx)!\n", Status);
96         return NULL;
97     }
98 
99     Length = sizeof(ACL) +
100                  sizeof(ACCESS_ALLOWED_ACE) + RtlLengthSid(LocalSystemSid);
101 
102     Dacl = RtlAllocateHeap(RtlGetProcessHeap(),
103                            HEAP_ZERO_MEMORY,
104                            Length);
105     if (!Dacl)
106     {
107         ok(0, "Failed to allocate from heap for DACL!\n");
108         RtlFreeHeap(RtlGetProcessHeap(), 0, LocalSystemSid);
109         return NULL;
110     }
111 
112     Status = RtlCreateAcl(Dacl,
113                           Length,
114                           ACL_REVISION);
115     if (!NT_SUCCESS(Status))
116     {
117         ok(0, "Failed to create ACL (Status code %lx)!\n", Status);
118         RtlFreeHeap(RtlGetProcessHeap(), 0, LocalSystemSid);
119         RtlFreeHeap(RtlGetProcessHeap(), 0, Dacl);
120         return NULL;
121     }
122 
123     Status = RtlAddAccessAllowedAce(Dacl,
124                                     ACL_REVISION,
125                                     GENERIC_ALL,
126                                     LocalSystemSid);
127     if (!NT_SUCCESS(Status))
128     {
129         ok(0, "Failed to add access allowed ACE (Status code %lx)!\n", Status);
130         RtlFreeHeap(RtlGetProcessHeap(), 0, LocalSystemSid);
131         RtlFreeHeap(RtlGetProcessHeap(), 0, Dacl);
132         return NULL;
133     }
134 
135     *DaclLength = Length;
136     RtlFreeHeap(RtlGetProcessHeap(), 0, LocalSystemSid);
137     return Dacl;
138 }
139 
140 static
141 VOID
142 SetTokenDefaultDaclTests(
143     _In_ HANDLE Token)
144 {
145     NTSTATUS Status;
146     PACL NewDacl;
147     TOKEN_DEFAULT_DACL NewDefaultDacl;
148     PTOKEN_DEFAULT_DACL DefaultDacl;
149     ULONG OriginalDaclLength, NewDaclLength;
150 
151     /*
152      * Query the original DACL of the token first,
153      * we don't want to leave the token tampered
154      * later on.
155      */
156     DefaultDacl = QueryOriginalDefaultDacl(Token, &OriginalDaclLength);
157     if (!DefaultDacl)
158     {
159         ok(0, "Failed to query token's default DACL!\n");
160         return;
161     }
162 
163     /* Allocate new DACL */
164     NewDacl = CreateNewDefaultDacl(&NewDaclLength);
165     if (!DefaultDacl)
166     {
167         ok(0, "Failed to allocate buffer for new DACL!\n");
168         RtlFreeHeap(RtlGetProcessHeap(), 0, DefaultDacl);
169         return;
170     }
171 
172     NewDefaultDacl.DefaultDacl = NewDacl;
173 
174     /*
175      * Set a new DACL for the token.
176      */
177     Status = NtSetInformationToken(Token,
178                                    TokenDefaultDacl,
179                                    &NewDefaultDacl,
180                                    NewDaclLength);
181     ok_ntstatus(Status, STATUS_SUCCESS);
182 
183     /* Now set the original DACL */
184     Status = NtSetInformationToken(Token,
185                                    TokenDefaultDacl,
186                                    DefaultDacl,
187                                    OriginalDaclLength);
188     ok_ntstatus(Status, STATUS_SUCCESS);
189 
190     RtlFreeHeap(RtlGetProcessHeap(), 0, DefaultDacl);
191     RtlFreeHeap(RtlGetProcessHeap(), 0, NewDacl);
192 }
193 
194 static
195 VOID
196 SetTokenSessionIdTests(
197     _In_ HANDLE Token)
198 {
199     NTSTATUS Status;
200     ULONG SessionId = 1;
201 
202     /*
203      * We're not allowed to set a session ID
204      * because we don't have the TCB privilege.
205      */
206     Status = NtSetInformationToken(Token,
207                                    TokenSessionId,
208                                    &SessionId,
209                                    sizeof(ULONG));
210     ok_ntstatus(Status, STATUS_PRIVILEGE_NOT_HELD);
211 }
212 
213 START_TEST(NtSetInformationToken)
214 {
215     NTSTATUS Status;
216     ULONG DummyReturnLength = 0;
217     HANDLE Token;
218 
219     /* Everything else is NULL */
220     Status = NtSetInformationToken(NULL,
221                                    TokenOwner,
222                                    NULL,
223                                    0);
224     ok_ntstatus(Status, STATUS_INVALID_HANDLE);
225 
226     /* We don't give a token */
227     Status = NtSetInformationToken(NULL,
228                                    TokenOwner,
229                                    NULL,
230                                    DummyReturnLength);
231     ok_ntstatus(Status, STATUS_INVALID_HANDLE);
232 
233     Token = OpenCurrentToken();
234 
235     /* We give a bogus token class */
236     Status = NtSetInformationToken(Token,
237                                    0xa0a,
238                                    NULL,
239                                    DummyReturnLength);
240     ok_ntstatus(Status, STATUS_INVALID_INFO_CLASS);
241 
242     /* Now perform tests for each class */
243     SetTokenDefaultDaclTests(Token);
244     SetTokenSessionIdTests(Token);
245 
246     CloseHandle(Token);
247 }
248