1#include "processes.h"
 2#include "common/debug.h"
 3#include "common/mallocHelper.h"
 4#include "common/windows/nt.h"
 5
 6#include <ntstatus.h>
 7#include <winternl.h>
 8
 9const char* ffDetectProcesses(const FFProcessesOptions* options, FFProcessesResult* result) {
10    if (options->countKprocs && ffIsSystemBasicProcessInfoAvailable()) {
11        // SystemHandleCountInformation reports the total process/thread counts directly,
12        // so we don't need to walk the whole process table.
13        SYSTEM_HANDLECOUNT_INFORMATION info = {}; // Seems that kernel only fills the lower 32 bits of the counts, leave the upper 32 bits untouched.
14        if (NT_SUCCESS(NtQuerySystemInformation(SystemHandleCountInformation, &info, sizeof(info), NULL))) {
15            result->processes = info.ProcessCount;
16            result->threads = info.ThreadCount;
17            return nullptr;
18        }
19        // Otherwise fall back to walking the process table
20    }
21
22    FF_AUTO_FREE SYSTEM_PROCESS_INFORMATION* pstart = nullptr;
23
24    // Multiple attempts in case processes change while
25    // we are in the middle of querying them.
26    ULONG size = 0;
27    for (int attempts = 0;; ++attempts) {
28        if (size) {
29            pstart = (SYSTEM_PROCESS_INFORMATION*) realloc(pstart, size);
30            assert(pstart);
31        }
32        NTSTATUS status = NtQuerySystemInformation(SystemProcessInformation, pstart, size, &size);
33        if (NT_SUCCESS(status)) {
34            break;
35        } else if (status == STATUS_INFO_LENGTH_MISMATCH && attempts < 4) {
36            size += sizeof(SYSTEM_PROCESS_INFORMATION) * 5;
37        } else {
38            FF_DEBUG("NtQuerySystemInformation(SystemProcessInformation) failed: %s", ffDebugNtStatus(status));
39            return "NtQuerySystemInformation(SystemProcessInformation) failed";
40        }
41    }
42
43    for (SYSTEM_PROCESS_INFORMATION* ptr = pstart; ; ptr = (SYSTEM_PROCESS_INFORMATION*) ((uint8_t*) ptr + ptr->NextEntryOffset)) {
44        if (options->countKprocs || (uintptr_t) ptr->InheritedFromUniqueProcessId > 0) {
45            ++result->processes;
46            result->threads += ptr->NumberOfThreads;
47        }
48        if (ptr->NextEntryOffset == 0) {
49            break;
50        }
51    }
52
53    return nullptr;
54}