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     HeapFree(GetProcessHeap(), 0, Buffer);
121     RegCloseKey(IFEOKey);
122 }
123 
124 static BOOL DisplayImageInfo(HKEY HandleKey, PWSTR SubKey, PBOOL Header)
125 {
126     LONG Ret;
127     BOOL Handled;
128     DWORD GlobalFlags;
129     HKEY HandleSubKey;
130 
131     Ret = RegOpenKeyExW(HandleKey, SubKey, 0, KEY_READ, &HandleSubKey);
132     if (Ret != ERROR_SUCCESS)
133     {
134         wprintf(L"DII: RegOpenKeyEx failed (%d)\n", Ret);
135         return FALSE;
136     }
137 
138     Handled = FALSE;
139     GlobalFlags = ReadSZFlagsFromRegistry(HandleSubKey, L"GlobalFlag");
140     if (GlobalFlags & FLG_HEAP_PAGE_ALLOCS)
141     {
142         DWORD PageHeapFlags;
143 
144         if (Image == NULL)
145         {
146             if (!*Header)
147             {
148                 wprintf(L"path: %s\n", ImageExecOptionsString);
149                 *Header = TRUE;
150             }
151             wprintf(L"\t%s: page heap enabled with flags (", SubKey);
152         }
153         else
154         {
155             wprintf(L"Page heap is enabled for %s with flags (", SubKey);
156         }
157 
158         PageHeapFlags = ReadSZFlagsFromRegistry(HandleSubKey, L"PageHeapFlags");
159         if (PageHeapFlags & 0x1)
160         {
161             wprintf(L"full ");
162         }
163 
164         if (PageHeapFlags & 0x2)
165         {
166             wprintf(L"traces");
167         }
168 
169         wprintf(L")\n");
170 
171         Handled = TRUE;
172     }
173 
174     RegCloseKey(HandleSubKey);
175 
176     return Handled;
177 }
178 
179 static VOID DisplayStatus(VOID)
180 {
181     LONG Ret;
182     HKEY HandleKey;
183     DWORD Index, MaxLen, Handled;
184     TCHAR * SubKey;
185     BOOL Header;
186 
187     if (!OpenImageFileExecOptions(KEY_READ, NULL, &HandleKey))
188     {
189         return;
190     }
191 
192     Ret = RegQueryInfoKeyW(HandleKey, NULL, NULL, NULL, NULL, &MaxLen, NULL, NULL, NULL, NULL, NULL, NULL);
193     if (Ret != ERROR_SUCCESS)
194     {
195         wprintf(L"DS: RegQueryInfoKey failed (%d)\n", Ret);
196         RegCloseKey(HandleKey);
197         return;
198     }
199 
200     ++MaxLen; // NULL-char
201     SubKey = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, MaxLen * sizeof(WCHAR));
202     if (SubKey == NULL)
203     {
204         wprintf(L"DS: HeapAlloc failed\n");
205         RegCloseKey(HandleKey);
206         return;
207     }
208 
209     Index = 0;
210     Handled = 0;
211     Header = FALSE;
212     do
213     {
214         Ret = RegEnumKeyW(HandleKey, Index, SubKey, MaxLen);
215         if (Ret != ERROR_NO_MORE_ITEMS)
216         {
217             if (Image == NULL || wcscmp(SubKey, Image) == 0)
218             {
219                 if (DisplayImageInfo(HandleKey, SubKey, &Header))
220                 {
221                     ++Handled;
222                 }
223             }
224 
225             ++Index;
226         }
227     } while (Ret != ERROR_NO_MORE_ITEMS);
228 
229     if (Handled == 0)
230     {
231         if (Image == NULL)
232         {
233             wprintf(L"No application has page heap enabled.\n");
234         }
235         else
236         {
237             wprintf(L"Page heap is not enabled for %s\n", Image);
238         }
239     }
240 
241     HeapFree(GetProcessHeap(), 0, SubKey);
242     RegCloseKey(HandleKey);
243 }
244 
245 BOOL PageHeap_ParseCmdline(INT i, int argc, LPWSTR argv[])
246 {
247     for (; i < argc; i++)
248     {
249         if (argv[i][0] == L'/')
250         {
251             if (wcscmp(argv[i], L"/enable") == 0)
252             {
253                 Set = TRUE;
254             }
255             else if (wcscmp(argv[i], L"/disable") == 0)
256             {
257                 Unset = TRUE;
258             }
259             else if (wcscmp(argv[i], L"/full") == 0)
260             {
261                 Full = TRUE;
262             }
263         }
264         else if (Image == NULL)
265         {
266             Image = argv[i];
267         }
268         else
269         {
270             wprintf(L"Invalid option: %s\n", argv[i]);
271             return FALSE;
272         }
273     }
274 
275     if (Set && Unset)
276     {
277         wprintf(L"ENABLE and DISABLED cannot be set together\n");
278         return FALSE;
279     }
280 
281     if (Image == NULL && (Set || Unset || Full))
282     {
283         wprintf(L"Can't ENABLE or DISABLE with no image\n");
284         return FALSE;
285     }
286 
287     if (!Set && !Unset && Full)
288     {
289         wprintf(L"Cannot deal with full traces with no other indication\n");
290         return FALSE;
291     }
292 
293     return TRUE;
294 }
295 
296 INT PageHeap_Execute()
297 {
298     if (!Set && !Unset)
299     {
300         DisplayStatus();
301     }
302     else
303     {
304         ModifyStatus();
305     }
306 
307     return 0;
308 }
309