1 /*
2  * COPYRIGHT:   See COPYING in the top level directory
3  * PROJECT:     ReactOS ping utility
4  * FILE:        applications/cmdutils/gflags/gflags.c
5  * PURPOSE:     Global Flags utility
6  * PROGRAMMERS: Pierre Schweitzer <pierre@reactos.org>
7  */
8 
9 #define WIN32_NO_STATUS
10 #include <stdarg.h>
11 #include <windef.h>
12 #include <winbase.h>
13 #include <winuser.h>
14 #include <winreg.h>
15 #include <stdio.h>
16 #include <stdlib.h>
17 
18 static BOOL Set = FALSE;
19 static BOOL Unset = FALSE;
20 static BOOL Full = FALSE;
21 static PWSTR Image = NULL;
22 static WCHAR ImageExecOptionsString[] = L"Software\\Microsoft\\Windows NT\\CurrentVersion\\Image File Execution Options";
23 
24 static DWORD ReagFlagsFromRegistry(HKEY SubKey, PVOID Buffer, PWSTR Value, DWORD MaxLen)
25 {
26     DWORD Len, Flags, Type;
27 
28     Len = MaxLen;
29     Flags = 0;
30     if (RegQueryValueEx(SubKey, Value, NULL, &Type, Buffer, &Len) == ERROR_SUCCESS && Type == REG_SZ)
31     {
32         Flags = wcstoul(Buffer, NULL, 16);
33     }
34 
35     return Flags;
36 }
37 
38 static VOID ModifyStatus(VOID)
39 {
40     LONG Ret;
41     DWORD MaxLen, GlobalFlags;
42     PVOID Buffer;
43     HKEY HandleKey, HandleSubKey;
44 
45     Ret = RegOpenKeyEx(HKEY_LOCAL_MACHINE, ImageExecOptionsString, 0, KEY_WRITE | KEY_READ, &HandleKey);
46     if (Ret != ERROR_SUCCESS)
47     {
48         wprintf(L"MS: RegOpenKeyEx failed (%d)\n", Ret);
49         return;
50     }
51 
52     Ret = RegCreateKeyEx(HandleKey, Image, 0, NULL, REG_OPTION_NON_VOLATILE, KEY_WRITE | KEY_READ, NULL, &HandleSubKey, NULL);
53     if (Ret != ERROR_SUCCESS)
54     {
55         wprintf(L"MS: RegCreateKeyEx failed (%d)\n", Ret);
56         return;
57     }
58 
59     Ret = RegQueryInfoKey(HandleSubKey, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, &MaxLen, NULL, NULL);
60     if (Ret != ERROR_SUCCESS)
61     {
62         wprintf(L"MS: RegQueryInfoKey failed (%d)\n", Ret);
63         RegCloseKey(HandleSubKey);
64         RegCloseKey(HandleKey);
65         return;
66     }
67 
68     MaxLen = max(MaxLen, 11 * sizeof(WCHAR));
69     Buffer = HeapAlloc(GetProcessHeap(), 0, MaxLen);
70     if (Buffer == NULL)
71     {
72         wprintf(L"MS: HeapAlloc failed\n");
73         RegCloseKey(HandleSubKey);
74         RegCloseKey(HandleKey);
75         return;
76     }
77 
78     GlobalFlags = ReagFlagsFromRegistry(HandleSubKey, Buffer, L"GlobalFlag", MaxLen);
79     if (Set)
80     {
81         GlobalFlags |= 0x02000000;
82     }
83     else
84     {
85         GlobalFlags &= ~0x02000000;
86     }
87 
88     if (GlobalFlags != 0)
89     {
90         wsprintf(Buffer, L"0x%08x", GlobalFlags);
91         Ret = RegSetValueEx(HandleSubKey, L"GlobalFlag", 0, REG_SZ, Buffer, 11 * sizeof(WCHAR));
92         if (Ret != ERROR_SUCCESS)
93         {
94             wprintf(L"MS: RegSetValueEx failed (%d)\n", Ret);
95         }
96     }
97     else
98     {
99         Ret = RegDeleteValue(HandleSubKey, L"GlobalFlag");
100         if (Ret != ERROR_SUCCESS)
101         {
102             wprintf(L"MS: RegDeleteValue failed (%d)\n", Ret);
103         }
104     }
105 
106     if (Unset)
107     {
108         Ret = RegDeleteValue(HandleSubKey, L"PageHeapFlags");
109         if (Ret != ERROR_SUCCESS)
110         {
111             wprintf(L"MS: RegDeleteValue failed (%d)\n", Ret);
112         }
113     }
114     else
115     {
116         DWORD PageHeapFlags;
117 
118         PageHeapFlags = ReagFlagsFromRegistry(HandleSubKey, Buffer, L"PageHeapFlags", MaxLen);
119         PageHeapFlags &= ~3;
120 
121         if (Full)
122         {
123             PageHeapFlags |= 1;
124         }
125         PageHeapFlags |= 2;
126 
127         wsprintf(Buffer, L"0x%x", PageHeapFlags);
128         Ret = RegSetValueEx(HandleSubKey, L"PageHeapFlags", 0, REG_SZ, Buffer, 11 * sizeof(WCHAR));
129         if (Ret != ERROR_SUCCESS)
130         {
131             wprintf(L"MS: RegSetValueEx failed (%d)\n", Ret);
132         }
133     }
134 
135     if (Set)
136     {
137         DWORD Type, VerifierFlags, Len;
138 
139         VerifierFlags = 0;
140         Len = MaxLen;
141         if (RegQueryValueEx(HandleSubKey, L"VerifierFlags", NULL, &Type, Buffer, &Len) == ERROR_SUCCESS &&
142             Type == REG_DWORD && Len == sizeof(DWORD))
143         {
144             VerifierFlags = ((DWORD *)Buffer)[0];
145             VerifierFlags &= ~0x8001;
146         }
147 
148         if (Full)
149         {
150             VerifierFlags |= 1;
151         }
152         else
153         {
154             VerifierFlags |= 0x8000;
155         }
156 
157         Ret = RegSetValueEx(HandleSubKey, L"VerifierFlags", 0, REG_DWORD, (const BYTE *)&VerifierFlags, sizeof(DWORD));
158         if (Ret != ERROR_SUCCESS)
159         {
160             wprintf(L"MS: RegSetValueEx failed (%d)\n", Ret);
161         }
162     }
163 
164     wprintf(L"path: %s\n", ImageExecOptionsString);
165     wprintf(L"\t%s: page heap %s\n", Image, (Set ? L"enabled" : L"disabled"));
166 
167     HeapFree(GetProcessHeap(), 0, Buffer);
168     RegCloseKey(HandleSubKey);
169     RegCloseKey(HandleKey);
170 }
171 
172 static BOOL DisplayImageInfo(HKEY HandleKey, PWSTR SubKey, PBOOL Header)
173 {
174     LONG Ret;
175     BOOL Handled;
176     DWORD MaxLen, GlobalFlags;
177     HKEY HandleSubKey;
178     PVOID Buffer;
179 
180     Ret = RegOpenKeyEx(HandleKey, SubKey, 0, KEY_READ, &HandleSubKey);
181     if (Ret != ERROR_SUCCESS)
182     {
183         wprintf(L"DII: RegOpenKeyEx failed (%d)\n", Ret);
184         return FALSE;
185     }
186 
187     Ret = RegQueryInfoKey(HandleSubKey, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, &MaxLen, NULL, NULL);
188     if (Ret != ERROR_SUCCESS)
189     {
190         wprintf(L"DII: RegQueryInfoKey failed (%d)\n", Ret);
191         RegCloseKey(HandleSubKey);
192         return FALSE;
193     }
194 
195     Buffer = HeapAlloc(GetProcessHeap(), 0, MaxLen);
196     if (Buffer == NULL)
197     {
198         wprintf(L"DII: HeapAlloc failed\n");
199         RegCloseKey(HandleSubKey);
200         return FALSE;
201     }
202 
203     Handled = FALSE;
204     GlobalFlags = ReagFlagsFromRegistry(HandleSubKey, Buffer, L"GlobalFlag", MaxLen);
205     if (GlobalFlags & 0x02000000)
206     {
207         DWORD PageHeapFlags;
208 
209         if (Image == NULL)
210         {
211             if (!*Header)
212             {
213                 wprintf(L"path: %s\n", ImageExecOptionsString);
214                 *Header = TRUE;
215             }
216             wprintf(L"\t%s: page heap enabled with flags (", SubKey);
217         }
218         else
219         {
220             wprintf(L"Page heap is enabled for %s with flags (", SubKey);
221         }
222 
223         PageHeapFlags = ReagFlagsFromRegistry(HandleSubKey, Buffer, L"PageHeapFlags", MaxLen);
224         if (PageHeapFlags & 0x1)
225         {
226             wprintf(L"full ");
227         }
228 
229         if (PageHeapFlags & 0x2)
230         {
231             wprintf(L"traces");
232         }
233 
234         wprintf(L")\n");
235 
236         Handled = TRUE;
237     }
238 
239     HeapFree(GetProcessHeap(), 0, Buffer);
240     RegCloseKey(HandleSubKey);
241 
242     return Handled;
243 }
244 
245 static VOID DisplayStatus(VOID)
246 {
247     LONG Ret;
248     HKEY HandleKey;
249     DWORD Index, MaxLen, Handled;
250     TCHAR * SubKey;
251     BOOL Header;
252 
253     Ret = RegOpenKeyEx(HKEY_LOCAL_MACHINE, ImageExecOptionsString, 0, KEY_READ, &HandleKey);
254     if (Ret != ERROR_SUCCESS)
255     {
256         wprintf(L"DS: RegOpenKeyEx failed (%d)\n", Ret);
257         return;
258     }
259 
260     Ret = RegQueryInfoKey(HandleKey, NULL, NULL, NULL, NULL, &MaxLen, NULL, NULL, NULL, NULL, NULL, NULL);
261     if (Ret != ERROR_SUCCESS)
262     {
263         wprintf(L"DS: RegQueryInfoKey failed (%d)\n", Ret);
264         RegCloseKey(HandleKey);
265         return;
266     }
267 
268     ++MaxLen; // NULL-char
269     SubKey = HeapAlloc(GetProcessHeap(), 0, MaxLen * sizeof(TCHAR));
270     if (SubKey == NULL)
271     {
272         wprintf(L"DS: HeapAlloc failed\n");
273         RegCloseKey(HandleKey);
274         return;
275     }
276 
277     Index = 0;
278     Handled = 0;
279     Header = FALSE;
280     do
281     {
282         Ret = RegEnumKey(HandleKey, Index, SubKey, MaxLen);
283         if (Ret != ERROR_NO_MORE_ITEMS)
284         {
285             if (Image == NULL || wcscmp(SubKey, Image) == 0)
286             {
287                 if (DisplayImageInfo(HandleKey, SubKey, &Header))
288                 {
289                     ++Handled;
290                 }
291             }
292 
293             ++Index;
294         }
295     } while (Ret != ERROR_NO_MORE_ITEMS);
296 
297     if (Handled == 0)
298     {
299         if (Image == NULL)
300         {
301             wprintf(L"No application has page heap enabled.\n");
302         }
303         else
304         {
305             wprintf(L"Page heap is not enabled for %s\n", Image);
306         }
307     }
308 
309     HeapFree(GetProcessHeap(), 0, SubKey);
310     RegCloseKey(HandleKey);
311 }
312 
313 static VOID Usage(VOID)
314 {
315     wprintf(L"Usage: gflags /p [image.exe] [/enable|/disable [/full]]\n"
316             L"\timage.exe:\tImage you want to deal with\n"
317             L"\t/enable:\tenable page heap for the image\n"
318             L"\t/disable:\tdisable page heap for the image\n"
319             L"\t/full:\t\tactivate full debug page heap\n");
320 }
321 
322 static BOOL ParseCmdline(int argc, LPWSTR argv[])
323 {
324     INT i;
325     BOOL UsePageHeap = FALSE;
326 
327     if (argc < 2)
328     {
329         wprintf(L"Not enough args!\n", argc);
330         Usage();
331         return FALSE;
332     }
333 
334     for (i = 1; i < argc; i++)
335     {
336         if (argv[i][0] == L'/')
337         {
338             if (argv[i][1] == L'p' && argv[i][2] == UNICODE_NULL)
339             {
340                 UsePageHeap = TRUE;
341             }
342             else if (argv[i][1] == L'p' && argv[i][2] != UNICODE_NULL)
343             {
344                 wprintf(L"Invalid option: %s\n", argv[i]);
345                 Usage();
346                 return FALSE;
347             }
348             else
349             {
350                 if (wcscmp(argv[i], L"/enable") == 0)
351                 {
352                     Set = TRUE;
353                 }
354                 else if (wcscmp(argv[i], L"/disable") == 0)
355                 {
356                     Unset = TRUE;
357                 }
358                 else if (wcscmp(argv[i], L"/full") == 0)
359                 {
360                     Full = TRUE;
361                 }
362             }
363         }
364         else if (Image == NULL)
365         {
366             Image = argv[i];
367         }
368         else
369         {
370             wprintf(L"Invalid option: %s\n", argv[i]);
371             Usage();
372             return FALSE;
373         }
374     }
375 
376     if (!UsePageHeap)
377     {
378         wprintf(L"Only page heap flags are supported\n");
379         Usage();
380         return FALSE;
381     }
382 
383     if (Set && Unset)
384     {
385         wprintf(L"ENABLE and DISABLED cannot be set together\n");
386         Usage();
387         return FALSE;
388     }
389 
390     if (Image == NULL && (Set || Unset || Full))
391     {
392         wprintf(L"Can't ENABLE or DISABLE with no image\n");
393         Usage();
394         return FALSE;
395     }
396 
397     if (!Set && !Unset && Full)
398     {
399         wprintf(L"Cannot deal with full traces with no other indication\n");
400         Usage();
401         return FALSE;
402     }
403 
404     return TRUE;
405 }
406 
407 int wmain(int argc, LPWSTR argv[])
408 {
409     if (!ParseCmdline(argc, argv))
410     {
411         return 1;
412     }
413 
414     if (!Set && !Unset)
415     {
416         DisplayStatus();
417     }
418     else
419     {
420         ModifyStatus();
421     }
422 
423     return 0;
424 }
425