1 /*
2  * PROJECT:     ReactOS Tasklist Command
3  * LICENSE:     GPL-2.0-or-later (https://spdx.org/licenses/GPL-2.0-or-later)
4  * PURPOSE:     Displays a list of currently running processes on the computer.
5  * COPYRIGHT:   Copyright 2021 He Yang <1160386205@qq.com>
6  */
7 
8 #include "tasklist.h"
9 
10 // the strings in OptionList are the command-line options.
11 // should always correspond with the defines below, in sequence (except OPTION_INVALID)
12 static PCWSTR OptionList[] = { L"?", L"nh" };
13 
14 #define OPTION_INVALID    -1
15 #define OPTION_HELP       0
16 #define OPTION_NOHEADER    1
17 
18 // the max string length PrintResString can handle
19 #define RES_STR_MAXLEN 64
20 
21 // Print split line
22 VOID PrintSplitLine(UINT Length)
23 {
24     for (; Length; Length--)
25     {
26         ConPuts(StdOut, L"=");
27     }
28 }
29 
30 // Print spaces
31 VOID PrintSpace(UINT Length)
32 {
33     ConPrintf(StdOut, L"%*ls", (INT)Length, L"");
34 }
35 
36 // Print a string.
37 // if bAlignLeft == TRUE then aligned to left, otherwise aligned to right
38 // MaxWidth is the width for printing.
39 VOID PrintString(LPCWSTR String, UINT MaxWidth, BOOL bAlignLeft)
40 {
41     ConPrintf(StdOut, bAlignLeft ? L"%-*.*ls" : L"%*.*ls", MaxWidth, MaxWidth, String);
42 }
43 
44 // Print a string from resource
45 // if bAlignLeft == TRUE then aligned to left, otherwise aligned to right
46 // MaxWidth is the width for printing.
47 // The string WILL be truncated if it's longer than RES_STR_MAXLEN
48 VOID PrintResString(HINSTANCE hInstance, UINT uID, UINT MaxWidth, BOOL bAlignLeft)
49 {
50     if (!hInstance)
51         return;
52 
53     WCHAR StringBuffer[RES_STR_MAXLEN];
54     LoadStringW(hInstance, uID, StringBuffer, _countof(StringBuffer));
55     PrintString(StringBuffer, MaxWidth, bAlignLeft);
56 }
57 
58 // Print a number, aligned to right.
59 // MaxWidth is the width for printing.
60 // the number WILL NOT be truncated if it's longer than MaxWidth
61 VOID PrintNum(LONGLONG Number, UINT MaxWidth)
62 {
63     ConPrintf(StdOut, L"%*lld", MaxWidth, Number);
64 }
65 
66 // Print memory size using KB as unit, with comma-separated number, aligned to right.
67 // MaxWidth is the width for printing.
68 // the number WILL be truncated if it's longer than MaxWidth
69 BOOL PrintMemory(SIZE_T MemorySizeByte, UINT MaxWidth, HINSTANCE hInstance)
70 {
71     if (!hInstance)
72         return FALSE;
73 
74     SIZE_T MemorySize = MemorySizeByte >> 10;
75 
76     WCHAR NumberString[27] = { 0 }; // length 26 is enough to display ULLONG_MAX in decimal with comma, one more for zero-terminated.
77     C_ASSERT(sizeof(SIZE_T) <= 8);
78 
79     PWCHAR pNumberStr = NumberString;
80 
81     // calculate the length
82     UINT PrintLength = 0;
83     SIZE_T Tmp = MemorySize;
84     UINT Mod = 1;
85 
86     do
87     {
88         Tmp /= 10;
89         PrintLength++;
90         Mod *= 10;
91     } while (Tmp);
92 
93     for (UINT i = PrintLength; i; i--)
94     {
95         Mod /= 10;
96         *pNumberStr = L'0' + (MemorySize / Mod);
97         MemorySize %= Mod;
98         pNumberStr++;
99 
100         if (i != 1 && i % 3 == 1)
101         {
102             *pNumberStr = L',';
103             pNumberStr++;
104         }
105     }
106 
107     WCHAR FormatStr[RES_STR_MAXLEN];
108     LoadStringW(hInstance, IDS_MEMORY_STR, FormatStr, _countof(FormatStr));
109 
110     WCHAR String[RES_STR_MAXLEN + _countof(NumberString)] = { 0 };
111 
112     StringCchPrintfW(String, _countof(String), FormatStr, NumberString);
113     PrintString(String, MaxWidth, FALSE);
114 
115     return TRUE;
116 }
117 
118 VOID PrintHeader(HINSTANCE hInstance)
119 {
120     if (!hInstance)
121         return;
122 
123     PrintResString(hInstance, IDS_HEADER_IMAGENAME, COLUMNWIDTH_IMAGENAME, TRUE);
124     PrintSpace(1);
125     PrintResString(hInstance, IDS_HEADER_PID,       COLUMNWIDTH_PID,       FALSE);
126     PrintSpace(1);
127     PrintResString(hInstance, IDS_HEADER_SESSION,   COLUMNWIDTH_SESSION,   FALSE);
128     PrintSpace(1);
129     PrintResString(hInstance, IDS_HEADER_MEMUSAGE,  COLUMNWIDTH_MEMUSAGE,  FALSE);
130 
131     ConPuts(StdOut, L"\n");
132 
133     PrintSplitLine(COLUMNWIDTH_IMAGENAME);
134     PrintSpace(1);
135     PrintSplitLine(COLUMNWIDTH_PID);
136     PrintSpace(1);
137     PrintSplitLine(COLUMNWIDTH_SESSION);
138     PrintSpace(1);
139     PrintSplitLine(COLUMNWIDTH_MEMUSAGE);
140 
141     ConPuts(StdOut, L"\n");
142 }
143 
144 BOOL EnumProcessAndPrint(BOOL bNoHeader)
145 {
146     // Call NtQuerySystemInformation for the process information
147     ULONG ProcessInfoBufferLength = 0;
148     ULONG ResultLength = 0;
149     PBYTE ProcessInfoBuffer = NULL;
150 
151     // Get the buffer size we need
152     NTSTATUS Status = NtQuerySystemInformation(SystemProcessInformation, NULL, 0, &ResultLength);
153 
154     // New process/thread might appear before we call for the actual data.
155     // Try to avoid this by retrying several times.
156     for (UINT Retry = 0; Retry < NT_SYSTEM_QUERY_MAX_RETRY; Retry++)
157     {
158         // (Re)allocate buffer
159         ProcessInfoBufferLength = ResultLength;
160         ResultLength = 0;
161         if (ProcessInfoBuffer)
162         {
163             PBYTE NewProcessInfoBuffer = HeapReAlloc(GetProcessHeap(), 0,
164                                                      ProcessInfoBuffer,
165                                                      ProcessInfoBufferLength);
166             if (NewProcessInfoBuffer)
167             {
168                 ProcessInfoBuffer = NewProcessInfoBuffer;
169             }
170             else
171             {
172                 // out of memory
173                 ConResMsgPrintf(StdErr, 0, IDS_OUT_OF_MEMORY);
174                 HeapFree(GetProcessHeap(), 0, ProcessInfoBuffer);
175                 return FALSE;
176             }
177         }
178         else
179         {
180             ProcessInfoBuffer = HeapAlloc(GetProcessHeap(), 0, ProcessInfoBufferLength);
181             if (!ProcessInfoBuffer)
182             {
183                 // out of memory
184                 ConResMsgPrintf(StdErr, 0, IDS_OUT_OF_MEMORY);
185                 return FALSE;
186             }
187         }
188 
189         // Query information
190         Status = NtQuerySystemInformation(SystemProcessInformation,
191                                           ProcessInfoBuffer,
192                                           ProcessInfoBufferLength,
193                                           &ResultLength);
194         if (Status != STATUS_INFO_LENGTH_MISMATCH)
195         {
196             break;
197         }
198     }
199 
200     if (!NT_SUCCESS(Status))
201     {
202         // tried NT_SYSTEM_QUERY_MAX_RETRY times, or failed with some other reason
203         ConResMsgPrintf(StdErr, 0, IDS_ENUM_FAILED);
204         HeapFree(GetProcessHeap(), 0, ProcessInfoBuffer);
205         return FALSE;
206     }
207 
208     HINSTANCE hInstance = GetModuleHandleW(NULL);
209     assert(hInstance);
210 
211     ConPuts(StdOut, L"\n");
212 
213     if (!bNoHeader)
214     {
215         PrintHeader(hInstance);
216     }
217 
218     PSYSTEM_PROCESS_INFORMATION pSPI;
219     pSPI = (PSYSTEM_PROCESS_INFORMATION)ProcessInfoBuffer;
220     while (pSPI)
221     {
222         PrintString(pSPI->UniqueProcessId ? pSPI->ImageName.Buffer : L"System Idle Process", COLUMNWIDTH_IMAGENAME, TRUE);
223         PrintSpace(1);
224         PrintNum((LONGLONG)(INT_PTR)pSPI->UniqueProcessId, COLUMNWIDTH_PID);
225         PrintSpace(1);
226         PrintNum((ULONGLONG)pSPI->SessionId, COLUMNWIDTH_SESSION);
227         PrintSpace(1);
228         PrintMemory(pSPI->WorkingSetSize, COLUMNWIDTH_MEMUSAGE, hInstance);
229 
230         ConPuts(StdOut, L"\n");
231 
232         if (pSPI->NextEntryOffset == 0)
233             break;
234         pSPI = (PSYSTEM_PROCESS_INFORMATION)((LPBYTE)pSPI + pSPI->NextEntryOffset);
235     }
236 
237     HeapFree(GetProcessHeap(), 0, ProcessInfoBuffer);
238     return TRUE;
239 }
240 
241 INT GetOptionType(LPCWSTR szOption)
242 {
243     if (szOption[0] != L'/' && szOption[0] != L'-')
244     {
245         return OPTION_INVALID;
246     }
247     szOption++;
248 
249     for (UINT i = 0; i < _countof(OptionList); i++)
250     {
251         if (!_wcsicmp(OptionList[i], szOption))
252         {
253             return i;
254         }
255     }
256     return OPTION_INVALID;
257 }
258 
259 BOOL ProcessArguments(INT argc, WCHAR **argv)
260 {
261     BOOL bHasHelp = FALSE, bHasNoHeader = FALSE;
262 
263     for (INT i = 1; i < argc; i++)
264     {
265         INT Option = GetOptionType(argv[i]);
266 
267         switch (Option)
268         {
269             case OPTION_HELP:
270             {
271                 if (bHasHelp)
272                 {
273                     // -? already specified
274                     ConResMsgPrintf(StdErr, 0, IDS_OPTION_TOO_MUCH, argv[i], 1);
275                     ConResMsgPrintf(StdErr, 0, IDS_USAGE);
276                     return FALSE;
277                 }
278                 bHasHelp = TRUE;
279                 break;
280             }
281             case OPTION_NOHEADER:
282             {
283                 if (bHasNoHeader)
284                 {
285                     // -nh already specified
286                     ConResMsgPrintf(StdErr, 0, IDS_OPTION_TOO_MUCH, argv[i], 1);
287                     ConResMsgPrintf(StdErr, 0, IDS_USAGE);
288                     return FALSE;
289                 }
290                 bHasNoHeader = TRUE;
291                 break;
292             }
293             case OPTION_INVALID:
294             default:
295             {
296                 ConResMsgPrintf(StdErr, 0, IDS_INVALID_OPTION);
297                 ConResMsgPrintf(StdErr, 0, IDS_USAGE);
298                 return FALSE;
299             }
300         }
301     }
302 
303     if (bHasHelp)
304     {
305         if (argc > 2) // any arguments other than -? exists
306         {
307             ConResMsgPrintf(StdErr, 0, IDS_INVALID_SYNTAX);
308             ConResMsgPrintf(StdErr, 0, IDS_USAGE);
309             return FALSE;
310         }
311         else
312         {
313             ConResMsgPrintf(StdOut, 0, IDS_USAGE);
314             ConResMsgPrintf(StdOut, 0, IDS_DESCRIPTION);
315             return FALSE;
316         }
317     }
318     else
319     {
320         EnumProcessAndPrint(bHasNoHeader);
321     }
322     return TRUE;
323 }
324 
325 int wmain(int argc, WCHAR **argv)
326 {
327     /* Initialize the Console Standard Streams */
328     ConInitStdStreams();
329 
330     if (!ProcessArguments(argc, argv))
331     {
332         return 1;
333     }
334     return 0;
335 }
336