1 /*
2  * PROJECT:     Global Flags utility
3  * LICENSE:     GPL-2.0 (https://spdx.org/licenses/GPL-2.0)
4  * PURPOSE:     Global Flags utility page heap options
5  * COPYRIGHT:   Copyright 2017 Pierre Schweitzer (pierre@reactos.org)
6  */
7 
8 #include "gflags.h"
9 
10 static BOOL Set = FALSE;
11 static BOOL Unset = FALSE;
12 static BOOL Full = FALSE;
13 static PWSTR Image = NULL;
14 
15 
16 static VOID ModifyStatus(VOID)
17 {
18     LONG Ret;
19     DWORD GlobalFlags;
20     HKEY IFEOKey;
21     WCHAR Buffer[11];
22 
23     if (!OpenImageFileExecOptions(KEY_WRITE | KEY_READ, Image, &IFEOKey))
24     {
25         return;
26     }
27 
28     GlobalFlags = ReadSZFlagsFromRegistry(IFEOKey, L"GlobalFlag");
29     if (Set)
30     {
31         GlobalFlags |= FLG_HEAP_PAGE_ALLOCS;
32     }
33     else
34     {
35         GlobalFlags &= ~FLG_HEAP_PAGE_ALLOCS;
36     }
37 
38     if (GlobalFlags != 0)
39     {
40         wsprintf(Buffer, L"0x%08x", GlobalFlags);
41         Ret = RegSetValueExW(IFEOKey, L"GlobalFlag", 0, REG_SZ, (BYTE*)Buffer, (wcslen(Buffer) + 1) * sizeof(WCHAR));
42         if (Ret != ERROR_SUCCESS)
43         {
44             wprintf(L"MS: RegSetValueEx failed (%d)\n", Ret);
45         }
46     }
47     else
48     {
49         Ret = RegDeleteValueW(IFEOKey, L"GlobalFlag");
50         if (Ret != ERROR_SUCCESS)
51         {
52             wprintf(L"MS: RegDeleteValue failed (%d)\n", Ret);
53         }
54     }
55 
56     if (Unset)
57     {
58         Ret = RegDeleteValueW(IFEOKey, L"PageHeapFlags");
59         if (Ret != ERROR_SUCCESS)
60         {
61             wprintf(L"MS: RegDeleteValue failed (%d)\n", Ret);
62         }
63     }
64     else
65     {
66         DWORD PageHeapFlags;
67 
68         PageHeapFlags = ReadSZFlagsFromRegistry(IFEOKey, L"PageHeapFlags");
69         PageHeapFlags &= ~3;
70 
71         if (Full)
72         {
73             PageHeapFlags |= 1;
74         }
75         PageHeapFlags |= 2;
76 
77         wsprintf(Buffer, L"0x%x", PageHeapFlags);
78         Ret = RegSetValueExW(IFEOKey, L"PageHeapFlags", 0, REG_SZ, (BYTE*)Buffer, (wcslen(Buffer) + 1) * sizeof(WCHAR));
79         if (Ret != ERROR_SUCCESS)
80         {
81             wprintf(L"MS: RegSetValueEx failed (%d)\n", Ret);
82         }
83     }
84 
85     if (Set)
86     {
87         DWORD Type, VerifierFlags, Len;
88 
89         VerifierFlags = 0;
90         Len = VerifierFlags;
91         if (RegQueryValueExW(IFEOKey, L"VerifierFlags", NULL, &Type, (BYTE *)&VerifierFlags, &Len) == ERROR_SUCCESS &&
92             Type == REG_DWORD && Len == sizeof(DWORD))
93         {
94             VerifierFlags &= ~0x8001;   /* RTL_VRF_FLG_FAST_FILL_HEAP | RTL_VRF_FLG_FULL_PAGE_HEAP */
95         }
96         else
97         {
98             VerifierFlags = 0;
99         }
100 
101         if (Full)
102         {
103             VerifierFlags |= 1; /* RTL_VRF_FLG_FULL_PAGE_HEAP */
104         }
105         else
106         {
107             VerifierFlags |= 0x8000;    /* RTL_VRF_FLG_FAST_FILL_HEAP */
108         }
109 
110         Ret = RegSetValueExW(IFEOKey, L"VerifierFlags", 0, REG_DWORD, (const BYTE *)&VerifierFlags, sizeof(DWORD));
111         if (Ret != ERROR_SUCCESS)
112         {
113             wprintf(L"MS: RegSetValueEx failed (%d)\n", Ret);
114         }
115     }
116 
117     wprintf(L"path: %s\n", ImageExecOptionsString);
118     wprintf(L"\t%s: page heap %s\n", Image, (Set ? L"enabled" : L"disabled"));
119 
120     RegCloseKey(IFEOKey);
121 }
122 
123 static BOOL DisplayImageInfo(HKEY HandleKey, PWSTR SubKey, PBOOL Header)
124 {
125     LONG Ret;
126     BOOL Handled;
127     DWORD GlobalFlags;
128     HKEY HandleSubKey;
129 
130     Ret = RegOpenKeyExW(HandleKey, SubKey, 0, KEY_READ, &HandleSubKey);
131     if (Ret != ERROR_SUCCESS)
132     {
133         wprintf(L"DII: RegOpenKeyEx failed (%d)\n", Ret);
134         return FALSE;
135     }
136 
137     Handled = FALSE;
138     GlobalFlags = ReadSZFlagsFromRegistry(HandleSubKey, L"GlobalFlag");
139     if (GlobalFlags & FLG_HEAP_PAGE_ALLOCS)
140     {
141         DWORD PageHeapFlags;
142 
143         if (Image == NULL)
144         {
145             if (!*Header)
146             {
147                 wprintf(L"path: %s\n", ImageExecOptionsString);
148                 *Header = TRUE;
149             }
150             wprintf(L"\t%s: page heap enabled with flags (", SubKey);
151         }
152         else
153         {
154             wprintf(L"Page heap is enabled for %s with flags (", SubKey);
155         }
156 
157         PageHeapFlags = ReadSZFlagsFromRegistry(HandleSubKey, L"PageHeapFlags");
158         if (PageHeapFlags & 0x1)
159         {
160             wprintf(L"full ");
161         }
162 
163         if (PageHeapFlags & 0x2)
164         {
165             wprintf(L"traces");
166         }
167 
168         wprintf(L")\n");
169 
170         Handled = TRUE;
171     }
172 
173     RegCloseKey(HandleSubKey);
174 
175     return Handled;
176 }
177 
178 static VOID DisplayStatus(VOID)
179 {
180     LONG Ret;
181     HKEY HandleKey;
182     DWORD Index, MaxLen, Handled;
183     TCHAR * SubKey;
184     BOOL Header;
185 
186     if (!OpenImageFileExecOptions(KEY_READ, NULL, &HandleKey))
187     {
188         return;
189     }
190 
191     Ret = RegQueryInfoKeyW(HandleKey, NULL, NULL, NULL, NULL, &MaxLen, NULL, NULL, NULL, NULL, NULL, NULL);
192     if (Ret != ERROR_SUCCESS)
193     {
194         wprintf(L"DS: RegQueryInfoKey failed (%d)\n", Ret);
195         RegCloseKey(HandleKey);
196         return;
197     }
198 
199     ++MaxLen; // NULL-char
200     SubKey = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, MaxLen * sizeof(WCHAR));
201     if (SubKey == NULL)
202     {
203         wprintf(L"DS: HeapAlloc failed\n");
204         RegCloseKey(HandleKey);
205         return;
206     }
207 
208     Index = 0;
209     Handled = 0;
210     Header = FALSE;
211     do
212     {
213         Ret = RegEnumKeyW(HandleKey, Index, SubKey, MaxLen);
214         if (Ret != ERROR_NO_MORE_ITEMS)
215         {
216             if (Image == NULL || wcscmp(SubKey, Image) == 0)
217             {
218                 if (DisplayImageInfo(HandleKey, SubKey, &Header))
219                 {
220                     ++Handled;
221                 }
222             }
223 
224             ++Index;
225         }
226     } while (Ret != ERROR_NO_MORE_ITEMS);
227 
228     if (Handled == 0)
229     {
230         if (Image == NULL)
231         {
232             wprintf(L"No application has page heap enabled.\n");
233         }
234         else
235         {
236             wprintf(L"Page heap is not enabled for %s\n", Image);
237         }
238     }
239 
240     HeapFree(GetProcessHeap(), 0, SubKey);
241     RegCloseKey(HandleKey);
242 }
243 
244 BOOL PageHeap_ParseCmdline(INT i, int argc, LPWSTR argv[])
245 {
246     for (; i < argc; i++)
247     {
248         if (argv[i][0] == L'/')
249         {
250             if (wcscmp(argv[i], L"/enable") == 0)
251             {
252                 Set = TRUE;
253             }
254             else if (wcscmp(argv[i], L"/disable") == 0)
255             {
256                 Unset = TRUE;
257             }
258             else if (wcscmp(argv[i], L"/full") == 0)
259             {
260                 Full = TRUE;
261             }
262         }
263         else if (Image == NULL)
264         {
265             Image = argv[i];
266         }
267         else
268         {
269             wprintf(L"Invalid option: %s\n", argv[i]);
270             return FALSE;
271         }
272     }
273 
274     if (Set && Unset)
275     {
276         wprintf(L"ENABLE and DISABLED cannot be set together\n");
277         return FALSE;
278     }
279 
280     if (Image == NULL && (Set || Unset || Full))
281     {
282         wprintf(L"Can't ENABLE or DISABLE with no image\n");
283         return FALSE;
284     }
285 
286     if (!Set && !Unset && Full)
287     {
288         wprintf(L"Cannot deal with full traces with no other indication\n");
289         return FALSE;
290     }
291 
292     return TRUE;
293 }
294 
295 INT PageHeap_Execute()
296 {
297     if (!Set && !Unset)
298     {
299         DisplayStatus();
300     }
301     else
302     {
303         ModifyStatus();
304     }
305 
306     return 0;
307 }
308