├── res
├── 1.png
├── 2.png
└── 3.png
├── src
├── UnitTests
│ ├── UnitTests.vcxproj.user
│ ├── pch.cpp
│ ├── raw_assembly.asm
│ ├── pch.h
│ ├── UnitTests.vcxproj.filters
│ ├── UnitTests.cpp
│ └── UnitTests.vcxproj
├── Hunt-Sleeping-Beacons
│ ├── Hunt-Sleeping-Beacons.vcxproj.user
│ ├── calltrace.hpp
│ ├── thread.hpp
│ ├── private_memory.cpp
│ ├── hunt-sleeping-beacons.hpp
│ ├── hardware_breakpoints.cpp
│ ├── scans.hpp
│ ├── non_executable_memory.cpp
│ ├── abnormal_intermodular_call.cpp
│ ├── process.hpp
│ ├── detection.hpp
│ ├── blocking_apc.cpp
│ ├── misc.hpp
│ ├── thread_builder.hpp
│ ├── stomped_module.cpp
│ ├── main.cpp
│ ├── return_address_spoofing.cpp
│ ├── logger.hpp
│ ├── process_scanner.hpp
│ ├── process_enumerator.hpp
│ ├── calltrace_builder.hpp
│ ├── blocking_timer.cpp
│ ├── Hunt-Sleeping-Beacons.vcxproj.filters
│ ├── process_builder.hpp
│ ├── suspicious_timer.cpp
│ └── Hunt-Sleeping-Beacons.vcxproj
└── Hunt-Sleeping-Beacons.sln
├── Readme.md
└── inc
├── BS_thread_pool_utils.hpp
├── threadpooling.h
└── BS_thread_pool.hpp
/res/1.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/thefLink/Hunt-Sleeping-Beacons/HEAD/res/1.png
--------------------------------------------------------------------------------
/res/2.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/thefLink/Hunt-Sleeping-Beacons/HEAD/res/2.png
--------------------------------------------------------------------------------
/res/3.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/thefLink/Hunt-Sleeping-Beacons/HEAD/res/3.png
--------------------------------------------------------------------------------
/src/UnitTests/UnitTests.vcxproj.user:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
--------------------------------------------------------------------------------
/src/UnitTests/pch.cpp:
--------------------------------------------------------------------------------
1 | // pch.cpp: source file corresponding to the pre-compiled header
2 |
3 | #include "pch.h"
4 |
5 | // When you are using pre-compiled headers, this source file is necessary for compilation to succeed.
6 |
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/Hunt-Sleeping-Beacons.vcxproj.user:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
--------------------------------------------------------------------------------
/src/UnitTests/raw_assembly.asm:
--------------------------------------------------------------------------------
1 | .code
2 |
3 | public retrbx
4 |
5 | retrbx PROC
6 | jmp QWORD PTR [rbx]
7 | retrbx ENDP
8 |
9 | getgadget PROC
10 | call next
11 | next:
12 | pop rax
13 | sub rax, 7
14 | ret
15 | getgadget ENDP
16 |
17 | END
--------------------------------------------------------------------------------
/src/UnitTests/pch.h:
--------------------------------------------------------------------------------
1 | // pch.h: This is a precompiled header file.
2 | // Files listed below are compiled only once, improving build performance for future builds.
3 | // This also affects IntelliSense performance, including code completion and many code browsing features.
4 | // However, files listed here are ALL re-compiled if any one of them is updated between builds.
5 | // Do not add files here that you will be updating frequently as this negates the performance advantage.
6 |
7 | #ifndef PCH_H
8 | #define PCH_H
9 |
10 | // add headers that you want to pre-compile here
11 |
12 | #endif //PCH_H
13 |
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/calltrace.hpp:
--------------------------------------------------------------------------------
1 | #pragma once
2 |
3 | #include
4 | #include
5 |
6 | namespace hsb::containers {
7 |
8 | class calltrace{
9 | public:
10 | inline calltrace();
11 | inline ~calltrace();
12 | std::vector raw_addresses;
13 | std::vector modules;
14 | std::vector syms;
15 | };
16 |
17 | //implementation
18 | //================================================================================================
19 |
20 | #pragma region constructor and destructor
21 | calltrace::calltrace() {}
22 | calltrace::~calltrace() {}
23 | #pragma endregion
24 |
25 |
26 |
27 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/thread.hpp:
--------------------------------------------------------------------------------
1 | #pragma once
2 |
3 | #include
4 | #include
5 |
6 | #include "calltrace.hpp"
7 |
8 | namespace hsb::containers {
9 | class thread{
10 | public:
11 | inline thread();
12 | inline ~thread();
13 | DWORD tid;
14 | HANDLE handle;
15 | uint64_t stackbase;
16 | std::unique_ptr calltrace;
17 | };
18 |
19 | //implementation
20 | //================================================================================================
21 |
22 | #pragma region constructor and destructor
23 | thread::thread()
24 | : tid(0)
25 | , handle(nullptr)
26 | , stackbase(0)
27 | {}
28 |
29 | thread::~thread() {
30 | CloseHandle(handle);
31 | }
32 |
33 | #pragma endregion
34 |
35 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/private_memory.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include
4 |
5 | #include "phnt.h"
6 | #include "scans.hpp"
7 |
8 | namespace hsb::scanning {
9 |
10 | thread_scan thread_scans::private_memory = [](process* process, thread* thread) {
11 |
12 | BOOL bSuspicious = FALSE, bSuccess = FALSE;
13 |
14 | for (std::string module : thread->calltrace->modules) {
15 |
16 | if (module == "unknown") {
17 | bSuspicious = TRUE;
18 | break;
19 | }
20 | }
21 |
22 | if (bSuspicious) {
23 |
24 | thread_detection detection;
25 | detection.name = L"Abnormal Page in Callstack";
26 | detection.description = L"Callstack contains private memory regions";
27 | detection.tid = thread->tid;
28 | detection.severity = hsb::containers::detections::MID;
29 |
30 | process->add_detection(detection);
31 |
32 | }
33 |
34 | };
35 |
36 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/hunt-sleeping-beacons.hpp:
--------------------------------------------------------------------------------
1 | #include
2 |
3 | #include "process_enumerator.hpp"
4 | #include "process_scanner.hpp"
5 | #include "misc.hpp"
6 |
7 |
8 | namespace hsb {
9 |
10 | using process = hsb::containers::process;
11 | using process_scanner = hsb::scanning::process_scanner;
12 | using process_enumerator = hsb::containers::process_enumerator;
13 | using token_helpers = hsb::misc::token_helpers;
14 |
15 | std::vector> hunt(uint16_t pid = 0, bool ignore_dotnet) {
16 |
17 | process_enumerator process_enumerator(pid, ignore_dotnet);
18 | process_scanner process_scanner;
19 | std::vector> scanned_processes;
20 |
21 | if (token_helpers::is_elevated() == FALSE)
22 | return scanned_processes;
23 |
24 | if (token_helpers::set_debug_privilege() == FALSE)
25 | return scanned_processes;
26 |
27 | scanned_processes = process_enumerator.enumerate_processes();
28 | return scanned_processes;
29 |
30 | }
31 |
32 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/hardware_breakpoints.cpp:
--------------------------------------------------------------------------------
1 | #pragma once
2 |
3 | #include "phnt.h"
4 | #include "scans.hpp"
5 |
6 | namespace hsb::scanning {
7 |
8 | thread_scan thread_scans::hardware_breakpoints = [](process* process, thread* thread) {
9 |
10 | CONTEXT context = { 0 };
11 | context.ContextFlags = CONTEXT_ALL;
12 | bool bFound = false;
13 |
14 |
15 | if (GetThreadContext(thread->handle, &context))
16 | {
17 |
18 | bFound = (context.Dr0 != 0) || (context.Dr1 != 0) || (context.Dr2 != 0) || (context.Dr3 != 0 || (context.Dr7 != 0));
19 | if (bFound)
20 | {
21 |
22 | thread_detection detection;
23 |
24 | detection.name = L"Identified enabled hardware-breakpoints";
25 | detection.description = L"Often used as part of patchless-modifications of code";
26 | detection.tid = thread->tid;
27 | detection.severity = hsb::containers::detections::CRITICAL;
28 |
29 | process->add_detection(detection);
30 |
31 | }
32 | }
33 |
34 | };
35 |
36 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/scans.hpp:
--------------------------------------------------------------------------------
1 | #pragma once
2 |
3 | #include
4 | #include
5 |
6 | #include "detection.hpp"
7 | #include "process.hpp"
8 |
9 | namespace hsb::scanning {
10 |
11 | using process = hsb::containers::process;
12 | using thread = hsb::containers::thread;
13 | using thread_detection = hsb::containers::detections::thread_detection;
14 | using process_detection = hsb::containers::detections::process_detection;
15 |
16 | typedef std::function process_scan;
17 | typedef std::function thread_scan;
18 |
19 | struct process_scans {
20 |
21 | process_scans() = delete;
22 |
23 | static process_scan suspicious_timer;
24 |
25 | };
26 |
27 | struct thread_scans {
28 |
29 | thread_scans() = delete;
30 |
31 | static thread_scan private_memory;
32 | static thread_scan stomped_module;
33 | static thread_scan blocking_apc;
34 | static thread_scan blocking_timer;
35 | static thread_scan abnormal_intermodular_call;
36 | static thread_scan return_address_spoofing;
37 | static thread_scan hardware_breakpoints;
38 | static thread_scan non_executable_memory;
39 |
40 | };
41 |
42 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/non_executable_memory.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include
4 |
5 | #include "phnt.h"
6 | #include "scans.hpp"
7 |
8 | namespace hsb::scanning {
9 |
10 | thread_scan thread_scans::non_executable_memory = [](process* process,thread* thread) {
11 |
12 | BOOL bSuspicious = FALSE,bSuccess = FALSE;
13 | SIZE_T s = 0;
14 | MEMORY_BASIC_INFORMATION mbi ={0};
15 |
16 | for(uint64_t ret : thread->calltrace->raw_addresses) {
17 |
18 | s = VirtualQueryEx(process->handle,(LPCVOID)ret,&mbi,sizeof(MEMORY_BASIC_INFORMATION));
19 | if(s == 0)
20 | continue;
21 |
22 | if((mbi.Protect & (PAGE_EXECUTE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY)) == 0) {
23 | bSuspicious = TRUE;
24 | break;
25 | }
26 |
27 | }
28 |
29 | if(bSuspicious) {
30 |
31 | thread_detection detection;
32 | detection.name = L"Non-Executable Page in Callstack";
33 | detection.description = L"Callstack contains memory regions marked as non-executable";
34 | detection.tid = thread->tid;
35 | detection.severity = hsb::containers::detections::MID;
36 |
37 | process->add_detection(detection);
38 |
39 | }
40 |
41 | };
42 |
43 | }
--------------------------------------------------------------------------------
/src/UnitTests/UnitTests.vcxproj.filters:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 | {4FC737F1-C7A5-4376-A066-2A32D752A2FF}
6 | cpp;c;cc;cxx;c++;cppm;ixx;def;odl;idl;hpj;bat;asm;asmx
7 |
8 |
9 | {93995380-89BD-4b04-88EB-625FBE52EBFB}
10 | h;hh;hpp;hxx;h++;hm;inl;inc;ipp;xsd
11 |
12 |
13 | {67DA6AB6-F800-4c08-8B7A-83BB121AAD01}
14 | rc;ico;cur;bmp;dlg;rc2;rct;bin;rgs;gif;jpg;jpeg;jpe;resx;tiff;tif;png;wav;mfcribbon-ms
15 |
16 |
17 |
18 |
19 | Source Files
20 |
21 |
22 | Source Files
23 |
24 |
25 |
26 |
27 | Header Files
28 |
29 |
30 |
31 |
32 | Source Files
33 |
34 |
35 |
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/abnormal_intermodular_call.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include
4 |
5 | #include "phnt.h"
6 | #include "misc.hpp"
7 | #include "scans.hpp"
8 |
9 | namespace hsb::scanning {
10 |
11 |
12 | thread_scan thread_scans::abnormal_intermodular_call = [](process* process, thread* thread) {
13 |
14 | BOOL bSuccess = FALSE;
15 |
16 | if (thread->calltrace->raw_addresses.size() <= 2)
17 | goto Cleanup;
18 |
19 | for (int i = 0; i < thread->calltrace->raw_addresses.size(); i++) {
20 |
21 | std::string module_tmp = thread->calltrace->modules.at(i);
22 |
23 | if (!_stricmp(module_tmp.c_str(), "kernelbase") || !_stricmp(module_tmp.c_str(), "kernel32")) {
24 |
25 | if (i == thread->calltrace->raw_addresses.size() - 2)
26 | break;
27 |
28 | module_tmp = thread->calltrace->modules.at(i + 1);
29 | if (!_stricmp(module_tmp.c_str(), "ntdll")) {
30 |
31 | thread_detection detection;
32 | detection.name = L"Abnormal Intermodular Call";
33 | detection.description = std::format(L"{} called {}. This indicates module-proxying.", misc::string_to_wstring(thread->calltrace->syms.at(i + 1)), misc::string_to_wstring(thread->calltrace->syms.at(i)));
34 | detection.tid = thread->tid;
35 | detection.severity = hsb::containers::detections::CRITICAL;
36 |
37 | process->add_detection(detection);
38 |
39 | }
40 |
41 | }
42 |
43 | }
44 |
45 | Cleanup:
46 |
47 | return;
48 |
49 | };
50 |
51 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/process.hpp:
--------------------------------------------------------------------------------
1 | #pragma once
2 |
3 | #include "dbghelp.h"
4 | #include
5 | #include
6 | #include
7 |
8 | #include "detection.hpp"
9 | #include "thread.hpp"
10 |
11 | namespace hsb::containers {
12 |
13 | class process{
14 |
15 | using detection = hsb::containers::detections::detection;
16 |
17 | public:
18 | inline process();
19 | inline ~process();
20 | HANDLE handle;
21 | DWORD pid;
22 | std::wstring imagename;
23 | std::wstring cmdline;
24 | std::vector> threads;
25 | std::vector> detections;
26 |
27 | template
28 | void add_detection(const T&);
29 |
30 | private:
31 | std::mutex mutex;
32 | };
33 |
34 | //implementation
35 | //================================================================================================
36 |
37 | #pragma region constructor and destructor
38 | process::process()
39 | : handle(nullptr)
40 | , pid(0)
41 | , imagename(L"")
42 | , cmdline(L"")
43 | { }
44 |
45 | process::~process() {
46 | CloseHandle(handle);
47 | }
48 |
49 | #pragma endregion
50 |
51 | #pragma region private methods
52 | template
53 | void process::add_detection(const T &d){
54 | static_assert(std::is_base_of::value,"T must inherit from detection");
55 |
56 | std::lock_guard lock(mutex);
57 | detections.push_back(std::make_unique(d));
58 |
59 | }
60 | #pragma endregion
61 |
62 |
63 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/detection.hpp:
--------------------------------------------------------------------------------
1 | #pragma once
2 |
3 | #include
4 | #include
5 | #include
6 | #include
7 |
8 | namespace hsb::containers::detections{
9 |
10 | enum severity {
11 | LOW,
12 | MID,
13 | HIGH,
14 | CRITICAL
15 | };
16 |
17 | struct detection{
18 |
19 | std::wstring name;
20 | std::wstring description;
21 | severity severity;
22 |
23 | virtual std::wstring to_string() const = 0;
24 |
25 | };
26 |
27 | static constexpr std::array, 4> severity_info = { {
28 | {L"\033[32m", L"LOW"}, // Green for LOW
29 | {L"\033[33m", L"MID"}, // Yellow for MID
30 | {L"\033[31m", L"HIGH"}, // Red for HIGH
31 | {L"\033[35m", L"CRITICAL"} // Magenta for CRITICAL
32 | } };
33 |
34 | struct process_detection: public detection {
35 |
36 | std::wstring to_string() const override {
37 |
38 | const auto& [color, severity_str] = severity_info[static_cast(severity)];
39 |
40 | return std::format(L"! {}{}{}[0m | {}{}{}[0m | Severity: {}{}{}[0m",
41 | color, name, L"\033",
42 | color, description, L"\033",
43 | color, severity_str, L"\033");
44 | }
45 |
46 | };
47 |
48 | struct thread_detection: public detection {
49 |
50 | DWORD tid = 0;
51 |
52 | std::wstring to_string() const override {
53 |
54 | const auto& [color, severity_str] = severity_info[static_cast(severity)];
55 |
56 | return std::format(L"! Thread {} | {}{}{}[0m | {}{}{}[0m | Severity: {}{}{}[0m",
57 | tid,
58 | color, name, L"\033",
59 | color, description, L"\033",
60 | color, severity_str, L"\033");
61 |
62 | }
63 |
64 | };
65 |
66 |
67 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/blocking_apc.cpp:
--------------------------------------------------------------------------------
1 | #pragma once
2 |
3 | #include "phnt.h"
4 | #include "scans.hpp"
5 |
6 | namespace hsb::scanning {
7 |
8 | static inline std::once_flag apc_resolved;
9 |
10 | static DWORD64 pKiUserAPCDispatcher = 0;
11 | static bool resolve_apc_dispatcher(void);
12 |
13 | thread_scan thread_scans::blocking_apc = [](process* process,thread* thread) {
14 |
15 | BOOL bSuccess = FALSE;
16 | PVOID readAddr = NULL;
17 | DWORD64 stackAddr = 0;
18 | size_t numRead = 0;
19 |
20 | CONTEXT context ={0};
21 | context.ContextFlags = CONTEXT_ALL;
22 |
23 | std::call_once(apc_resolved, resolve_apc_dispatcher);
24 | if (pKiUserAPCDispatcher == 0) {
25 | return;
26 | }
27 |
28 | bSuccess = GetThreadContext(thread->handle,&context);
29 | if(bSuccess == FALSE)
30 | return;
31 |
32 | for(DWORD64 i = context.Rsp ; i <= (DWORD64)thread->stackbase; i += 8) {
33 |
34 | bSuccess = ReadProcessMemory(process->handle,(LPCVOID)i,&stackAddr,sizeof(DWORD64),&numRead);
35 | if(bSuccess == FALSE)
36 | break;
37 |
38 | if(stackAddr >= (DWORD64)pKiUserAPCDispatcher && (pKiUserAPCDispatcher + 120) > stackAddr) {
39 |
40 | thread_detection detection;
41 | detection.name = L"Blocking APC detected";
42 | detection.description = L"Thread's state triggered by ntdll!kiuserapcdispatcher";
43 | detection.tid = thread->tid;
44 | detection.severity = hsb::containers::detections::HIGH;
45 |
46 | process->add_detection(detection);
47 |
48 | }
49 | }
50 |
51 | };
52 |
53 | static bool resolve_apc_dispatcher(void) {
54 |
55 | HMODULE hNtdll = NULL;
56 |
57 | hNtdll = GetModuleHandleA("ntdll.dll");
58 | if (hNtdll == NULL)
59 | return false;
60 |
61 | pKiUserAPCDispatcher = (DWORD64)GetProcAddress(hNtdll, "KiUserApcDispatcher");
62 | if (pKiUserAPCDispatcher == 0)
63 | return false;
64 |
65 | return true;
66 |
67 | }
68 |
69 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/misc.hpp:
--------------------------------------------------------------------------------
1 | #pragma once
2 |
3 | #include "phnt.h"
4 |
5 | #include
6 | #include
7 |
8 | namespace hsb::misc {
9 |
10 | class token_helpers{
11 |
12 | private:
13 | token_helpers() = delete;
14 |
15 | public:
16 |
17 | static BOOL is_elevated(VOID) {
18 |
19 | BOOL fRet = FALSE;
20 | HANDLE hToken = NULL;
21 | if(OpenProcessToken(GetCurrentProcess(),TOKEN_QUERY,&hToken)) {
22 | TOKEN_ELEVATION Elevation ={0};
23 | DWORD cbSize = sizeof(TOKEN_ELEVATION);
24 | if(GetTokenInformation(hToken,TokenElevation,&Elevation,sizeof(Elevation),&cbSize)) {
25 | fRet = Elevation.TokenIsElevated;
26 | }
27 | }
28 | if(hToken) {
29 | CloseHandle(hToken);
30 | }
31 |
32 | return fRet;
33 |
34 | }
35 |
36 | //https://github.com/outflanknl/Dumpert/blob/master/Dumpert/Outflank-Dumpert/Dumpert.c Is Elevated() and SetDebugPrivilege was taken from here :).
37 | static BOOL set_debug_privilege(VOID) {
38 | HANDLE hToken = NULL;
39 | TOKEN_PRIVILEGES TokenPrivileges ={0};
40 |
41 | if(!OpenProcessToken(GetCurrentProcess(),TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES,&hToken)) {
42 | return FALSE;
43 | }
44 |
45 | TokenPrivileges.PrivilegeCount = 1;
46 | TokenPrivileges.Privileges[0].Attributes = TRUE ? SE_PRIVILEGE_ENABLED : 0;
47 |
48 | const wchar_t* lpwPriv = L"SeDebugPrivilege";
49 | if(!LookupPrivilegeValueW(NULL,(LPCWSTR)lpwPriv,&TokenPrivileges.Privileges[0].Luid)) {
50 | CloseHandle(hToken);
51 | return FALSE;
52 | }
53 |
54 | if(!AdjustTokenPrivileges(hToken,FALSE,&TokenPrivileges,sizeof(TOKEN_PRIVILEGES),NULL,NULL)) {
55 | CloseHandle(hToken);
56 | return FALSE;
57 | }
58 |
59 | CloseHandle(hToken);
60 | return TRUE;
61 | }
62 |
63 | };
64 |
65 | inline std::wstring string_to_wstring(const std::string& str) {
66 | std::wstring wsTmp(str.begin(), str.end());
67 | return wsTmp;
68 | }
69 |
70 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/thread_builder.hpp:
--------------------------------------------------------------------------------
1 | #pragma once
2 |
3 | #include
4 | #include
5 | #include
6 | #include "phnt.h"
7 |
8 | #include "calltrace_builder.hpp"
9 | #include "thread.hpp"
10 |
11 | namespace hsb::containers {
12 | class thread_builder{
13 | public:
14 | thread_builder();
15 | std::unique_ptr build(HANDLE, PSYSTEM_THREAD_INFORMATION);
16 | private:
17 | calltrace_builder calltrace_builder_;
18 | uint64_t enum_stackbase(HANDLE,HANDLE);
19 | };
20 |
21 | //implementation
22 | //================================================================================================
23 |
24 | #pragma region constructor and destructor
25 | thread_builder::thread_builder()
26 | {}
27 | #pragma endregion
28 |
29 | #pragma region public methods
30 | std::unique_ptr thread_builder::build(HANDLE h_process, PSYSTEM_THREAD_INFORMATION sti){
31 |
32 | auto t = std::make_unique();
33 | #pragma warning(suppress: 4302)
34 | #pragma warning(suppress: 4311)
35 | t->tid = (DWORD)sti->ClientId.UniqueThread;
36 | t->handle = OpenThread(THREAD_ALL_ACCESS,FALSE,t->tid);
37 | if(t->handle == nullptr)
38 | return nullptr;
39 |
40 | t->stackbase = enum_stackbase(h_process,t->handle);
41 | t->calltrace = calltrace_builder_.build(h_process, t->handle);
42 | if (t->calltrace == nullptr)
43 | return nullptr;
44 |
45 | return t;
46 |
47 | }
48 |
49 | uint64_t thread_builder::enum_stackbase(HANDLE h_process,HANDLE h_thread) {
50 |
51 | TEB teb ={0};
52 | PTEB pTeb = NULL;
53 | THREAD_BASIC_INFORMATION tbi ={0};
54 | SIZE_T sRead = 0;
55 |
56 | NTSTATUS status = STATUS_SUCCESS;
57 | BOOL bSuccess = FALSE;
58 |
59 | status = NtQueryInformationThread(h_thread,(THREADINFOCLASS)ThreadBasicInformation,&tbi,sizeof(THREAD_BASIC_INFORMATION),NULL);
60 | if(status == STATUS_SUCCESS) {
61 |
62 | bSuccess = ReadProcessMemory(h_process,tbi.TebBaseAddress,&teb,sizeof(TEB),&sRead);
63 | if(bSuccess == FALSE)
64 | return 0;
65 |
66 | return (uint64_t)teb.NtTib.StackBase;
67 |
68 | }
69 |
70 | return 0;
71 |
72 | }
73 |
74 |
75 | #pragma endregion
76 |
77 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/stomped_module.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include
4 |
5 | #include "phnt.h"
6 | #include "misc.hpp"
7 | #include "scans.hpp"
8 |
9 | namespace hsb::scanning {
10 |
11 | static BOOL ModuleIsSharedOriginal(HANDLE, PVOID);
12 |
13 | thread_scan thread_scans::stomped_module = [](process* process, thread* thread) {
14 |
15 | BOOL bSuspicious = FALSE, bSuccess = FALSE;
16 | std::wstring message;
17 | std::string moduleName;
18 |
19 | for (int i = 0; i < thread->calltrace->raw_addresses.size(); i++) {
20 |
21 | DWORD64 savedRet = thread->calltrace->raw_addresses.at(i);
22 | moduleName = thread->calltrace->modules.at(i);
23 |
24 | if (ModuleIsSharedOriginal(process->handle, (PVOID)savedRet) || moduleName == "unknown")
25 | continue;
26 |
27 | if (!_strcmpi(moduleName.c_str(), "ntdll") ||
28 | !_strcmpi(moduleName.c_str(), "kernel32") ||
29 | !_strcmpi(moduleName.c_str(), "kernelbase") ||
30 | !_strcmpi(moduleName.c_str(), "user32") ||
31 | !_strcmpi(moduleName.c_str(), "win32u")
32 | )
33 | continue;
34 |
35 | bSuspicious = TRUE;
36 | break;
37 |
38 | }
39 |
40 | if (bSuspicious) {
41 |
42 | thread_detection detection;
43 | detection.name = L"Module Stomping";
44 | detection.description = std::format(L"Callstack contains stomped module: {}", misc::string_to_wstring(moduleName));
45 | detection.tid = thread->tid;
46 | detection.severity = hsb::containers::detections::LOW;
47 | process->add_detection(detection);
48 |
49 | }
50 |
51 | };
52 |
53 | static BOOL ModuleIsSharedOriginal(HANDLE hProcess, PVOID pAddr) {
54 |
55 | BOOL bIsSharedOrig = TRUE;
56 | SIZE_T s = 0;
57 |
58 | PMEMORY_WORKING_SET_EX_INFORMATION workingSets = { 0 };
59 | MEMORY_BASIC_INFORMATION mbi = { 0 };
60 | MEMORY_WORKING_SET_EX_INFORMATION mwsi = { 0 };
61 |
62 | s = VirtualQueryEx(hProcess, (LPCVOID)pAddr, &mbi, sizeof(MEMORY_BASIC_INFORMATION));
63 | if (s == 0)
64 | goto Cleanup;
65 |
66 | if (mbi.Type != MEM_IMAGE)
67 | goto Cleanup;
68 |
69 | mwsi.VirtualAddress = mbi.BaseAddress;
70 | if (NtQueryVirtualMemory(hProcess, NULL, MemoryWorkingSetExInformation, &mwsi, sizeof(MEMORY_WORKING_SET_EX_INFORMATION), 0) != STATUS_SUCCESS)
71 | goto Cleanup;
72 |
73 | if (mwsi.u1.VirtualAttributes.SharedOriginal == 0)
74 | bIsSharedOrig = FALSE;
75 |
76 | Cleanup:
77 |
78 | return bIsSharedOrig;
79 |
80 | }
81 |
82 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/main.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 |
4 | #include "logger.hpp"
5 | #include "process_enumerator.hpp"
6 | #include "process_scanner.hpp"
7 | #include "misc.hpp"
8 |
9 | using logger = hsb::logger::logger;
10 | using process = hsb::containers::process;
11 | using process_scanner = hsb::scanning::process_scanner;
12 | using process_enumerator = hsb::containers::process_enumerator;
13 | using token_helpers = hsb::misc::token_helpers;
14 |
15 | void parse_args(int argc, char** argv, bool*, bool*, uint16_t*);
16 |
17 | int main (int argc,char** argv)
18 | {
19 |
20 | bool print_cmdline = false, ignore_dotnet = true;
21 | uint16_t scan_pid = 0;
22 | parse_args(argc, argv, &print_cmdline, &ignore_dotnet, &scan_pid);
23 |
24 | process_enumerator process_enumerator(scan_pid, ignore_dotnet);
25 | process_scanner process_scanner;
26 | std::vector> processes;
27 | std::pair scan_stats;
28 |
29 | if (token_helpers::is_elevated() == FALSE) {
30 | std::wcout << L"! Not elevated" << std::endl;
31 | return 0;
32 | }
33 |
34 | if (token_helpers::set_debug_privilege() == FALSE) {
35 | std::wcout << L"! Failed to enable debug privilege" << std::endl;
36 | return 0;
37 | }
38 |
39 | logger::init(print_cmdline);
40 | logger::logo();
41 |
42 | auto t1 = std::chrono::high_resolution_clock::now();
43 |
44 | processes = process_enumerator.enumerate_processes();
45 | scan_stats = process_scanner.scan_processes(processes);
46 |
47 | for(auto& process : processes){
48 | if(process->detections.size()){
49 | logger::print_suspicious_process(process.get());
50 | }
51 | }
52 |
53 | auto t2 = std::chrono::high_resolution_clock::now();
54 | auto ms_int = duration_cast(t2 - t1);
55 |
56 | logger::print_stats(scan_stats, ms_int.count());
57 |
58 | return 0;
59 |
60 | }
61 |
62 | void parse_args(int argc, char** argv, bool* print_cmdline, bool* ignore_dotnet, uint16_t* pid)
63 | {
64 | for (int i = 1; i < argc; i++) {
65 |
66 | if (!_strcmpi(argv[i], "-p") || !strcmp(argv[i], "--pid")) {
67 | *pid = atoi(argv[i + 1]);
68 | i++;
69 | }
70 | else if (!_strcmpi(argv[i], "--dotnet"))
71 | *ignore_dotnet = false;
72 | else if (!_strcmpi(argv[i], "--commandline"))
73 | *print_cmdline = true;
74 | else if (!_strcmpi(argv[i], "-h") || !strcmp(argv[i], "--help"))
75 | logger::help();
76 | else
77 | logger::help();
78 |
79 | }
80 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/return_address_spoofing.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include
4 |
5 | #include "phnt.h"
6 | #include "misc.hpp"
7 | #include "scans.hpp"
8 |
9 | namespace hsb::scanning {
10 |
11 |
12 | thread_scan thread_scans::return_address_spoofing = [](process* process, thread* thread) {
13 |
14 |
15 | BOOL bSuspicious = FALSE, bSuccess = FALSE;
16 | SIZE_T nRead = 0, s = 0;
17 |
18 | BYTE instructions[4] = { 0x00 };
19 |
20 | BYTE patternJmpDerefRbx[2] = { 0xFF, 0x23 };
21 | BYTE patternJmpDerefRbp[3] = { 0xFF, 0x65, 0x00 };
22 | BYTE patternJmpDerefRdi[2] = { 0xFF, 0x27 };
23 | BYTE patternJmpDerefRsi[2] = { 0xFF, 0x26 };
24 | BYTE patternJmpDerefR12[4] = { 0x41, 0xff, 0x24, 0x24 };
25 | BYTE patternJmpDerefR13[4] = { 0x41, 0xff, 0x65, 0x00 };
26 | BYTE patternJmpDerefR14[3] = { 0x41, 0xff, 0x26 };
27 | BYTE patternJmpDerefR15[3] = { 0x41, 0xff, 0x27 };
28 |
29 | for (int i = 0; i < thread->calltrace->raw_addresses.size(); i++) {
30 |
31 | bSuccess = ReadProcessMemory(process->handle, (PVOID)thread->calltrace->raw_addresses.at(i), instructions, sizeof(instructions), &nRead);
32 | if (bSuccess == FALSE)
33 | goto Cleanup;
34 |
35 | if (memcmp(instructions, patternJmpDerefRbx, sizeof(patternJmpDerefRbx)) == 0)
36 | bSuspicious = TRUE;
37 | else if (memcmp(instructions, patternJmpDerefRbp, sizeof(patternJmpDerefRbp)) == 0)
38 | bSuspicious = TRUE;
39 | else if (memcmp(instructions, patternJmpDerefRdi, sizeof(patternJmpDerefRdi)) == 0)
40 | bSuspicious = TRUE;
41 | else if (memcmp(instructions, patternJmpDerefRsi, sizeof(patternJmpDerefRsi)) == 0)
42 | bSuspicious = TRUE;
43 | else if (memcmp(instructions, patternJmpDerefR12, sizeof(patternJmpDerefR12)) == 0)
44 | bSuspicious = TRUE;
45 | else if (memcmp(instructions, patternJmpDerefR13, sizeof(patternJmpDerefR13)) == 0)
46 | bSuspicious = TRUE;
47 | else if (memcmp(instructions, patternJmpDerefR14, sizeof(patternJmpDerefR14)) == 0)
48 | bSuspicious = TRUE;
49 | else if (memcmp(instructions, patternJmpDerefR15, sizeof(patternJmpDerefR15)) == 0)
50 | bSuspicious = TRUE;
51 |
52 | if (bSuspicious) {
53 |
54 | thread_detection detection;
55 | detection.name = L"Return Address Spoofing";
56 | detection.description = std::format(L"Thread {} returns to JMP gadget. Gadget in: {}", thread->tid, misc::string_to_wstring(thread->calltrace->syms.at(i)));
57 | detection.tid = thread->tid;
58 | detection.severity = hsb::containers::detections::CRITICAL;
59 |
60 | process->add_detection(detection);
61 |
62 | break;
63 |
64 | }
65 |
66 | }
67 | Cleanup:
68 |
69 | return;
70 |
71 | };
72 |
73 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/logger.hpp:
--------------------------------------------------------------------------------
1 | #include "phnt.h"
2 | #include
3 | #include
4 |
5 | #include "detection.hpp"
6 | #include "process.hpp"
7 |
8 | namespace hsb::logger {
9 |
10 | class logger {
11 |
12 | using detection = hsb::containers::detections::detection;
13 | using process = hsb::containers::process;
14 |
15 | private:
16 |
17 | logger() = delete;
18 | static inline bool cmdline_;
19 |
20 | public:
21 |
22 | static void init(bool);
23 | static void logo(void);
24 | static void help(void);
25 | static void print_suspicious_process(process*);
26 | static void print_stats(std::pair, long long);
27 |
28 | };
29 |
30 | //implementation
31 | //================================================================================================
32 |
33 | #pragma region public methods
34 |
35 | // Copy paste from https://cboard.cprogramming.com/cplusplus-programming/181215-printing-colored-text-code-blocks-cplusplus.html
36 | void logger::init(bool cmdline)
37 | {
38 | HANDLE h;
39 | DWORD mode;
40 |
41 | h = GetStdHandle(STD_OUTPUT_HANDLE);
42 | GetConsoleMode(h, &mode);
43 | SetConsoleMode(h, mode | ENABLE_VIRTUAL_TERMINAL_PROCESSING);
44 |
45 | cmdline_ = cmdline;
46 |
47 | }
48 |
49 | void logger::logo(void)
50 | {
51 | std::wcout <<
52 | L" _ _ _____ ______\r\n"
53 | L"| | | | / ___| | ___ \\\r\n"
54 | L"| |_| | \\ `--. | |_/ /\r\n"
55 | L"| _ | `--. \\ | ___ \\\r\n"
56 | L"| | | | /\\__/ / | |_/ /\r\n"
57 | L"\\_| |_/ \\____/ \\____/\r\n"
58 | L"\r\n"
59 | L"Hunt-Sleeping-Beacons | @thefLinkk\r\n"
60 | << std::endl;
61 | }
62 |
63 | void logger::help(void)
64 | {
65 |
66 | std::wcout << std::endl;
67 | std::wcout << L"-p / --pid {PID}" << std::endl;;
68 | std::wcout << std::endl;
69 | std::wcout << L"--dotnet | Set to also include dotnet processes. ( Prone to false positivies )" << std::endl;
70 | std::wcout << L"--commandline | Enables output of cmdline for suspicious processes" << std::endl;
71 | std::wcout << L"-h / --help | Prints this message?" << std::endl;
72 | std::wcout << std::endl;
73 |
74 | exit(0);
75 |
76 | }
77 |
78 | void logger::print_suspicious_process(process* process)
79 | {
80 |
81 | std::wcout << std::format(L"\033[36m* Detections for: {} ({}) {}\033[0m", process->imagename, process->pid, (cmdline_ ? process->cmdline : L"")) << std::endl;
82 | for (std::unique_ptr& detection : process->detections) {
83 | std::wcout << "\t" << detection->to_string() << std::endl;
84 | }
85 |
86 | }
87 |
88 | void logger::print_stats(std::pair stats, long long time)
89 | {
90 | std::wcout << std::endl << std::format(L"* Scanned: {} processes and {} threads in {} seconds", stats.first, stats.second, (double)time / 1000) << std::endl;
91 | }
92 | #pragma endregion
93 |
94 | };
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/process_scanner.hpp:
--------------------------------------------------------------------------------
1 | #pragma once
2 |
3 | #include
4 |
5 | #include "BS_thread_pool.hpp"
6 |
7 | #include "scans.hpp"
8 | #include "process.hpp"
9 |
10 | namespace hsb::scanning {
11 |
12 | class process_scanner{
13 |
14 | using process_scan = hsb::scanning::process_scan;
15 | using process = hsb::containers::process;
16 | using multi_future = BS::multi_future;
17 | using thread_pool = BS::thread_pool;
18 |
19 | public:
20 |
21 | process_scanner();
22 | ~process_scanner();
23 |
24 | std::pair scan_processes(std::vector>&);
25 |
26 | private:
27 | thread_pool thread_pool_;
28 | multi_future multi_future_;
29 | std::vector process_scans_;
30 | std::vector thread_scans_;
31 | std::atomic n_scanned_processes_;
32 | std::atomic n_scanned_threads_;
33 | };
34 |
35 | //implementation
36 | //================================================================================================
37 |
38 | #pragma region constructor and destructor
39 | process_scanner::process_scanner()
40 | : n_scanned_processes_(0)
41 | , n_scanned_threads_(0)
42 | {
43 |
44 | process_scans_.push_back(process_scans::suspicious_timer);
45 |
46 | thread_scans_.push_back(thread_scans::blocking_apc);
47 | thread_scans_.push_back(thread_scans::blocking_timer);
48 | thread_scans_.push_back(thread_scans::abnormal_intermodular_call);
49 | thread_scans_.push_back(thread_scans::return_address_spoofing);
50 | thread_scans_.push_back(thread_scans::private_memory);
51 | thread_scans_.push_back(thread_scans::stomped_module);
52 | thread_scans_.push_back(thread_scans::hardware_breakpoints);
53 | thread_scans_.push_back(thread_scans::non_executable_memory);
54 |
55 | }
56 | process_scanner::~process_scanner() {}
57 | #pragma endregion
58 |
59 | #pragma region public methods
60 |
61 | std::pair process_scanner::scan_processes(std::vector>& processes)
62 | {
63 |
64 | for(auto& process : processes){
65 |
66 | n_scanned_processes_.fetch_add(1, std::memory_order::relaxed);
67 |
68 | for(auto& process_scan : process_scans_){
69 | multi_future_.push_back(thread_pool_.submit_task([this, &process_scan, &process]() {
70 | process_scan(process.get());
71 | }));
72 | }
73 |
74 | for(auto& thread : process->threads){
75 | for(auto& thread_scan : thread_scans_){
76 | multi_future_.push_back(thread_pool_.submit_task([this, &thread, &thread_scan,&process]() {
77 | thread_scan(process.get(), thread.get());
78 | }));
79 | }
80 |
81 | n_scanned_threads_.fetch_add(1, std::memory_order::relaxed);
82 |
83 | }
84 | }
85 |
86 | multi_future_.get();
87 |
88 | return std::make_pair(n_scanned_processes_.load(std::memory_order::relaxed), n_scanned_threads_.load(std::memory_order::relaxed));
89 |
90 | }
91 |
92 | #pragma endregion
93 |
94 | }
95 |
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons.sln:
--------------------------------------------------------------------------------
1 |
2 | Microsoft Visual Studio Solution File, Format Version 12.00
3 | # Visual Studio Version 17
4 | VisualStudioVersion = 17.11.35327.3
5 | MinimumVisualStudioVersion = 10.0.40219.1
6 | Project("{8BC9CEB8-8B4A-11D0-8D11-00A0C91BC942}") = "Hunt-Sleeping-Beacons", "Hunt-Sleeping-Beacons\Hunt-Sleeping-Beacons.vcxproj", "{70A45DBF-68FE-4D6B-864E-D5E239340047}"
7 | EndProject
8 | Project("{8BC9CEB8-8B4A-11D0-8D11-00A0C91BC942}") = "UnitTests", "UnitTests\UnitTests.vcxproj", "{47D29CEF-58D6-418E-AC77-397E90D40748}"
9 | EndProject
10 | Global
11 | GlobalSection(SolutionConfigurationPlatforms) = preSolution
12 | Debug|ARM64 = Debug|ARM64
13 | Debug|x64 = Debug|x64
14 | Debug|x86 = Debug|x86
15 | Release|ARM64 = Release|ARM64
16 | Release|x64 = Release|x64
17 | Release|x86 = Release|x86
18 | EndGlobalSection
19 | GlobalSection(ProjectConfigurationPlatforms) = postSolution
20 | {70A45DBF-68FE-4D6B-864E-D5E239340047}.Debug|ARM64.ActiveCfg = Debug|ARM64
21 | {70A45DBF-68FE-4D6B-864E-D5E239340047}.Debug|ARM64.Build.0 = Debug|ARM64
22 | {70A45DBF-68FE-4D6B-864E-D5E239340047}.Debug|x64.ActiveCfg = Debug|x64
23 | {70A45DBF-68FE-4D6B-864E-D5E239340047}.Debug|x64.Build.0 = Debug|x64
24 | {70A45DBF-68FE-4D6B-864E-D5E239340047}.Debug|x86.ActiveCfg = Debug|Win32
25 | {70A45DBF-68FE-4D6B-864E-D5E239340047}.Debug|x86.Build.0 = Debug|Win32
26 | {70A45DBF-68FE-4D6B-864E-D5E239340047}.Release|ARM64.ActiveCfg = Release|ARM64
27 | {70A45DBF-68FE-4D6B-864E-D5E239340047}.Release|ARM64.Build.0 = Release|ARM64
28 | {70A45DBF-68FE-4D6B-864E-D5E239340047}.Release|x64.ActiveCfg = Release|x64
29 | {70A45DBF-68FE-4D6B-864E-D5E239340047}.Release|x64.Build.0 = Release|x64
30 | {70A45DBF-68FE-4D6B-864E-D5E239340047}.Release|x86.ActiveCfg = Release|Win32
31 | {70A45DBF-68FE-4D6B-864E-D5E239340047}.Release|x86.Build.0 = Release|Win32
32 | {47D29CEF-58D6-418E-AC77-397E90D40748}.Debug|ARM64.ActiveCfg = Debug|ARM64
33 | {47D29CEF-58D6-418E-AC77-397E90D40748}.Debug|ARM64.Build.0 = Debug|ARM64
34 | {47D29CEF-58D6-418E-AC77-397E90D40748}.Debug|x64.ActiveCfg = Debug|x64
35 | {47D29CEF-58D6-418E-AC77-397E90D40748}.Debug|x64.Build.0 = Debug|x64
36 | {47D29CEF-58D6-418E-AC77-397E90D40748}.Debug|x86.ActiveCfg = Debug|Win32
37 | {47D29CEF-58D6-418E-AC77-397E90D40748}.Debug|x86.Build.0 = Debug|Win32
38 | {47D29CEF-58D6-418E-AC77-397E90D40748}.Release|ARM64.ActiveCfg = Release|ARM64
39 | {47D29CEF-58D6-418E-AC77-397E90D40748}.Release|ARM64.Build.0 = Release|ARM64
40 | {47D29CEF-58D6-418E-AC77-397E90D40748}.Release|x64.ActiveCfg = Release|x64
41 | {47D29CEF-58D6-418E-AC77-397E90D40748}.Release|x64.Build.0 = Release|x64
42 | {47D29CEF-58D6-418E-AC77-397E90D40748}.Release|x86.ActiveCfg = Release|Win32
43 | {47D29CEF-58D6-418E-AC77-397E90D40748}.Release|x86.Build.0 = Release|Win32
44 | EndGlobalSection
45 | GlobalSection(SolutionProperties) = preSolution
46 | HideSolutionNode = FALSE
47 | EndGlobalSection
48 | GlobalSection(ExtensibilityGlobals) = postSolution
49 | SolutionGuid = {2F0C055E-B04E-4871-A33F-628A61E058D3}
50 | EndGlobalSection
51 | EndGlobal
52 |
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/process_enumerator.hpp:
--------------------------------------------------------------------------------
1 | #pragma once
2 |
3 | #include
4 | #include
5 |
6 | #include "phnt.h"
7 |
8 | #include "process.hpp"
9 | #include "process_builder.hpp"
10 |
11 | namespace hsb::containers {
12 |
13 | using process = hsb::containers::process;
14 | using process_builder = hsb::containers::process_builder;
15 |
16 | class process_enumerator{
17 | public:
18 |
19 | process_enumerator(uint16_t pid_to_scan = 0, bool ignore_dotnet = true);
20 | ~process_enumerator();
21 | std::vector> enumerate_processes();
22 |
23 | private:
24 | process_builder process_builder_;
25 | uint16_t pid_to_scan_;
26 | };
27 |
28 | //implementation
29 | //================================================================================================
30 |
31 | #pragma region constructor and destructor
32 | process_enumerator::process_enumerator(uint16_t pid_to_scan, bool ignore_dotnet)
33 | : pid_to_scan_(pid_to_scan)
34 | , process_builder_(ignore_dotnet)
35 | {
36 | SymSetOptions(SYMOPT_UNDNAME | SYMOPT_DEFERRED_LOADS);
37 | }
38 |
39 | process_enumerator::~process_enumerator()
40 | {}
41 |
42 | #pragma endregion
43 |
44 | #pragma region public methods
45 |
46 | std::vector> process_enumerator::enumerate_processes(){
47 |
48 | NTSTATUS status = STATUS_UNSUCCESSFUL;
49 | PVOID pBuffer = NULL;
50 | ULONG uBufferSize = 0;
51 | BOOL bSuccess = FALSE;
52 |
53 | PSYSTEM_PROCESS_INFORMATION pProcessInformation = NULL;
54 | SYSTEM_THREAD_INFORMATION thread_information = {0};
55 |
56 | std::vector> processes;
57 | std::unique_ptr tmp = nullptr;
58 |
59 | do {
60 |
61 | status = NtQuerySystemInformation((SYSTEM_INFORMATION_CLASS)SystemProcessInformation,pBuffer,uBufferSize,&uBufferSize);
62 | if(!NT_SUCCESS(status)) {
63 |
64 | if(status == STATUS_INFO_LENGTH_MISMATCH) {
65 | if(pBuffer != NULL)
66 | LocalFree(pBuffer);
67 | pBuffer = LocalAlloc(LMEM_ZEROINIT,uBufferSize);
68 | if(pBuffer == NULL)
69 | goto Cleanup;
70 |
71 | continue;
72 | }
73 | break;
74 |
75 | }
76 | else {
77 |
78 | pProcessInformation = (PSYSTEM_PROCESS_INFORMATION)pBuffer;
79 | break;
80 |
81 | }
82 |
83 | } while(1);
84 |
85 | while(pProcessInformation && pProcessInformation->NextEntryOffset) {
86 |
87 | #pragma warning(suppress: 4302) // 'type cast' : truncation
88 | #pragma warning(suppress: 4311) // pointer truncation
89 | if ( pid_to_scan_ && pid_to_scan_ != (uint16_t)pProcessInformation->UniqueProcessId) {
90 | goto Next;
91 | }
92 |
93 | tmp = process_builder_.build(pProcessInformation);
94 | if(tmp)
95 | processes.push_back(std::move(tmp));
96 |
97 | Next:
98 | pProcessInformation = (PSYSTEM_PROCESS_INFORMATION)((LPBYTE)pProcessInformation + pProcessInformation->NextEntryOffset);
99 |
100 | }
101 |
102 |
103 | Cleanup:
104 |
105 | return processes;
106 |
107 | }
108 |
109 | #pragma endregion
110 |
111 |
112 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/calltrace_builder.hpp:
--------------------------------------------------------------------------------
1 | #pragma once
2 |
3 | #include "dbghelp.h"
4 | #include
5 | #include
6 | #include "tlhelp32.h"
7 | #include
8 |
9 | #include "phnt.h"
10 | #include "calltrace.hpp"
11 |
12 | namespace hsb::containers {
13 |
14 | class calltrace_builder{
15 | public:
16 | calltrace_builder();
17 | ~calltrace_builder();
18 | std::unique_ptr build(HANDLE, HANDLE);
19 | };
20 |
21 | //implementation
22 | //================================================================================================
23 |
24 | #pragma region constructor and destructor
25 | calltrace_builder::calltrace_builder() {}
26 | calltrace_builder::~calltrace_builder() {}
27 | #pragma endregion
28 |
29 | #pragma region public methods
30 | std::unique_ptr calltrace_builder::build(HANDLE h_process, HANDLE h_thread) {
31 |
32 | CONTEXT context ={0x00};
33 | STACKFRAME64 stackframe ={0x00};
34 | PIMAGEHLP_SYMBOL64 pSymbol = NULL;
35 | PIMAGEHLP_MODULE64 pModInfo = NULL;
36 | BOOL bSuccess = FALSE, bModuleFound = FALSE;
37 | char cSymName[256] ={0x00},line[512] ={0};
38 | DWORD64 dw64Displacement = 0x00;
39 |
40 | auto c = std::make_unique();
41 |
42 | context.ContextFlags = CONTEXT_ALL;
43 | bSuccess = GetThreadContext(h_thread,&context);
44 | if(bSuccess == FALSE)
45 | return nullptr;
46 |
47 | stackframe.AddrPC.Offset = context.Rip;
48 | stackframe.AddrPC.Mode = AddrModeFlat;
49 | stackframe.AddrStack.Offset = context.Rsp;
50 | stackframe.AddrStack.Mode = AddrModeFlat;
51 | stackframe.AddrFrame.Offset = context.Rbp;
52 | stackframe.AddrFrame.Mode = AddrModeFlat;
53 |
54 | pSymbol = (PIMAGEHLP_SYMBOL64)LocalAlloc(LMEM_ZEROINIT,sizeof(IMAGEHLP_SYMBOL64) + 256 * sizeof(WCHAR));
55 | if(pSymbol == NULL)
56 | return nullptr;
57 |
58 | pModInfo = (PIMAGEHLP_MODULE64)LocalAlloc(LMEM_ZEROINIT,sizeof(IMAGEHLP_MODULE64) + 256 * sizeof(WCHAR));
59 | if(pModInfo == NULL)
60 | return nullptr;
61 |
62 | pSymbol->SizeOfStruct = sizeof(IMAGEHLP_SYMBOL64);
63 | pSymbol->MaxNameLength = 255;
64 | pModInfo->SizeOfStruct = sizeof(IMAGEHLP_MODULE64);
65 |
66 | while(1) {
67 |
68 | bSuccess = StackWalk64(IMAGE_FILE_MACHINE_AMD64,h_process,h_thread,&stackframe,&context,NULL,NULL,NULL,NULL);
69 | if(bSuccess == FALSE) {
70 | break;
71 | }
72 |
73 | memset(line,0,512);
74 | memset(cSymName,0,256);
75 |
76 | c->raw_addresses.push_back(stackframe.AddrPC.Offset);
77 |
78 | bModuleFound = SymGetModuleInfo64(h_process,(ULONG64)stackframe.AddrPC.Offset,pModInfo);
79 | if(bModuleFound == FALSE) {
80 | c->modules.push_back("unknown");
81 | c->syms.push_back("unknown");
82 | }
83 | else {
84 |
85 | SymGetSymFromAddr64(h_process,(ULONG64)stackframe.AddrPC.Offset,&dw64Displacement,pSymbol);
86 | UnDecorateSymbolName(pSymbol->Name,cSymName,256,UNDNAME_COMPLETE);
87 |
88 | wsprintfA(line,"%s!%s",pModInfo->ModuleName,cSymName);
89 | c->syms.push_back(line);
90 | c->modules.push_back(pModInfo->ModuleName);
91 |
92 |
93 | }
94 |
95 | }
96 |
97 | if (pSymbol)
98 | LocalFree(pSymbol);
99 |
100 | if (pModInfo)
101 | LocalFree(pModInfo);
102 |
103 | return c;
104 |
105 | }
106 | #pragma endregion
107 |
108 |
109 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/blocking_timer.cpp:
--------------------------------------------------------------------------------
1 | #pragma once
2 |
3 | #include "phnt.h"
4 | #include "scans.hpp"
5 |
6 | namespace hsb::scanning {
7 |
8 | static DWORD64 callback_dispatcher = 0;
9 | static HANDLE hEvent = NULL;
10 | static inline std::once_flag dispatcher_resolved;
11 | static void my_callback(void);
12 | static bool resolve_callback_dispatcher(void);
13 |
14 | thread_scan thread_scans::blocking_timer = [](process* process,thread* thread) {
15 |
16 | BOOL bSuccess = FALSE;
17 | PVOID readAddr = NULL;
18 | DWORD64 stackAddr = 0;
19 | size_t numRead = 0;
20 |
21 | CONTEXT context = { 0 };
22 | context.ContextFlags = CONTEXT_ALL;
23 |
24 | std::call_once(dispatcher_resolved, resolve_callback_dispatcher);
25 | if (callback_dispatcher == 0) {
26 | return;
27 | }
28 |
29 | bSuccess = GetThreadContext(thread->handle, &context);
30 | if (bSuccess == FALSE)
31 | goto Cleanup;
32 |
33 | for (DWORD64 i = context.Rsp; i <= (DWORD64)thread->stackbase; i += 8) {
34 |
35 | bSuccess = ReadProcessMemory(process->handle, (LPCVOID)i, &stackAddr, sizeof(DWORD64), &numRead);
36 | if (bSuccess == FALSE)
37 | break;
38 |
39 | if (stackAddr == callback_dispatcher) {
40 |
41 | thread_detection detection;
42 | detection.name = L"Blocking Timer detected";
43 | detection.description = L"Thread's blocking state triggered by ntdll!RtlpTpTimerCallback";
44 | detection.tid = thread->tid;
45 | detection.severity = hsb::containers::detections::HIGH;
46 |
47 | process->add_detection(detection);
48 |
49 | }
50 |
51 | }
52 |
53 | Cleanup:
54 |
55 | return;
56 |
57 | };
58 |
59 | static void my_callback(void) {
60 |
61 | CONTEXT context = { 0 };
62 | STACKFRAME64 stackframe = { 0x00 };
63 |
64 | BOOL bSuccess = FALSE;
65 |
66 | RtlCaptureContext(&context);
67 |
68 | stackframe.AddrPC.Offset = context.Rip;
69 | stackframe.AddrPC.Mode = AddrModeFlat;
70 | stackframe.AddrStack.Offset = context.Rsp;
71 | stackframe.AddrStack.Mode = AddrModeFlat;
72 | stackframe.AddrFrame.Offset = context.Rbp;
73 | stackframe.AddrFrame.Mode = AddrModeFlat;
74 |
75 | SymInitialize(GetCurrentProcess(), NULL, TRUE);
76 |
77 | bSuccess = StackWalk64(IMAGE_FILE_MACHINE_AMD64, GetCurrentProcess(), GetCurrentThread(), &stackframe, &context, NULL, SymFunctionTableAccess64, SymGetModuleBase64, NULL);
78 | if (bSuccess == FALSE)
79 | return;
80 |
81 | bSuccess = StackWalk64(IMAGE_FILE_MACHINE_AMD64, GetCurrentProcess(), GetCurrentThread(), &stackframe, &context, NULL, SymFunctionTableAccess64, SymGetModuleBase64, NULL);
82 | if (bSuccess == FALSE)
83 | return;
84 |
85 | SymCleanup(GetCurrentProcess());
86 |
87 | callback_dispatcher = stackframe.AddrPC.Offset;
88 |
89 | SetEvent(hEvent);
90 |
91 |
92 | }
93 |
94 | static bool resolve_callback_dispatcher(void) {
95 |
96 | BOOL bSuccess = FALSE;
97 | PVOID retDispatcher = NULL;
98 |
99 | HANDLE hNewTimer = NULL, hTimerQueue = NULL;
100 | HMODULE hNtdll = NULL;
101 |
102 | hEvent = CreateEventW(0, 0, 0, 0);
103 | if (hEvent == NULL)
104 | return FALSE;
105 |
106 | hTimerQueue = CreateTimerQueue();
107 | if (hTimerQueue == NULL)
108 | return FALSE;
109 |
110 | CreateTimerQueueTimer(&hNewTimer, hTimerQueue, (WAITORTIMERCALLBACK)(my_callback), NULL, 0, 0, WT_EXECUTEINTIMERTHREAD);
111 | WaitForSingleObject(hEvent, INFINITE);
112 |
113 | hNtdll = GetModuleHandleA("ntdll.dll");
114 | if (hNtdll == NULL)
115 | goto exit;
116 |
117 | bSuccess = TRUE;
118 |
119 | exit:
120 |
121 | //if ( hNewTimer )
122 | // CloseHandle ( hNewTimer );
123 |
124 | if (hEvent)
125 | CloseHandle(hEvent);
126 |
127 | //if ( hTimerQueue )
128 | // CloseHandle ( hTimerQueue );
129 |
130 | return bSuccess;
131 |
132 | }
133 |
134 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/Hunt-Sleeping-Beacons.vcxproj.filters:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 | {4FC737F1-C7A5-4376-A066-2A32D752A2FF}
6 | cpp;c;cc;cxx;c++;cppm;ixx;def;odl;idl;hpj;bat;asm;asmx
7 |
8 |
9 | {93995380-89BD-4b04-88EB-625FBE52EBFB}
10 | h;hh;hpp;hxx;h++;hm;inl;inc;ipp;xsd
11 |
12 |
13 | {67DA6AB6-F800-4c08-8B7A-83BB121AAD01}
14 | rc;ico;cur;bmp;dlg;rc2;rct;bin;rgs;gif;jpg;jpeg;jpe;resx;tiff;tif;png;wav;mfcribbon-ms
15 |
16 |
17 | {ae639502-6954-4e94-8c7c-a0ce533a0be4}
18 |
19 |
20 | {a26944a4-a039-4ad4-b810-a8447703e1da}
21 |
22 |
23 | {94477e87-8215-483a-bd61-893c5c5c8376}
24 |
25 |
26 | {b84b4892-11b9-4bb8-bfab-109e82403ede}
27 |
28 |
29 |
30 |
31 | Source Files
32 |
33 |
34 | Header Files\scanning\scans
35 |
36 |
37 | Header Files\scanning\scans
38 |
39 |
40 | Header Files\scanning\scans
41 |
42 |
43 | Header Files\scanning\scans
44 |
45 |
46 | Header Files\scanning\scans
47 |
48 |
49 | Header Files\scanning\scans
50 |
51 |
52 | Header Files\scanning\scans
53 |
54 |
55 | Header Files\scanning\scans
56 |
57 |
58 | Header Files\scanning\scans
59 |
60 |
61 |
62 |
63 | Header Files\container
64 |
65 |
66 | Header Files\container
67 |
68 |
69 | Header Files\container
70 |
71 |
72 | Header Files\container
73 |
74 |
75 | Header Files\container
76 |
77 |
78 | Header Files\container
79 |
80 |
81 | Header Files\container
82 |
83 |
84 | Header Files\container
85 |
86 |
87 | Header Files\scanning
88 |
89 |
90 | Header Files\scanning
91 |
92 |
93 | Header Files\misc
94 |
95 |
96 | Header Files\misc
97 |
98 |
99 |
--------------------------------------------------------------------------------
/Readme.md:
--------------------------------------------------------------------------------
1 | # Hunt-Sleeping-Beacons
2 |
3 | This project is ( mostly ) a callstack scanner which tries to identify IOCs indicating an unpacked or injected C2 agent.
4 |
5 | All checks are based on the observation that C2 agents wait between their callbacks causing the beacons thread to idle and this tool aims to analyze what potentially caused the thread to idle.
6 |
7 | This includes traditional IOCs, such as unbacked memory or stomped modules, but also attempts to detect multiple implementation of sleepmasks using APCs or Timers. The latter is done by both: analyzing the callstack but also **enumerating timers and their exact callbacks from userland**.
8 |
9 | (Almost) none of those IOCs can be considered a 100% true positive, the module stomping detection e.g. is very prone to false positives. Yet, the results might raise suspicion about the behaviour of a process.
10 |
11 | DotNet and 32Bit binaries are ignored.
12 |
13 | 
14 |
15 | ## Checks
16 |
17 | ### Unbacked Memory
18 |
19 | A private r(w)x page in a callstack might indicate a beacon which was unpacked or injected at runtime.
20 |
21 | ### Non-Executable Memory
22 |
23 | Multiple Sleepmasks change the page permissions of the beacon's page to non-executable. This leads to a suspicious non-executable page in the callstack.
24 |
25 | ### Module Stomping
26 |
27 | Often, beacons avoid private memory pages by loading and overwriting a legitimate module from disk.
28 | Thanks to the ``copy on write`` mechanism, manipulated images can be identified by checking the field ``VirtualAttributes.SharedOriginal`` of ``MEMORY_WORKING_SET_EX_INFORMATION``. If any page in the callstack is not private and ``SharedOriginal == 0``, it is considered an IOC.
29 |
30 | This is probably the detection the most prone to false positives. :'(
31 |
32 | ### Suspicious APC
33 |
34 | Multiple implementations of sleepmasks queue a series of APCs to ``Ntdll!NtContinue`` one of which triggers the execution of ``Ntdll!WaitForSingleObject``. Thus, if ``Ntdll!KiUserApcDispatcher`` can be found on the callstack to a blocking function, this tool considers it an IOC.
35 |
36 | ### Suspicious Timers
37 |
38 | Similar to the suspicious usage of APCs, this tool also checks for ``ntdll!RtlpTpTimerCallback`` on the callstack to a blocking function to detect timer-based sleepmasks.
39 | ### Enumerating Timers and Callbacks
40 |
41 | To my understanding, Timers are implemented on top of ThreadPools. As [Alon Leviev has demonstrated](https://github.com/SafeBreach-Labs/PoolParty) those can be enumerated using ``NtQueryInformationWorkerFactory`` with ``WorkerFactoryBasicInformation``.
42 |
43 | The ``WORKER_FACTORY_BASIC_INFORMATION`` struct embeds a ``FULL_TP_POOL`` which in turn links to a ``TimerQueue`` double linked list. Traversing that list of ``PFULL_TP_TIMER`` allows accessing each registered callback. If any callback is found pointing to a set of suspicious api calls, such as ``ntdll!ntcontinue``, it can be considered a strong IOC.
44 |
45 | 
46 |
47 | ### Abnormal Intermodular Calls ( Module Proxying )
48 |
49 | Originally module proxying was introduced as a method to [bypass suspicious callstacks](https://0xdarkvortex.dev/proxying-dll-loads-for-hiding-etwti-stack-tracing/).
50 | While the bypass works, it introduces an other strong IOC, as the NTAPI is used to call the WINAPI. This is odd, as WINAPI is an abstraction for NTAPI. Thus, if a callstack is observed in which a sequence of ntdll.dll->kernel32.dll->ntdll.dll is found ending up calling a blocking function it can be considered an IOC.
51 |
52 | ### Return Address Spoofing
53 |
54 | Most Returnaddress spoofing implementations I am aware of make use of a technique in which the called function returns to a ``jmp [Nonvolatile-Register]`` gadget. This project simply iterates every return address in callstacks and searches for patterns indicating the return to a jmp gadget.
55 |
56 | 
57 |
58 | # Usage
59 |
60 | ```
61 | _ _ _____ ______
62 | | | | | / ___| | ___ \
63 | | |_| | \ `--. | |_/ /
64 | | _ | `--. \ | ___ \
65 | | | | | /\__/ / | |_/ /
66 | \_| |_/ \____/ \____/
67 |
68 | Hunt-Sleeping-Beacons | @thefLinkk
69 |
70 | -p / --pid {PID}
71 |
72 | --dotnet | Set to also include dotnet processes. ( Prone to false positivies )
73 | --commandline | Enables output of cmdline for suspicious processes
74 | -h / --help | Prints this message?
75 | ```
76 |
77 | # Credits
78 |
79 | - https://urien.gitbook.io/diago-lima/a-deep-dive-into-exploiting-windows-thread-pools/attacking-timer-queues
80 | - https://github.com/mrexodia/phnt-single-header
81 | - https://github.com/SafeBreach-Labs/PoolParty
82 | - https://github.com/bshoshany/thread-pool
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/process_builder.hpp:
--------------------------------------------------------------------------------
1 | #pragma once
2 |
3 | #include "phnt.h"
4 |
5 | #include
6 | #include
7 | #include
8 | #include "tlhelp32.h"
9 | #include "shlwapi.h"
10 | #include "strsafe.h"
11 |
12 | #include "process.hpp"
13 | #include "thread_builder.hpp"
14 |
15 | namespace hsb::containers {
16 | class process_builder{
17 | public:
18 | process_builder(bool ignore_dotnet = true);
19 | std::unique_ptr build(PSYSTEM_PROCESS_INFORMATION);
20 | private:
21 | std::vector> enumerate_threads(HANDLE, uint32_t, PSYSTEM_THREAD_INFORMATION);
22 | std::wstring enumerate_commandline(HANDLE);
23 | bool is_managed_process(DWORD);
24 | thread_builder thread_builder_;
25 | bool ignore_dotnet_;
26 |
27 | };
28 |
29 | //implementation
30 | //================================================================================================
31 |
32 | #pragma region constructor and destructor
33 | process_builder::process_builder(bool ignore_dotnet)
34 | : ignore_dotnet_(ignore_dotnet)
35 | {}
36 | #pragma endregion
37 |
38 | #pragma region public methods
39 | std::unique_ptr process_builder::build(PSYSTEM_PROCESS_INFORMATION spi){
40 |
41 | std::unique_ptr p = nullptr;
42 | HANDLE h_process = nullptr;
43 | BOOL bIs32Bit = false;
44 |
45 | p = std::make_unique();
46 |
47 | #pragma warning(suppress: 4302) // 'type cast' : truncation
48 | #pragma warning(suppress: 4311) // pointer truncation
49 | p->pid = (DWORD)spi->UniqueProcessId;
50 | p->handle = OpenProcess(PROCESS_ALL_ACCESS,FALSE,p->pid);
51 | if(p->handle == nullptr)
52 | return nullptr;
53 |
54 | IsWow64Process(p->handle, &bIs32Bit);
55 | if (bIs32Bit)
56 | return nullptr;
57 |
58 | if (ignore_dotnet_ && is_managed_process(p->pid))
59 | return nullptr;
60 |
61 | SymInitialize(p->handle, NULL, TRUE);
62 |
63 | p->imagename = std::wstring(spi->ImageName.Buffer);
64 | p->cmdline = enumerate_commandline(p->handle);
65 | p->threads = enumerate_threads(p->handle, spi->NumberOfThreads, spi->Threads);
66 |
67 | SymCleanup(p->handle);
68 |
69 | return p;
70 |
71 | }
72 | #pragma endregion
73 | #pragma region private methods
74 | std::vector> process_builder::enumerate_threads(HANDLE h_process, uint32_t n_threads, PSYSTEM_THREAD_INFORMATION sti) {
75 |
76 | std::vector> threads;
77 |
78 | for(uint32_t i = 0; i < n_threads; i++){
79 |
80 | auto t = thread_builder_.build(h_process, sti);
81 | if(t == nullptr)
82 | continue;
83 |
84 | threads.push_back(std::move(t));
85 |
86 | sti = (PSYSTEM_THREAD_INFORMATION)((PBYTE)sti + sizeof(SYSTEM_THREAD_INFORMATION));
87 |
88 | }
89 |
90 | return threads;
91 |
92 | }
93 |
94 | std::wstring process_builder::enumerate_commandline(HANDLE handle) {
95 |
96 | NTSTATUS status = STATUS_UNSUCCESSFUL;
97 | BOOL bSuccess = FALSE;
98 | ULONG uLen = 0;
99 | SIZE_T len = 0;
100 |
101 | PWSTR buf = NULL;
102 | PEB peb = {0};
103 | RTL_USER_PROCESS_PARAMETERS parameters = {0};
104 | PROCESS_BASIC_INFORMATION processInfo = {0};
105 |
106 | std::wstring cmdLine;
107 |
108 | status = NtQueryInformationProcess(handle,(PROCESSINFOCLASS)0,&processInfo,sizeof(PROCESS_BASIC_INFORMATION),&uLen);
109 | if(status != STATUS_SUCCESS)
110 | goto Cleanup;
111 |
112 | bSuccess = ReadProcessMemory(handle,processInfo.PebBaseAddress,&peb,sizeof(PEB),&len);
113 | if(bSuccess == FALSE)
114 | goto Cleanup;
115 |
116 | bSuccess = ReadProcessMemory(handle,peb.ProcessParameters,¶meters,sizeof(RTL_USER_PROCESS_PARAMETERS),&len);
117 | if(bSuccess == FALSE)
118 | goto Cleanup;
119 |
120 | buf = (PWSTR)HeapAlloc(GetProcessHeap(),HEAP_ZERO_MEMORY,parameters.CommandLine.Length * sizeof(WCHAR) + 2);
121 | if(buf == nullptr)
122 | goto Cleanup;
123 |
124 | if(parameters.CommandLine.Buffer == nullptr)
125 | goto Cleanup;
126 |
127 | bSuccess = ReadProcessMemory(handle,parameters.CommandLine.Buffer,buf,parameters.CommandLine.Length,&len);
128 | if(bSuccess == FALSE)
129 | goto Cleanup;
130 |
131 | cmdLine = std::wstring(buf);
132 |
133 | Cleanup:
134 | return cmdLine;
135 |
136 | }
137 |
138 | bool process_builder::is_managed_process(DWORD pid) { // Based on the idea of processhacker: https://processhacker.sourceforge.io/doc/native_8c_source.html#l04766
139 |
140 | std::vector ManagedDlls ={L"clr.dll",L"mscorwks.dll",L"mscorsvr.dll",L"mscorlib.dll",L"mscorlib.ni.dll",L"coreclr.dll",L"clrjit.dll"};
141 |
142 | PCWSTR fmt_v4 = L"\\BaseNamedObjects\\Cor_Private_IPCBlock_v4_%d";
143 | PCWSTR fmt_v2 = L"\\BaseNamedObjects\\Cor_Private_IPCBlock_%d";
144 | WCHAR sectionName[MAX_PATH] ={0};
145 | UNICODE_STRING UsSectionName ={0};
146 | OBJECT_ATTRIBUTES objectAttributes ={0};
147 |
148 | BOOL bIsManaged = FALSE;
149 | NTSTATUS status = STATUS_UNSUCCESSFUL;
150 | HANDLE hSection = NULL;
151 |
152 | StringCbPrintfW(sectionName,MAX_PATH,fmt_v4,pid);
153 | RtlInitUnicodeString(&UsSectionName,sectionName);
154 |
155 | InitializeObjectAttributes(
156 | &objectAttributes,
157 | &UsSectionName,
158 | OBJ_CASE_INSENSITIVE,
159 | NULL,
160 | NULL
161 | );
162 |
163 | status = NtOpenSection(
164 | &hSection,
165 | SECTION_QUERY,
166 | &objectAttributes
167 | );
168 |
169 | if(NT_SUCCESS(status) || status == STATUS_ACCESS_DENIED) {
170 | bIsManaged = TRUE;
171 | }
172 | else {
173 |
174 | StringCbPrintfW(sectionName,MAX_PATH,fmt_v2,pid);
175 | RtlInitUnicodeString(&UsSectionName,sectionName);
176 |
177 | InitializeObjectAttributes(
178 | &objectAttributes,
179 | &UsSectionName,
180 | OBJ_CASE_INSENSITIVE,
181 | NULL,
182 | NULL
183 | );
184 |
185 | status = NtOpenSection(
186 | &hSection,
187 | SECTION_QUERY,
188 | &objectAttributes
189 | );
190 |
191 | if(NT_SUCCESS(status) || status == STATUS_ACCESS_DENIED) {
192 | bIsManaged = TRUE;
193 | }
194 | else {
195 |
196 | MODULEENTRY32 me32;
197 | auto hModuleSnap = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE,pid);
198 | if(hModuleSnap == INVALID_HANDLE_VALUE)
199 | goto Cleanup;
200 |
201 | me32.dwSize = sizeof(MODULEENTRY32);
202 | if(!Module32First(hModuleSnap,&me32))
203 | {
204 | CloseHandle(hModuleSnap);
205 | goto Cleanup;
206 | }
207 |
208 | do {
209 | if(std::find(ManagedDlls.begin(),ManagedDlls.end(),me32.szModule) != ManagedDlls.end())
210 | {
211 | bIsManaged = TRUE;
212 | break;
213 | }
214 |
215 | } while(Module32Next(hModuleSnap,&me32));
216 |
217 | CloseHandle(hModuleSnap);
218 |
219 | }
220 |
221 | }
222 |
223 | Cleanup:
224 |
225 | if(hSection)
226 | CloseHandle(hSection);
227 |
228 | return bIsManaged;
229 | }
230 |
231 | #pragma endregion
232 |
233 | }
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/suspicious_timer.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include
4 |
5 | #include "phnt.h"
6 | #include "shlwapi.h"
7 | #include "psapi.h"
8 | #include "misc.hpp"
9 | #include "scans.hpp"
10 | #include "threadpooling.h"
11 |
12 | namespace hsb::scanning {
13 |
14 | typedef struct _SUSPICIOUS_CALLBACK {
15 |
16 | std::wstring name;
17 | PVOID addr;
18 |
19 | } SUSPICIOUS_CALLBACK, * PSUSPICIOUS_CALLBACK;
20 |
21 |
22 | SUSPICIOUS_CALLBACK make_callback(std::wstring name, PVOID addr) {
23 | SUSPICIOUS_CALLBACK callback;
24 | callback.name = name;
25 | callback.addr = addr;
26 | return callback;
27 | }
28 |
29 | static std::vector suspicious_callbacks;
30 | static inline std::once_flag callbacks_resolved;
31 |
32 | static void initialize_suspicious_callbacks(void);
33 | static BOOL get_workerfactory_handles(process* process, ACCESS_MASK access, std::vector& pHandleList);
34 | static std::string addr_to_module(HANDLE, DWORD64);
35 |
36 | process_scan process_scans::suspicious_timer = [](process* process) {
37 |
38 | std::vector workerFactories;
39 | WORKER_FACTORY_BASIC_INFORMATION wfbi = { 0 };
40 | FULL_TP_POOL full_tp_pool = { 0 };
41 | PFULL_TP_TIMER p_tp_timer = NULL, p_head = NULL;
42 | FULL_TP_TIMER tp_timer = { 0 };
43 | TPP_CLEANUP_GROUP_MEMBER ctx = { 0 };
44 | SIZE_T len = 0;
45 |
46 | BOOL bSuccess = FALSE;
47 |
48 | std::call_once(callbacks_resolved, initialize_suspicious_callbacks);
49 |
50 | bSuccess = get_workerfactory_handles(process, WORKER_FACTORY_ALL_ACCESS, workerFactories);
51 | if (bSuccess == FALSE)
52 | return;
53 |
54 | for (HANDLE hWorkerFactory : workerFactories) {
55 |
56 | if (NtQueryInformationWorkerFactory(hWorkerFactory, WorkerFactoryBasicInformation, &wfbi, sizeof(WORKER_FACTORY_BASIC_INFORMATION), NULL) == STATUS_SUCCESS) {
57 |
58 | bSuccess = ReadProcessMemory(process->handle, wfbi.StartParameter, &full_tp_pool, sizeof(FULL_TP_POOL), &len);
59 | if (bSuccess == FALSE)
60 | continue;
61 |
62 | if (full_tp_pool.TimerQueue.RelativeQueue.WindowStart.Root)
63 | p_tp_timer = CONTAINING_RECORD(full_tp_pool.TimerQueue.RelativeQueue.WindowStart.Root, FULL_TP_TIMER, WindowStartLinks);
64 | else if (full_tp_pool.TimerQueue.AbsoluteQueue.WindowStart.Root)
65 | p_tp_timer = CONTAINING_RECORD(full_tp_pool.TimerQueue.AbsoluteQueue.WindowStart.Root, FULL_TP_TIMER, WindowStartLinks);
66 | else
67 | continue;
68 |
69 |
70 | bSuccess = ReadProcessMemory(process->handle, p_tp_timer, &tp_timer, sizeof(FULL_TP_TIMER), &len);
71 | if (bSuccess == FALSE)
72 | continue;
73 |
74 | PLIST_ENTRY pHead = tp_timer.WindowStartLinks.Children.Flink;
75 | PLIST_ENTRY pFwd = tp_timer.WindowStartLinks.Children.Flink;
76 | LIST_ENTRY entry = { 0 };
77 |
78 | do {
79 |
80 | bSuccess = ReadProcessMemory(process->handle, tp_timer.Work.CleanupGroupMember.Context, &ctx, sizeof(TPP_CLEANUP_GROUP_MEMBER), &len);
81 | if (bSuccess == FALSE)
82 | break;
83 |
84 | for (SUSPICIOUS_CALLBACK suspiciousCallback : suspicious_callbacks) {
85 | if (suspiciousCallback.addr == ctx.FinalizationCallback) {
86 |
87 | process_detection p;
88 | p.name = L"Suspicious Timer";
89 | p.description = std::format(L"A suspicious timer callback was identified pointing to {}", suspiciousCallback.name);
90 | p.severity = hsb::containers::detections::CRITICAL;
91 |
92 | process->add_detection(p);
93 |
94 | }
95 | }
96 |
97 | p_tp_timer = CONTAINING_RECORD(pFwd, FULL_TP_TIMER, WindowStartLinks);
98 | bSuccess = ReadProcessMemory(process->handle, p_tp_timer, &tp_timer, sizeof(FULL_TP_TIMER), &len);
99 | if (bSuccess == FALSE)
100 | break;
101 |
102 | ReadProcessMemory(process->handle, pFwd, &entry, sizeof(LIST_ENTRY), &len);
103 | pFwd = entry.Flink;
104 |
105 | } while (pHead != pFwd);
106 |
107 | }
108 |
109 | CloseHandle(hWorkerFactory);
110 |
111 | }
112 |
113 | };
114 |
115 | static BOOL get_workerfactory_handles(process* process, ACCESS_MASK access, std::vector& pHandleList) {
116 |
117 | BOOL bSuccess = FALSE;
118 | PVOID pBuffer = NULL;
119 | ULONG uBufferSize = 0;
120 | NTSTATUS status = 0;
121 | HANDLE dupHandle = NULL;
122 |
123 | CLIENT_ID clientId = { 0 };
124 |
125 | PSYSTEM_HANDLE_INFORMATION handleInfo = NULL;
126 | PSYSTEM_HANDLE_TABLE_ENTRY_INFO entryInfo = NULL;
127 | POBJECT_TYPE_INFORMATION objectTypeInfo = NULL;
128 |
129 | do {
130 |
131 | status = NtQuerySystemInformation((SYSTEM_INFORMATION_CLASS)SystemHandleInformation, pBuffer, uBufferSize, &uBufferSize);
132 | if (!NT_SUCCESS(status)) {
133 |
134 | if (status == STATUS_INFO_LENGTH_MISMATCH) {
135 | if (pBuffer != NULL)
136 | LocalFree(pBuffer);
137 | pBuffer = LocalAlloc(LMEM_ZEROINIT, uBufferSize);
138 | if (pBuffer == NULL)
139 | goto Cleanup;
140 |
141 | continue;
142 | }
143 | break;
144 |
145 | }
146 | else {
147 |
148 | handleInfo = (PSYSTEM_HANDLE_INFORMATION)pBuffer;
149 | break;
150 |
151 | }
152 |
153 | } while (1);
154 |
155 | if (handleInfo == NULL)
156 | goto Cleanup;
157 |
158 | objectTypeInfo = (POBJECT_TYPE_INFORMATION)LocalAlloc(LMEM_ZEROINIT, sizeof(OBJECT_TYPE_INFORMATION) * 2);
159 | if (objectTypeInfo == NULL)
160 | goto Cleanup;
161 |
162 | for (UINT i = 0; i < handleInfo->NumberOfHandles; i++) {
163 |
164 | entryInfo = &handleInfo->Handles[i];
165 |
166 | if (process->pid != entryInfo->UniqueProcessId)
167 | continue;
168 |
169 | clientId.UniqueProcess = (HANDLE)entryInfo->UniqueProcessId;
170 | clientId.UniqueThread = 0;
171 |
172 | if (NtDuplicateObject(process->handle, (HANDLE)(uint64_t)entryInfo->HandleValue, NtCurrentProcess(), &dupHandle, access, 0, 0) == STATUS_SUCCESS) {
173 |
174 | memset(objectTypeInfo, 0, sizeof(OBJECT_TYPE_INFORMATION) * 2);
175 |
176 | if (NtQueryObject(dupHandle, (OBJECT_INFORMATION_CLASS)ObjectTypeInformation, objectTypeInfo, sizeof(OBJECT_TYPE_INFORMATION) * 2, NULL) == STATUS_SUCCESS) {
177 |
178 | if (!lstrcmpW(objectTypeInfo->TypeName.Buffer, L"TpWorkerFactory")) {
179 | pHandleList.push_back(dupHandle);
180 | }
181 | else {
182 | CloseHandle(dupHandle);
183 | }
184 |
185 | }
186 |
187 | }
188 |
189 | }
190 |
191 | bSuccess = TRUE;
192 |
193 | Cleanup:
194 |
195 | if (objectTypeInfo)
196 | LocalFree(objectTypeInfo);
197 |
198 | if (pBuffer)
199 | LocalFree(pBuffer);
200 |
201 | return bSuccess;
202 |
203 | }
204 |
205 | static std::string addr_to_module(HANDLE hProcess, DWORD64 pAddr) {
206 |
207 | PSTR name = NULL;
208 | CHAR buffer[MAX_PATH] = { 0 };
209 | SIZE_T s = 0;
210 | MEMORY_BASIC_INFORMATION mbe = { 0 };
211 |
212 | s = VirtualQueryEx(hProcess, (PVOID)pAddr, &mbe, sizeof(MEMORY_BASIC_INFORMATION));
213 | if (s == 0)
214 | return std::string("");
215 |
216 | if (GetMappedFileNameA(hProcess, mbe.AllocationBase, buffer, MAX_PATH) == 0 ) {
217 | return std::string("");
218 | }
219 |
220 | name = PathFindFileNameA(buffer);
221 |
222 | return std::string(name);
223 |
224 | }
225 |
226 |
227 | static void initialize_suspicious_callbacks(void) {
228 |
229 | HMODULE hNtdll = GetModuleHandleA("ntdll.dll");
230 | HMODULE hKernel32 = GetModuleHandleA("kernel32.dll");
231 |
232 | if (hNtdll && hKernel32) {
233 | suspicious_callbacks.push_back(make_callback(L"ntdll!NtContinue", GetProcAddress(hNtdll, "NtContinue")));
234 | suspicious_callbacks.push_back(make_callback(L"ntdll!RtlCaptureContext", GetProcAddress(hNtdll, "RtlCaptureContext")));
235 | suspicious_callbacks.push_back(make_callback(L"ntdll!RtlCopyMemory", GetProcAddress(hNtdll, "RtlCopyMemory")));
236 | suspicious_callbacks.push_back(make_callback(L"ntdll!RtlMoveMemory", GetProcAddress(hNtdll, "RtlMoveMemory")));
237 | suspicious_callbacks.push_back(make_callback(L"ntdll!NtFreeVirtualMemory", GetProcAddress(hNtdll, "NtFreeVirtualMemory")));
238 | suspicious_callbacks.push_back(make_callback(L"ntdll!NtAllocateVirtualMemory", GetProcAddress(hNtdll, "NtAllocateVirtualMemory")));
239 | suspicious_callbacks.push_back(make_callback(L"ntdll!NtCreateThread", GetProcAddress(hNtdll, "NtCreateThread")));
240 | suspicious_callbacks.push_back(make_callback(L"kernel32!ResumeThread", GetProcAddress(hKernel32, "ResumeThread")));
241 | }
242 |
243 | }
244 |
245 | }
--------------------------------------------------------------------------------
/inc/BS_thread_pool_utils.hpp:
--------------------------------------------------------------------------------
1 | #ifndef BS_THREAD_POOL_UTILS_HPP
2 | #define BS_THREAD_POOL_UTILS_HPP
3 | /**
4 | * @file BS_thread_pool_utils.hpp
5 | * @author Barak Shoshany (baraksh@gmail.com) (https://baraksh.com)
6 | * @version 4.1.0
7 | * @date 2024-03-22
8 | * @copyright Copyright (c) 2024 Barak Shoshany. Licensed under the MIT license. If you found this project useful, please consider starring it on GitHub! If you use this library in software of any kind, please provide a link to the GitHub repository https://github.com/bshoshany/thread-pool in the source code and documentation. If you use this library in published research, please cite it as follows: Barak Shoshany, "A C++17 Thread Pool for High-Performance Scientific Computing", doi:10.1016/j.softx.2024.101687, SoftwareX 26 (2024) 101687, arXiv:2105.00613
9 | *
10 | * @brief BS::thread_pool: a fast, lightweight, and easy-to-use C++17 thread pool library. This header file contains independent utility classes that are part of the library, but are not needed to use the thread pool itself.
11 | */
12 |
13 | #include // std::chrono
14 | #include // std::promise, std::shared_future
15 | #include // std::initializer_list
16 | #include // std::cout
17 | #include // std::make_unique, std::unique_ptr
18 | #include // std::mutex, std::scoped_lock
19 | #include // std::endl, std::flush, std::ostream
20 | #include // std::forward
21 |
22 | /**
23 | * @brief A namespace used by Barak Shoshany's projects.
24 | */
25 | namespace BS {
26 | // Macros indicating the version of the thread pool utilities library.
27 | #define BS_THREAD_POOL_UTILS_VERSION_MAJOR 4
28 | #define BS_THREAD_POOL_UTILS_VERSION_MINOR 1
29 | #define BS_THREAD_POOL_UTILS_VERSION_PATCH 0
30 |
31 | /**
32 | * @brief A utility class to allow simple signalling between threads.
33 | */
34 | class [[nodiscard]] signaller
35 | {
36 | public:
37 | /**
38 | * @brief Construct a new signaller.
39 | */
40 | signaller(): promise(),future(promise.get_future()) {}
41 |
42 | // The copy constructor and copy assignment operator are deleted. The signaller works using a promise, which cannot be copied.
43 | signaller(const signaller&) = delete;
44 | signaller& operator=(const signaller&) = delete;
45 |
46 | // The move constructor and move assignment operator are defaulted.
47 | signaller(signaller&&) = default;
48 | signaller& operator=(signaller&&) = default;
49 |
50 | /**
51 | * @brief Inform any waiting threads that the signaller is ready.
52 | */
53 | void ready()
54 | {
55 | promise.set_value();
56 | }
57 |
58 | /**
59 | * @brief Wait until the signaller is ready.
60 | */
61 | void wait()
62 | {
63 | future.wait();
64 | }
65 |
66 | private:
67 | /**
68 | * @brief A promise used to set the state of the signaller.
69 | */
70 | std::promise promise;
71 |
72 | /**
73 | * @brief A future used to wait for the signaller.
74 | */
75 | std::shared_future future;
76 | }; // class signaller
77 |
78 | /**
79 | * @brief A utility class to synchronize printing to an output stream by different threads.
80 | */
81 | class [[nodiscard]] synced_stream
82 | {
83 | public:
84 | /**
85 | * @brief Construct a new synced stream.
86 | *
87 | * @param stream The output stream to print to. The default value is `std::cout`.
88 | */
89 | explicit synced_stream(std::ostream& stream = std::cout): out_stream(stream) {}
90 |
91 | // The copy and move constructors and assignment operators are deleted. The synced stream uses a mutex, which cannot be copied or moved.
92 | synced_stream(const synced_stream&) = delete;
93 | synced_stream(synced_stream&&) = delete;
94 | synced_stream& operator=(const synced_stream&) = delete;
95 | synced_stream& operator=(synced_stream&&) = delete;
96 |
97 | /**
98 | * @brief Print any number of items into the output stream. Ensures that no other threads print to this stream simultaneously, as long as they all exclusively use the same `synced_stream` object to print.
99 | *
100 | * @tparam T The types of the items.
101 | * @param items The items to print.
102 | */
103 | template
104 | void print(T&&... items)
105 | {
106 | const std::scoped_lock stream_lock(stream_mutex);
107 | (out_stream << ... << std::forward(items));
108 | }
109 |
110 | /**
111 | * @brief Print any number of items into the output stream, followed by a newline character. Ensures that no other threads print to this stream simultaneously, as long as they all exclusively use the same `synced_stream` object to print.
112 | *
113 | * @tparam T The types of the items.
114 | * @param items The items to print.
115 | */
116 | template
117 | void println(T&&... items)
118 | {
119 | print(std::forward(items)...,'\n');
120 | }
121 |
122 | /**
123 | * @brief A stream manipulator to pass to a `synced_stream` (an explicit cast of `std::endl`). Prints a newline character to the stream, and then flushes it. Should only be used if flushing is desired, otherwise a newline character should be used instead.
124 | */
125 | inline static std::ostream& (&endl)(std::ostream&) = static_cast(std::endl);
126 |
127 | /**
128 | * @brief A stream manipulator to pass to a `synced_stream` (an explicit cast of `std::flush`). Used to flush the stream.
129 | */
130 | inline static std::ostream& (&flush)(std::ostream&) = static_cast(std::flush);
131 |
132 | private:
133 | /**
134 | * @brief The output stream to print to.
135 | */
136 | std::ostream& out_stream;
137 |
138 | /**
139 | * @brief A mutex to synchronize printing.
140 | */
141 | mutable std::mutex stream_mutex ={};
142 | }; // class synced_stream
143 |
144 | /**
145 | * @brief A utility class to measure execution time for benchmarking purposes.
146 | */
147 | class [[nodiscard]] timer
148 | {
149 | public:
150 | /**
151 | * @brief Construct a new timer and immediately start measuring time.
152 | */
153 | timer() = default;
154 |
155 | /**
156 | * @brief Get the number of milliseconds that have elapsed since the object was constructed or since `start()` was last called, but keep the timer ticking.
157 | *
158 | * @return The number of milliseconds.
159 | */
160 | [[nodiscard]] std::chrono::milliseconds::rep current_ms() const
161 | {
162 | return (std::chrono::duration_cast(std::chrono::steady_clock::now() - start_time)).count();
163 | }
164 |
165 | /**
166 | * @brief Start (or restart) measuring time. Note that the timer starts ticking as soon as the object is created, so this is only necessary if we want to restart the clock later.
167 | */
168 | void start()
169 | {
170 | start_time = std::chrono::steady_clock::now();
171 | }
172 |
173 | /**
174 | * @brief Stop measuring time and store the elapsed time since the object was constructed or since `start()` was last called.
175 | */
176 | void stop()
177 | {
178 | elapsed_time = std::chrono::steady_clock::now() - start_time;
179 | }
180 |
181 | /**
182 | * @brief Get the number of milliseconds stored when `stop()` was last called.
183 | *
184 | * @return The number of milliseconds.
185 | */
186 | [[nodiscard]] std::chrono::milliseconds::rep ms() const
187 | {
188 | return (std::chrono::duration_cast(elapsed_time)).count();
189 | }
190 |
191 | private:
192 | /**
193 | * @brief The time point when measuring started.
194 | */
195 | std::chrono::time_point start_time = std::chrono::steady_clock::now();
196 |
197 | /**
198 | * @brief The duration that has elapsed between `start()` and `stop()`.
199 | */
200 | std::chrono::duration elapsed_time = std::chrono::duration::zero();
201 | }; // class timer
202 | } // namespace BS
203 | #endif
204 |
--------------------------------------------------------------------------------
/inc/threadpooling.h:
--------------------------------------------------------------------------------
1 | #pragma once
2 |
3 | /*
4 |
5 | BSD 3-Clause License
6 |
7 | Copyright (c) 2023, SafeBreach Labs
8 |
9 | Redistribution and use in source and binary forms, with or without
10 | modification, are permitted provided that the following conditions are met:
11 |
12 | 1. Redistributions of source code must retain the above copyright notice, this
13 | list of conditions and the following disclaimer.
14 |
15 | 2. Redistributions in binary form must reproduce the above copyright notice,
16 | this list of conditions and the following disclaimer in the documentation
17 | and/or other materials provided with the distribution.
18 |
19 | 3. Neither the name of the copyright holder nor the names of its
20 | contributors may be used to endorse or promote products derived from
21 | this software without specific prior written permission.
22 |
23 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
24 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
25 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
26 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
27 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
28 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
29 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
30 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
31 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
32 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
33 |
34 | */
35 |
36 |
37 | // Taken from: https://github.com/SafeBreach-Labs/PoolParty
38 |
39 | #include "Phnt.h"
40 |
41 | typedef struct _TP_TASK_CALLBACKS
42 | {
43 | void* ExecuteCallback;
44 | void* Unposted;
45 | } TP_TASK_CALLBACKS, * PTP_TASK_CALLBACKS;
46 |
47 |
48 | typedef struct _TP_TASK
49 | {
50 | struct _TP_TASK_CALLBACKS* Callbacks;
51 | UINT32 NumaNode;
52 | UINT8 IdealProcessor;
53 | char Padding_242[3];
54 | struct _LIST_ENTRY ListEntry;
55 | } TP_TASK, * PTP_TASK;
56 |
57 |
58 | typedef struct _TPP_REFCOUNT
59 | {
60 | volatile INT32 Refcount;
61 | } TPP_REFCOUNT, * PTPP_REFCOUNT;
62 |
63 |
64 | typedef struct _TPP_CALLER
65 | {
66 | void* ReturnAddress;
67 | } TPP_CALLER, * PTPP_CALLER;
68 |
69 |
70 | typedef struct _TPP_PH
71 | {
72 | struct _TPP_PH_LINKS* Root;
73 | } TPP_PH, * PTPP_PH;
74 |
75 |
76 | typedef struct _TP_DIRECT
77 | {
78 | struct _TP_TASK Task;
79 | UINT64 Lock;
80 | struct _LIST_ENTRY IoCompletionInformationList;
81 | void* Callback;
82 | UINT32 NumaNode;
83 | UINT8 IdealProcessor;
84 | char __PADDING__[3];
85 | } TP_DIRECT, * PTP_DIRECT;
86 |
87 |
88 | typedef struct _TPP_TIMER_SUBQUEUE
89 | {
90 | INT64 Expiration;
91 | struct _TPP_PH WindowStart;
92 | struct _TPP_PH WindowEnd;
93 | void* Timer;
94 | void* TimerPkt;
95 | struct _TP_DIRECT Direct;
96 | UINT32 ExpirationWindow;
97 | INT32 __PADDING__[1];
98 | } TPP_TIMER_SUBQUEUE, * PTPP_TIMER_SUBQUEUE;
99 |
100 |
101 | typedef struct _TPP_TIMER_QUEUE
102 | {
103 | struct _RTL_SRWLOCK Lock;
104 | struct _TPP_TIMER_SUBQUEUE AbsoluteQueue;
105 | struct _TPP_TIMER_SUBQUEUE RelativeQueue;
106 | INT32 AllocatedTimerCount;
107 | INT32 __PADDING__[1];
108 | } TPP_TIMER_QUEUE, * PTPP_TIMER_QUEUE;
109 |
110 |
111 | typedef struct _TPP_NUMA_NODE
112 | {
113 | INT32 WorkerCount;
114 | } TPP_NUMA_NODE, * PTPP_NUMA_NODE;
115 |
116 |
117 | typedef union _TPP_POOL_QUEUE_STATE
118 | {
119 | union
120 | {
121 | INT64 Exchange;
122 | struct
123 | {
124 | INT32 RunningThreadGoal : 16;
125 | UINT32 PendingReleaseCount : 16;
126 | UINT32 QueueLength;
127 | };
128 | };
129 | } TPP_POOL_QUEUE_STATE, * PTPP_POOL_QUEUE_STATE;
130 |
131 |
132 | typedef struct _TPP_QUEUE
133 | {
134 | struct _LIST_ENTRY Queue;
135 | struct _RTL_SRWLOCK Lock;
136 | } TPP_QUEUE, * PTPP_QUEUE;
137 |
138 |
139 | typedef struct _FULL_TP_POOL
140 | {
141 | struct _TPP_REFCOUNT Refcount;
142 | long Padding_239;
143 | union _TPP_POOL_QUEUE_STATE QueueState;
144 | struct _TPP_QUEUE* TaskQueue[3];
145 | struct _TPP_NUMA_NODE* NumaNode;
146 | struct _GROUP_AFFINITY* ProximityInfo;
147 | void* WorkerFactory;
148 | void* CompletionPort;
149 | struct _RTL_SRWLOCK Lock;
150 | struct _LIST_ENTRY PoolObjectList;
151 | struct _LIST_ENTRY WorkerList;
152 | struct _TPP_TIMER_QUEUE TimerQueue;
153 | struct _RTL_SRWLOCK ShutdownLock;
154 | UINT8 ShutdownInitiated;
155 | UINT8 Released;
156 | UINT16 PoolFlags;
157 | long Padding_240;
158 | struct _LIST_ENTRY PoolLinks;
159 | struct _TPP_CALLER AllocCaller;
160 | struct _TPP_CALLER ReleaseCaller;
161 | volatile INT32 AvailableWorkerCount;
162 | volatile INT32 LongRunningWorkerCount;
163 | UINT32 LastProcCount;
164 | volatile INT32 NodeStatus;
165 | volatile INT32 BindingCount;
166 | UINT32 CallbackChecksDisabled : 1;
167 | UINT32 TrimTarget : 11;
168 | UINT32 TrimmedThrdCount : 11;
169 | UINT32 SelectedCpuSetCount;
170 | long Padding_241;
171 | struct _RTL_CONDITION_VARIABLE TrimComplete;
172 | struct _LIST_ENTRY TrimmedWorkerList;
173 | } FULL_TP_POOL, * PFULL_TP_POOL;
174 |
175 | typedef union _TPP_WORK_STATE
176 | {
177 | union
178 | {
179 | INT32 Exchange;
180 | UINT32 Insertable : 1;
181 | UINT32 PendingCallbackCount : 31;
182 | };
183 | } TPP_WORK_STATE, * PTPP_WORK_STATE;
184 |
185 |
186 | typedef struct _TPP_ITE_WAITER
187 | {
188 | struct _TPP_ITE_WAITER* Next;
189 | void* ThreadId;
190 | } TPP_ITE_WAITER, * PTPP_ITE_WAITER;
191 |
192 | typedef struct _TPP_PH_LINKS
193 | {
194 | struct _LIST_ENTRY Siblings;
195 | struct _LIST_ENTRY Children;
196 | INT64 Key;
197 | } TPP_PH_LINKS, * PTPP_PH_LINKS;
198 |
199 |
200 | typedef struct _TPP_ITE
201 | {
202 | struct _TPP_ITE_WAITER* First;
203 | } TPP_ITE, * PTPP_ITE;
204 |
205 |
206 | typedef union _TPP_FLAGS_COUNT
207 | {
208 | union
209 | {
210 | UINT64 Count : 60;
211 | UINT64 Flags : 4;
212 | INT64 Data;
213 | };
214 | } TPP_FLAGS_COUNT, * PTPP_FLAGS_COUNT;
215 |
216 |
217 | typedef struct _TPP_BARRIER
218 | {
219 | volatile union _TPP_FLAGS_COUNT Ptr;
220 | struct _RTL_SRWLOCK WaitLock;
221 | struct _TPP_ITE WaitList;
222 | } TPP_BARRIER, * PTPP_BARRIER;
223 |
224 |
225 | typedef struct _TP_CLEANUP_GROUP
226 | {
227 | struct _TPP_REFCOUNT Refcount;
228 | INT32 Released;
229 | struct _RTL_SRWLOCK MemberLock;
230 | struct _LIST_ENTRY MemberList;
231 | struct _TPP_BARRIER Barrier;
232 | struct _RTL_SRWLOCK CleanupLock;
233 | struct _LIST_ENTRY CleanupList;
234 | } TP_CLEANUP_GROUP, * PTP_CLEANUP_GROUP;
235 |
236 |
237 | typedef struct _TPP_CLEANUP_GROUP_MEMBER
238 | {
239 | struct _TPP_REFCOUNT Refcount;
240 | long Padding_233;
241 | const struct _TPP_CLEANUP_GROUP_MEMBER_VFUNCS* VFuncs;
242 | struct _TP_CLEANUP_GROUP* CleanupGroup;
243 | void* CleanupGroupCancelCallback;
244 | void* FinalizationCallback;
245 | struct _LIST_ENTRY CleanupGroupMemberLinks;
246 | struct _TPP_BARRIER CallbackBarrier;
247 | union
248 | {
249 | void* Callback;
250 | void* WorkCallback;
251 | void* SimpleCallback;
252 | void* TimerCallback;
253 | void* WaitCallback;
254 | void* IoCallback;
255 | void* AlpcCallback;
256 | void* AlpcCallbackEx;
257 | void* JobCallback;
258 | };
259 | void* Context;
260 | struct _ACTIVATION_CONTEXT* ActivationContext;
261 | void* SubProcessTag;
262 | struct _GUID ActivityId;
263 | struct _ALPC_WORK_ON_BEHALF_TICKET WorkOnBehalfTicket;
264 | void* RaceDll;
265 | FULL_TP_POOL* Pool;
266 | struct _LIST_ENTRY PoolObjectLinks;
267 | union
268 | {
269 | volatile INT32 Flags;
270 | UINT32 LongFunction : 1;
271 | UINT32 Persistent : 1;
272 | UINT32 UnusedPublic : 14;
273 | UINT32 Released : 1;
274 | UINT32 CleanupGroupReleased : 1;
275 | UINT32 InCleanupGroupCleanupList : 1;
276 | UINT32 UnusedPrivate : 13;
277 | };
278 | long Padding_234;
279 | struct _TPP_CALLER AllocCaller;
280 | struct _TPP_CALLER ReleaseCaller;
281 | enum _TP_CALLBACK_PRIORITY CallbackPriority;
282 | INT32 __PADDING__[1];
283 | } TPP_CLEANUP_GROUP_MEMBER, * PTPP_CLEANUP_GROUP_MEMBER;
284 |
285 |
286 | typedef struct _FULL_TP_WORK
287 | {
288 | struct _TPP_CLEANUP_GROUP_MEMBER CleanupGroupMember;
289 | struct _TP_TASK Task;
290 | volatile union _TPP_WORK_STATE WorkState;
291 | INT32 __PADDING__[1];
292 | } FULL_TP_WORK, * PFULL_TP_WORK;
293 |
294 |
295 | typedef struct _FULL_TP_TIMER
296 | {
297 | struct _FULL_TP_WORK Work;
298 | struct _RTL_SRWLOCK Lock;
299 | union
300 | {
301 | struct _TPP_PH_LINKS WindowEndLinks;
302 | struct _LIST_ENTRY ExpirationLinks;
303 | };
304 | struct _TPP_PH_LINKS WindowStartLinks;
305 | INT64 DueTime;
306 | struct _TPP_ITE Ite;
307 | UINT32 Window;
308 | UINT32 Period;
309 | UINT8 Inserted;
310 | UINT8 WaitTimer;
311 | union
312 | {
313 | UINT8 TimerStatus;
314 | UINT8 InQueue : 1;
315 | UINT8 Absolute : 1;
316 | UINT8 Cancelled : 1;
317 | };
318 | UINT8 BlockInsert;
319 | INT32 __PADDING__[1];
320 | } FULL_TP_TIMER, * PFULL_TP_TIMER;
321 |
322 |
323 | typedef struct _FULL_TP_WAIT
324 | {
325 | struct _FULL_TP_TIMER Timer;
326 | void* Handle;
327 | void* WaitPkt;
328 | void* NextWaitHandle;
329 | union _LARGE_INTEGER NextWaitTimeout;
330 | struct _TP_DIRECT Direct;
331 | union
332 | {
333 | union
334 | {
335 | UINT8 AllFlags;
336 | UINT8 NextWaitActive : 1;
337 | UINT8 NextTimeoutActive : 1;
338 | UINT8 CallbackCounted : 1;
339 | UINT8 Spare : 5;
340 | };
341 | } WaitFlags;
342 | char __PADDING__[7];
343 | } FULL_TP_WAIT, * PFULL_TP_WAIT;
344 |
345 |
346 | typedef struct _FULL_TP_IO
347 | {
348 | struct _TPP_CLEANUP_GROUP_MEMBER CleanupGroupMember;
349 | struct _TP_DIRECT Direct;
350 | void* File;
351 | volatile INT32 PendingIrpCount;
352 | INT32 __PADDING__[1];
353 | } FULL_TP_IO, * PFULL_TP_IO;
354 |
355 |
356 | typedef struct _FULL_TP_ALPC
357 | {
358 | struct _TP_DIRECT Direct;
359 | struct _TPP_CLEANUP_GROUP_MEMBER CleanupGroupMember;
360 | void* AlpcPort;
361 | INT32 DeferredSendCount;
362 | INT32 LastConcurrencyCount;
363 | union
364 | {
365 | UINT32 Flags;
366 | UINT32 ExTypeCallback : 1;
367 | UINT32 CompletionListRegistered : 1;
368 | UINT32 Reserved : 30;
369 | };
370 | INT32 __PADDING__[1];
371 | } FULL_TP_ALPC, * PFULL_TP_ALPC;
--------------------------------------------------------------------------------
/src/UnitTests/UnitTests.cpp:
--------------------------------------------------------------------------------
1 | #include "pch.h"
2 | #include "CppUnitTest.h"
3 |
4 | #include "process_enumerator.hpp"
5 | #include "process_scanner.hpp"
6 |
7 | #include "blocking_apc.cpp"
8 | #include "blocking_timer.cpp"
9 | #include "suspicious_timer.cpp"
10 | #include "abnormal_intermodular_call.cpp"
11 | #include "hardware_breakpoints.cpp"
12 | #include "private_memory.cpp"
13 | #include "return_address_spoofing.cpp"
14 | #include "stomped_module.cpp"
15 | #include "non_executable_memory.cpp"
16 |
17 | using namespace Microsoft::VisualStudio::CppUnitTestFramework;
18 | using namespace hsb;
19 |
20 | namespace HSBUnittests
21 | {
22 |
23 | extern "C" DWORD64 getgadget(void);
24 |
25 | void blocking_function(void) {
26 |
27 | DWORD64 actualAddress = getgadget();
28 | void** pvAddressOfReturnAddress = (void**) _AddressOfReturnAddress();
29 | *pvAddressOfReturnAddress = (void*)actualAddress;
30 |
31 | Sleep(1000 * 60 * 60);
32 |
33 | }
34 |
35 | void queue_blocking_timers(void){
36 |
37 | HANDLE hTimerQueue = NULL;
38 | HANDLE hNewTimer = NULL, hNewTimer2 = NULL;
39 |
40 | DWORD sleeptime = 1000 * 60 * 60;
41 |
42 | hTimerQueue = CreateTimerQueue();
43 |
44 | CreateTimerQueueTimer(&hNewTimer,hTimerQueue,(WAITORTIMERCALLBACK) blocking_function, NULL,0,0,WT_EXECUTELONGFUNCTION);
45 | CreateTimerQueueTimer(&hNewTimer2,hTimerQueue,(WAITORTIMERCALLBACK)Sleep,(PVOID)&sleeptime,0,0,WT_EXECUTELONGFUNCTION);
46 |
47 | }
48 |
49 | void queue_future_timer(void){
50 | HANDLE hTimerQueue = NULL;
51 | HANDLE hNewTimer = NULL;
52 |
53 | PVOID pNtContinue = NULL;
54 | pNtContinue = GetProcAddress(GetModuleHandleA("ntdll.dll"),"NtContinue");
55 |
56 | hTimerQueue = CreateTimerQueue();
57 | CreateTimerQueueTimer(&hNewTimer,hTimerQueue,(WAITORTIMERCALLBACK)pNtContinue,NULL,1000 * 60,0,WT_EXECUTEINTIMERTHREAD);
58 | }
59 |
60 | void queue_blocking_apc(void) {
61 | HANDLE h_thread = CreateThread(NULL,0,(LPTHREAD_START_ROUTINE)queue_blocking_apc,NULL,CREATE_SUSPENDED,NULL);
62 | QueueUserAPC((PAPCFUNC)blocking_function,h_thread,NULL);
63 | ResumeThread(h_thread);
64 | }
65 |
66 | void exec_msgbox_shellcode(void) {
67 |
68 | //msgbox
69 | unsigned char data[] ={0x56,0x48,0x89,0xe6,0x48,0x83,0xe4,0xf0,0x48,0x83,0xec,0x20,0xe8,0x7f,0x01,0x00,0x00,0x48,0x89,0xf4,0x5e,0xc3,0x66,0x2e,0x0f,0x1f,0x84,0x00,0x00,0x00,0x00,0x00,0x65,0x48,0x8b,0x04,0x25,0x60,0x00,0x00,0x00,0x48,0x8b,0x40,0x18,0x41,0x89,0xca,0x4c,0x8b,0x58,0x20,0x4d,0x89,0xd9,0x66,0x0f,0x1f,0x84,0x00,0x00,0x00,0x00,0x00,0x49,0x8b,0x49,0x50,0x48,0x85,0xc9,0x74,0x63,0x0f,0xb7,0x01,0x66,0x85,0xc0,0x74,0x5f,0x48,0x89,0xca,0x0f,0x1f,0x40,0x00,0x44,0x8d,0x40,0xbf,0x66,0x41,0x83,0xf8,0x19,0x77,0x06,0x83,0xc0,0x20,0x66,0x89,0x02,0x0f,0xb7,0x42,0x02,0x48,0x83,0xc2,0x02,0x66,0x85,0xc0,0x75,0xe2,0x0f,0xb7,0x01,0x66,0x85,0xc0,0x74,0x32,0x41,0xb8,0x05,0x15,0x00,0x00,0x0f,0x1f,0x40,0x00,0x44,0x89,0xc2,0x48,0x83,0xc1,0x02,0xc1,0xe2,0x05,0x01,0xd0,0x41,0x01,0xc0,0x0f,0xb7,0x01,0x66,0x85,0xc0,0x75,0xe9,0x45,0x39,0xc2,0x74,0x17,0x4d,0x8b,0x09,0x4d,0x39,0xcb,0x75,0x94,0x31,0xc0,0xc3,0x90,0x41,0xb8,0x05,0x15,0x00,0x00,0x45,0x39,0xc2,0x75,0xe9,0x49,0x8b,0x41,0x20,0xc3,0x41,0x54,0x41,0x89,0xd4,0x53,0x89,0xcb,0x48,0x83,0xec,0x38,0xe8,0x4f,0xff,0xff,0xff,0x48,0x85,0xc0,0x75,0x22,0xb9,0x75,0xee,0x40,0x70,0xe8,0x40,0xff,0xff,0xff,0x48,0x89,0xc1,0x48,0x85,0xc0,0x75,0x28,0x48,0x83,0xc4,0x38,0x31,0xc0,0x5b,0x41,0x5c,0xc3,0x66,0x0f,0x1f,0x44,0x00,0x00,0x48,0x89,0xc1,0x48,0x83,0xc4,0x38,0x44,0x89,0xe2,0x5b,0x41,0x5c,0xe9,0xd6,0x00,0x00,0x00,0x66,0x0f,0x1f,0x44,0x00,0x00,0xba,0xfb,0xf0,0xbf,0x5f,0xe8,0xc6,0x00,0x00,0x00,0x48,0x85,0xc0,0x74,0xc9,0x81,0xfb,0xf3,0xd3,0x6b,0x5a,0x74,0x31,0x81,0xfb,0x6d,0x9c,0xbd,0x8d,0x75,0xb9,0x48,0xbb,0x57,0x69,0x6e,0x69,0x6e,0x65,0x74,0x2e,0x48,0x8d,0x4c,0x24,0x24,0xc7,0x44,0x24,0x2c,0x64,0x6c,0x6c,0x00,0x48,0x89,0x5c,0x24,0x24,0xff,0xd0,0x48,0x89,0xc1,0xeb,0x2e,0x66,0x0f,0x1f,0x44,0x00,0x00,0xba,0x6c,0x6c,0x00,0x00,0x48,0x8d,0x4c,0x24,0x24,0xc6,0x44,0x24,0x2e,0x00,0x48,0xbb,0x55,0x73,0x65,0x72,0x33,0x32,0x2e,0x64,0x48,0x89,0x5c,0x24,0x24,0x66,0x89,0x54,0x24,0x2c,0xff,0xd0,0x48,0x89,0xc1,0x48,0x85,0xc9,0x0f,0x85,0x72,0xff,0xff,0xff,0xe9,0x5a,0xff,0xff,0xff,0x90,0x90,0x48,0x83,0xec,0x38,0xba,0xb4,0x14,0x4f,0x38,0xb9,0xf3,0xd3,0x6b,0x5a,0xe8,0x1d,0xff,0xff,0xff,0x45,0x31,0xc0,0x48,0x85,0xc0,0x74,0x25,0x48,0x8d,0x54,0x24,0x2b,0xc7,0x44,0x24,0x2b,0x4d,0x6f,0x69,0x6e,0x41,0xb9,0x01,0x00,0x00,0x00,0x31,0xc9,0x49,0x89,0xd0,0xc6,0x44,0x24,0x2f,0x00,0xff,0xd0,0x41,0xb8,0x01,0x00,0x00,0x00,0x44,0x89,0xc0,0x48,0x83,0xc4,0x38,0xc3,0x90,0x90,0x90,0x90,0x90,0x90,0x90,0x90,0x57,0x56,0x53,0x48,0x63,0x41,0x3c,0x8b,0xbc,0x01,0x88,0x00,0x00,0x00,0x48,0x01,0xcf,0x44,0x8b,0x4f,0x20,0x8b,0x5f,0x18,0x49,0x01,0xc9,0x85,0xdb,0x74,0x59,0x49,0x89,0xcb,0x89,0xd6,0x45,0x31,0xd2,0x66,0x0f,0x1f,0x84,0x00,0x00,0x00,0x00,0x00,0x41,0x8b,0x01,0xb9,0x05,0x15,0x00,0x00,0x4c,0x01,0xd8,0x4c,0x8d,0x40,0x01,0x0f,0xb6,0x00,0x84,0xc0,0x74,0x21,0x66,0x2e,0x0f,0x1f,0x84,0x00,0x00,0x00,0x00,0x00,0x89,0xca,0xc1,0xe2,0x05,0x01,0xd0,0x01,0xc1,0x4c,0x89,0xc0,0x49,0x83,0xc0,0x01,0x0f,0xb6,0x00,0x84,0xc0,0x75,0xe9,0x39,0xce,0x74,0x13,0x49,0x83,0xc2,0x01,0x49,0x83,0xc1,0x04,0x4c,0x39,0xd3,0x75,0xb8,0x5b,0x31,0xc0,0x5e,0x5f,0xc3,0x8b,0x57,0x24,0x4b,0x8d,0x0c,0x53,0x8b,0x47,0x1c,0x5b,0x5e,0x0f,0xb7,0x14,0x11,0x5f,0x49,0x8d,0x14,0x93,0x8b,0x04,0x02,0x4c,0x01,0xd8,0xc3,0x90,0x90,0x90,0x90,0x90,0x90,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0x00,0x00,0x00,0x00,0x00,0x00,0x00};
70 | unsigned int data_size = sizeof(data);
71 |
72 | PVOID pBuffer = VirtualAlloc(0,2048,MEM_COMMIT,PAGE_EXECUTE_READWRITE);
73 | memcpy(pBuffer,data,data_size);
74 | CreateThread(NULL,0,(LPTHREAD_START_ROUTINE)pBuffer,NULL,0,NULL);
75 |
76 | }
77 |
78 | void thread_with_hw_br(void) {
79 |
80 | DWORD tid = 0;
81 | CONTEXT context ={0};
82 | context.ContextFlags = CONTEXT_DEBUG_REGISTERS;
83 | HANDLE hThread = CreateThread(NULL,0,(LPTHREAD_START_ROUTINE)blocking_function,NULL,CREATE_SUSPENDED,&tid);
84 | if(!GetThreadContext(hThread,&context)) {
85 | return;
86 | }
87 |
88 | context.Dr0 = (DWORD_PTR)blocking_function;
89 | context.Dr1 = (DWORD_PTR)blocking_function;
90 | context.Dr2 = (DWORD_PTR)blocking_function;
91 | context.Dr3 = (DWORD_PTR)blocking_function;
92 | context.Dr7 |= (1 << 0);
93 | context.Dr7 &= ~(1 << 16);
94 | context.Dr7 &= ~(1 << 17);
95 |
96 | context.ContextFlags = CONTEXT_DEBUG_REGISTERS;
97 | if(!SetThreadContext(hThread,&context)) {
98 | return;
99 | }
100 |
101 | }
102 |
103 | TEST_CLASS(scan_tests)
104 | {
105 | public:
106 |
107 | using process_enumerator = hsb::containers::process_enumerator;
108 | using process_scanner = hsb::scanning::process_scanner;
109 | using process = hsb::containers::process;
110 |
111 | scan_tests() {
112 | exec_msgbox_shellcode();
113 | queue_blocking_apc();
114 | thread_with_hw_br();
115 | queue_blocking_timers();
116 | queue_future_timer();
117 | }
118 |
119 | ~scan_tests() { }
120 |
121 | TEST_METHOD(future_timer)
122 | {
123 |
124 | process_enumerator process_enumerator((uint16_t)GetCurrentProcessId(),false);
125 | process_scanner process_scanner;
126 | std::vector> processes;
127 |
128 | processes = process_enumerator.enumerate_processes();
129 | Assert::AreEqual((int)processes.size(),1);
130 |
131 | hsb::scanning::process_scans::suspicious_timer(processes[0].get());
132 |
133 | Assert::AreNotEqual((int)processes[0]->detections.size(),0);
134 | }
135 |
136 | TEST_METHOD(blocking_apc)
137 | {
138 |
139 | process_enumerator process_enumerator((uint16_t)GetCurrentProcessId(), false);
140 | process_scanner process_scanner;
141 | std::vector> processes;
142 |
143 | processes = process_enumerator.enumerate_processes();
144 | Assert::AreEqual((int)processes.size(), 1);
145 |
146 | for (auto& thread : processes[0]->threads) {
147 | hsb::scanning::thread_scans::blocking_apc(processes[0].get(), thread.get());
148 | }
149 |
150 | Assert::AreNotEqual((int)processes[0]->detections.size(),0);
151 |
152 | }
153 |
154 | TEST_METHOD(blocking_timer)
155 | {
156 |
157 | process_enumerator process_enumerator((uint16_t)GetCurrentProcessId(), false);
158 | process_scanner process_scanner;
159 | std::vector> processes;
160 |
161 | processes = process_enumerator.enumerate_processes();
162 | Assert::AreEqual((int)processes.size(), 1);
163 |
164 | for (auto& thread : processes[0]->threads) {
165 | hsb::scanning::thread_scans::blocking_timer(processes[0].get(), thread.get());
166 | }
167 |
168 | Assert::AreNotEqual((int)processes[0]->detections.size(), 0);
169 | }
170 |
171 | TEST_METHOD(private_memory)
172 | {
173 |
174 | process_enumerator process_enumerator((uint16_t)GetCurrentProcessId(), false);
175 | process_scanner process_scanner;
176 | std::vector> processes;
177 |
178 | processes = process_enumerator.enumerate_processes();
179 | Assert::AreEqual((int)processes.size(), 1);
180 |
181 | for (auto& thread : processes[0]->threads) {
182 | hsb::scanning::thread_scans::private_memory(processes[0].get(), thread.get());
183 | }
184 |
185 | Assert::AreNotEqual((int)processes[0]->detections.size(), 0);
186 | }
187 |
188 | TEST_METHOD(hw_breakpoints)
189 | {
190 |
191 | process_enumerator process_enumerator((uint16_t)GetCurrentProcessId(), false);
192 | process_scanner process_scanner;
193 | std::vector> processes;
194 |
195 | processes = process_enumerator.enumerate_processes();
196 | Assert::AreEqual((int)processes.size(), 1);
197 |
198 | for (auto& thread : processes[0]->threads) {
199 | hsb::scanning::thread_scans::hardware_breakpoints(processes[0].get(), thread.get());
200 | }
201 |
202 | Assert::AreNotEqual((int)processes[0]->detections.size(), 0);
203 | }
204 |
205 | TEST_METHOD(intermodular_call)
206 | {
207 |
208 | process_enumerator process_enumerator((uint16_t)GetCurrentProcessId(),false);
209 | process_scanner process_scanner;
210 | std::vector> processes;
211 |
212 | processes = process_enumerator.enumerate_processes();
213 | Assert::AreEqual((int)processes.size(),1);
214 |
215 | for(auto& thread : processes[0]->threads) {
216 | hsb::scanning::thread_scans::abnormal_intermodular_call(processes[0].get(),thread.get());
217 | }
218 |
219 | Assert::AreNotEqual((int)processes[0]->detections.size(),0);
220 | }
221 |
222 | TEST_METHOD(stomped_module)
223 | {
224 |
225 | process_enumerator process_enumerator((uint16_t)GetCurrentProcessId(),false);
226 | process_scanner process_scanner;
227 | std::vector> processes;
228 |
229 | processes = process_enumerator.enumerate_processes();
230 | Assert::AreEqual((int)processes.size(),1);
231 |
232 | for(auto& thread : processes[0]->threads) {
233 | hsb::scanning::thread_scans::stomped_module(processes[0].get(),thread.get());
234 | }
235 |
236 | Assert::AreNotEqual((int)processes[0]->detections.size(),0);
237 | }
238 |
239 | TEST_METHOD(non_executable_memory)
240 | {
241 |
242 | process_enumerator process_enumerator((uint16_t)GetCurrentProcessId(),false);
243 | process_scanner process_scanner;
244 | std::vector> processes;
245 |
246 | processes = process_enumerator.enumerate_processes();
247 | Assert::AreEqual((int)processes.size(),1);
248 |
249 | for(auto& thread : processes[0]->threads) {
250 | hsb::scanning::thread_scans::non_executable_memory(processes[0].get(),thread.get());
251 | }
252 |
253 | Assert::AreNotEqual((int)processes[0]->detections.size(),0);
254 | }
255 |
256 | TEST_METHOD(return_addr_spoofing)
257 | {
258 |
259 | process_enumerator process_enumerator((uint16_t)GetCurrentProcessId(),false);
260 | process_scanner process_scanner;
261 | std::vector> processes;
262 |
263 | processes = process_enumerator.enumerate_processes();
264 | Assert::AreEqual((int)processes.size(),1);
265 |
266 | for(auto& thread : processes[0]->threads) {
267 | hsb::scanning::thread_scans::return_address_spoofing(processes[0].get(),thread.get());
268 | }
269 |
270 | Assert::AreNotEqual((int)processes[0]->detections.size(),0);
271 | }
272 |
273 | };
274 | }
275 |
--------------------------------------------------------------------------------
/src/Hunt-Sleeping-Beacons/Hunt-Sleeping-Beacons.vcxproj:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 | Debug
6 | Win32
7 |
8 |
9 | Release
10 | Win32
11 |
12 |
13 | Debug
14 | x64
15 |
16 |
17 | Release
18 | x64
19 |
20 |
21 | Debug
22 | ARM64
23 |
24 |
25 | Release
26 | ARM64
27 |
28 |
29 |
30 | 17.0
31 | Win32Proj
32 | {70a45dbf-68fe-4d6b-864e-d5e239340047}
33 | HuntSleepingBeacons
34 | 10.0
35 |
36 |
37 |
38 | Application
39 | true
40 | v143
41 | Unicode
42 |
43 |
44 | Application
45 | false
46 | v143
47 | true
48 | Unicode
49 |
50 |
51 | Application
52 | true
53 | v143
54 | Unicode
55 |
56 |
57 | Application
58 | false
59 | v143
60 | true
61 | Unicode
62 |
63 |
64 | Application
65 | true
66 | v143
67 | Unicode
68 |
69 |
70 | Application
71 | false
72 | v143
73 | true
74 | Unicode
75 |
76 |
77 |
78 |
79 |
80 |
81 |
82 |
83 |
84 |
85 |
86 |
87 |
88 |
89 |
90 |
91 |
92 |
93 |
94 |
95 |
96 |
97 |
98 |
99 |
100 |
101 |
102 | Level3
103 | true
104 | WIN32;_DEBUG;_CONSOLE;%(PreprocessorDefinitions)
105 | true
106 | stdcpp20
107 | $(ProjectDir)..\..\inc;%(AdditionalIncludeDirectories)
108 |
109 |
110 | Console
111 | true
112 | $(CoreLibraryDependencies);%(AdditionalDependencies);ntdll.lib;dbghelp.lib;psapi.lib;
113 |
114 |
115 |
116 |
117 | Level3
118 | true
119 | true
120 | true
121 | WIN32;NDEBUG;_CONSOLE;%(PreprocessorDefinitions)
122 | true
123 | stdcpp20
124 | $(ProjectDir)..\..\inc;%(AdditionalIncludeDirectories)
125 | MultiThreaded
126 |
127 |
128 | Console
129 | true
130 | true
131 | true
132 | $(CoreLibraryDependencies);%(AdditionalDependencies);ntdll.lib;dbghelp.lib;psapi.lib;
133 |
134 |
135 |
136 |
137 | Level3
138 | true
139 | _DEBUG;_CONSOLE;%(PreprocessorDefinitions)
140 | true
141 | stdcpp20
142 | $(ProjectDir)..\..\inc;%(AdditionalIncludeDirectories)
143 |
144 |
145 | Console
146 | true
147 | $(CoreLibraryDependencies);%(AdditionalDependencies);ntdll.lib;dbghelp.lib;psapi.lib;Shlwapi.lib;
148 |
149 |
150 |
151 |
152 | Level3
153 | true
154 | true
155 | true
156 | NDEBUG;_CONSOLE;%(PreprocessorDefinitions)
157 | true
158 | stdcpp20
159 | $(ProjectDir)..\..\inc;%(AdditionalIncludeDirectories)
160 | MultiThreaded
161 |
162 |
163 | Console
164 | true
165 | true
166 | true
167 | $(CoreLibraryDependencies);%(AdditionalDependencies);ntdll.lib;dbghelp.lib;psapi.lib;Shlwapi.lib;
168 |
169 |
170 |
171 |
172 | Level3
173 | true
174 | _DEBUG;_CONSOLE;%(PreprocessorDefinitions)
175 | true
176 | stdcpp20
177 | $(ProjectDir)..\..\inc;%(AdditionalIncludeDirectories)
178 |
179 |
180 | Console
181 | true
182 | $(CoreLibraryDependencies);%(AdditionalDependencies);ntdll.lib;dbghelp.lib;psapi.lib;
183 |
184 |
185 |
186 |
187 | Level3
188 | true
189 | true
190 | true
191 | NDEBUG;_CONSOLE;%(PreprocessorDefinitions)
192 | true
193 | stdcpp20
194 | $(ProjectDir)..\..\inc;%(AdditionalIncludeDirectories)
195 | MultiThreaded
196 |
197 |
198 | Console
199 | true
200 | true
201 | true
202 | $(CoreLibraryDependencies);%(AdditionalDependencies);ntdll.lib;dbghelp.lib;psapi.lib;
203 |
204 |
205 |
206 |
207 |
208 |
209 |
210 |
211 |
212 |
213 |
214 |
215 |
216 |
217 |
218 |
219 |
220 |
221 |
222 |
223 |
224 |
225 |
226 |
227 |
228 |
229 |
230 |
231 |
232 |
233 |
234 |
--------------------------------------------------------------------------------
/src/UnitTests/UnitTests.vcxproj:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 | Debug
6 | Win32
7 |
8 |
9 | Release
10 | Win32
11 |
12 |
13 | Debug
14 | x64
15 |
16 |
17 | Release
18 | x64
19 |
20 |
21 | Debug
22 | ARM64
23 |
24 |
25 | Release
26 | ARM64
27 |
28 |
29 |
30 | 17.0
31 | {47D29CEF-58D6-418E-AC77-397E90D40748}
32 | Win32Proj
33 | UnitTests
34 | 10.0
35 | NativeUnitTestProject
36 |
37 |
38 |
39 | DynamicLibrary
40 | true
41 | v143
42 | Unicode
43 | false
44 |
45 |
46 | DynamicLibrary
47 | false
48 | v143
49 | true
50 | Unicode
51 | false
52 |
53 |
54 | DynamicLibrary
55 | true
56 | v143
57 | Unicode
58 | false
59 |
60 |
61 | DynamicLibrary
62 | false
63 | v143
64 | true
65 | Unicode
66 | false
67 |
68 |
69 | DynamicLibrary
70 | true
71 | v143
72 | Unicode
73 | false
74 |
75 |
76 | DynamicLibrary
77 | false
78 | v143
79 | true
80 | Unicode
81 | false
82 |
83 |
84 |
85 |
86 |
87 |
88 |
89 |
90 |
91 |
92 |
93 |
94 |
95 |
96 |
97 |
98 |
99 |
100 |
101 |
102 |
103 |
104 |
105 |
106 |
107 |
108 |
109 | true
110 |
111 |
112 | true
113 |
114 |
115 | true
116 |
117 |
118 | false
119 |
120 |
121 | false
122 |
123 |
124 | false
125 |
126 |
127 |
128 | Use
129 | Level3
130 | true
131 | $(VCInstallDir)UnitTest\include;%(AdditionalIncludeDirectories);../Hunt-Sleeping-Beacons;$(ProjectDir)..\..\inc;
132 | _DEBUG;%(PreprocessorDefinitions)
133 | true
134 | pch.h
135 | stdcpp20
136 |
137 |
138 | Windows
139 | $(VCInstallDir)UnitTest\lib;%(AdditionalLibraryDirectories)
140 | $(CoreLibraryDependencies);%(AdditionalDependencies);ntdll.lib;dbghelp.lib;psapi.lib;
141 |
142 |
143 |
144 |
145 | Use
146 | Level3
147 | true
148 | $(VCInstallDir)UnitTest\include;%(AdditionalIncludeDirectories);../Hunt-Sleeping-Beacons;$(ProjectDir)..\..\inc;
149 | WIN32;_DEBUG;%(PreprocessorDefinitions)
150 | true
151 | pch.h
152 | stdcpp20
153 |
154 |
155 | Windows
156 | $(VCInstallDir)UnitTest\lib;%(AdditionalLibraryDirectories)
157 | $(CoreLibraryDependencies);%(AdditionalDependencies);ntdll.lib;dbghelp.lib;psapi.lib;
158 |
159 |
160 |
161 |
162 | Use
163 | Level3
164 | true
165 | $(VCInstallDir)UnitTest\include;%(AdditionalIncludeDirectories);../Hunt-Sleeping-Beacons;$(ProjectDir)..\..\inc;
166 | _DEBUG;%(PreprocessorDefinitions)
167 | true
168 | pch.h
169 | stdcpp20
170 |
171 |
172 | Windows
173 | $(VCInstallDir)UnitTest\lib;%(AdditionalLibraryDirectories)
174 | $(CoreLibraryDependencies);%(AdditionalDependencies);ntdll.lib;dbghelp.lib;psapi.lib;
175 |
176 |
177 |
178 |
179 | Use
180 | Level3
181 | true
182 | true
183 | true
184 | $(VCInstallDir)UnitTest\include;%(AdditionalIncludeDirectories);../Hunt-Sleeping-Beacons;$(ProjectDir)..\..\inc;
185 | WIN32;NDEBUG;%(PreprocessorDefinitions)
186 | true
187 | pch.h
188 | stdcpp20
189 |
190 |
191 | Windows
192 | true
193 | true
194 | $(VCInstallDir)UnitTest\lib;%(AdditionalLibraryDirectories)
195 | $(CoreLibraryDependencies);%(AdditionalDependencies);ntdll.lib;dbghelp.lib;psapi.lib;
196 |
197 |
198 |
199 |
200 | Use
201 | Level3
202 | true
203 | true
204 | true
205 | $(VCInstallDir)UnitTest\include;%(AdditionalIncludeDirectories);../Hunt-Sleeping-Beacons;$(ProjectDir)..\..\inc;
206 | NDEBUG;%(PreprocessorDefinitions)
207 | true
208 | pch.h
209 | stdcpp20
210 |
211 |
212 | Windows
213 | true
214 | true
215 | $(VCInstallDir)UnitTest\lib;%(AdditionalLibraryDirectories)
216 | $(CoreLibraryDependencies);%(AdditionalDependencies);ntdll.lib;dbghelp.lib;psapi.lib;
217 |
218 |
219 |
220 |
221 | Use
222 | Level3
223 | true
224 | true
225 | true
226 | $(VCInstallDir)UnitTest\include;%(AdditionalIncludeDirectories);../Hunt-Sleeping-Beacons;$(ProjectDir)..\..\inc;
227 | NDEBUG;%(PreprocessorDefinitions)
228 | true
229 | pch.h
230 | stdcpp20
231 |
232 |
233 | Windows
234 | true
235 | true
236 | $(VCInstallDir)UnitTest\lib;%(AdditionalLibraryDirectories)
237 | $(CoreLibraryDependencies);%(AdditionalDependencies);ntdll.lib;dbghelp.lib;psapi.lib;
238 |
239 |
240 |
241 |
242 | Create
243 | Create
244 | Create
245 | Create
246 | Create
247 | Create
248 |
249 |
250 |
251 |
252 |
253 |
254 |
255 |
256 | {70a45dbf-68fe-4d6b-864e-d5e239340047}
257 |
258 |
259 |
260 |
261 | Document
262 | false
263 | false
264 |
265 |
266 |
267 |
268 |
269 |
270 |
--------------------------------------------------------------------------------
/inc/BS_thread_pool.hpp:
--------------------------------------------------------------------------------
1 | #ifndef BS_THREAD_POOL_HPP
2 | #define BS_THREAD_POOL_HPP
3 | /**
4 | * @file BS_thread_pool.hpp
5 | * @author Barak Shoshany (baraksh@gmail.com) (https://baraksh.com)
6 | * @version 4.1.0
7 | * @date 2024-03-22
8 | * @copyright Copyright (c) 2024 Barak Shoshany. Licensed under the MIT license. If you found this project useful, please consider starring it on GitHub! If you use this library in software of any kind, please provide a link to the GitHub repository https://github.com/bshoshany/thread-pool in the source code and documentation. If you use this library in published research, please cite it as follows: Barak Shoshany, "A C++17 Thread Pool for High-Performance Scientific Computing", doi:10.1016/j.softx.2024.101687, SoftwareX 26 (2024) 101687, arXiv:2105.00613
9 | *
10 | * @brief BS::thread_pool: a fast, lightweight, and easy-to-use C++17 thread pool library. This header file contains the main thread pool class and some additional classes and definitions. No other files are needed in order to use the thread pool itself.
11 | */
12 |
13 | #ifndef __cpp_exceptions
14 | #define BS_THREAD_POOL_DISABLE_EXCEPTION_HANDLING
15 | #undef BS_THREAD_POOL_ENABLE_WAIT_DEADLOCK_CHECK
16 | #endif
17 |
18 | #include // std::chrono
19 | #include // std::condition_variable
20 | #include // std::size_t
21 | #ifdef BS_THREAD_POOL_ENABLE_PRIORITY
22 | #include // std::int_least16_t
23 | #endif
24 | #ifndef BS_THREAD_POOL_DISABLE_EXCEPTION_HANDLING
25 | #include // std::current_exception
26 | #endif
27 | #include // std::function
28 | #include // std::future, std::future_status, std::promise
29 | #include // std::make_shared, std::make_unique, std::shared_ptr, std::unique_ptr
30 | #include // std::mutex, std::scoped_lock, std::unique_lock
31 | #include // std::nullopt, std::optional
32 | #include // std::priority_queue (if priority enabled), std::queue
33 | #ifdef BS_THREAD_POOL_ENABLE_WAIT_DEADLOCK_CHECK
34 | #include // std::runtime_error
35 | #endif
36 | #include // std::thread
37 | #include // std::conditional_t, std::decay_t, std::invoke_result_t, std::is_void_v, std::remove_const_t (if priority enabled)
38 | #include // std::forward, std::move
39 | #include // std::vector
40 |
41 | /**
42 | * @brief A namespace used by Barak Shoshany's projects.
43 | */
44 | namespace BS {
45 | // Macros indicating the version of the thread pool library.
46 | #define BS_THREAD_POOL_VERSION_MAJOR 4
47 | #define BS_THREAD_POOL_VERSION_MINOR 1
48 | #define BS_THREAD_POOL_VERSION_PATCH 0
49 |
50 | class thread_pool;
51 |
52 | /**
53 | * @brief A type to represent the size of things.
54 | */
55 | using size_t = std::size_t;
56 |
57 | /**
58 | * @brief A convenient shorthand for the type of `std::thread::hardware_concurrency()`. Should evaluate to unsigned int.
59 | */
60 | using concurrency_t = std::invoke_result_t;
61 |
62 | #ifdef BS_THREAD_POOL_ENABLE_PRIORITY
63 | /**
64 | * @brief A type used to indicate the priority of a task. Defined to be an integer with a width of (at least) 16 bits.
65 | */
66 | using priority_t = std::int_least16_t;
67 |
68 | /**
69 | * @brief A namespace containing some pre-defined priorities for convenience.
70 | */
71 | namespace pr {
72 | constexpr priority_t highest = 32767;
73 | constexpr priority_t high = 16383;
74 | constexpr priority_t normal = 0;
75 | constexpr priority_t low = -16384;
76 | constexpr priority_t lowest = -32768;
77 | } // namespace pr
78 |
79 | // Macros used internally to enable or disable the priority arguments in the relevant functions.
80 | #define BS_THREAD_POOL_PRIORITY_INPUT , const priority_t priority = 0
81 | #define BS_THREAD_POOL_PRIORITY_OUTPUT , priority
82 | #else
83 | #define BS_THREAD_POOL_PRIORITY_INPUT
84 | #define BS_THREAD_POOL_PRIORITY_OUTPUT
85 | #endif
86 |
87 | /**
88 | * @brief A namespace used to obtain information about the current thread.
89 | */
90 | namespace this_thread {
91 | /**
92 | * @brief A type returned by `BS::this_thread::get_index()` which can optionally contain the index of a thread, if that thread belongs to a `BS::thread_pool`. Otherwise, it will contain no value.
93 | */
94 | using optional_index = std::optional;
95 |
96 | /**
97 | * @brief A type returned by `BS::this_thread::get_pool()` which can optionally contain the pointer to the pool that owns a thread, if that thread belongs to a `BS::thread_pool`. Otherwise, it will contain no value.
98 | */
99 | using optional_pool = std::optional;
100 |
101 | /**
102 | * @brief A helper class to store information about the index of the current thread.
103 | */
104 | class [[nodiscard]] thread_info_index
105 | {
106 | friend class BS::thread_pool;
107 |
108 | public:
109 | /**
110 | * @brief Get the index of the current thread. If this thread belongs to a `BS::thread_pool` object, it will have an index from 0 to `BS::thread_pool::get_thread_count() - 1`. Otherwise, for example if this thread is the main thread or an independent `std::thread`, `std::nullopt` will be returned.
111 | *
112 | * @return An `std::optional` object, optionally containing a thread index. Unless you are 100% sure this thread is in a pool, first use `std::optional::has_value()` to check if it contains a value, and if so, use `std::optional::value()` to obtain that value.
113 | */
114 | [[nodiscard]] optional_index operator()() const
115 | {
116 | return index;
117 | }
118 |
119 | private:
120 | /**
121 | * @brief The index of the current thread.
122 | */
123 | optional_index index = std::nullopt;
124 | }; // class thread_info_index
125 |
126 | /**
127 | * @brief A helper class to store information about the thread pool that owns the current thread.
128 | */
129 | class [[nodiscard]] thread_info_pool
130 | {
131 | friend class BS::thread_pool;
132 |
133 | public:
134 | /**
135 | * @brief Get the pointer to the thread pool that owns the current thread. If this thread belongs to a `BS::thread_pool` object, a pointer to that object will be returned. Otherwise, for example if this thread is the main thread or an independent `std::thread`, `std::nullopt` will be returned.
136 | *
137 | * @return An `std::optional` object, optionally containing a pointer to a thread pool. Unless you are 100% sure this thread is in a pool, first use `std::optional::has_value()` to check if it contains a value, and if so, use `std::optional::value()` to obtain that value.
138 | */
139 | [[nodiscard]] optional_pool operator()() const
140 | {
141 | return pool;
142 | }
143 |
144 | private:
145 | /**
146 | * @brief A pointer to the thread pool that owns the current thread.
147 | */
148 | optional_pool pool = std::nullopt;
149 | }; // class thread_info_pool
150 |
151 | /**
152 | * @brief A `thread_local` object used to obtain information about the index of the current thread.
153 | */
154 | inline thread_local thread_info_index get_index;
155 |
156 | /**
157 | * @brief A `thread_local` object used to obtain information about the thread pool that owns the current thread.
158 | */
159 | inline thread_local thread_info_pool get_pool;
160 | } // namespace this_thread
161 |
162 | /**
163 | * @brief A helper class to facilitate waiting for and/or getting the results of multiple futures at once.
164 | *
165 | * @tparam T The return type of the futures.
166 | */
167 | template
168 | class [[nodiscard]] multi_future: public std::vector>
169 | {
170 | public:
171 | // Inherit all constructors from the base class `std::vector`.
172 | using std::vector>::vector;
173 |
174 | // The copy constructor and copy assignment operator are deleted. The elements stored in a `multi_future` are futures, which cannot be copied.
175 | multi_future(const multi_future&) = delete;
176 | multi_future& operator=(const multi_future&) = delete;
177 |
178 | // The move constructor and move assignment operator are defaulted.
179 | multi_future(multi_future&&) = default;
180 | multi_future& operator=(multi_future&&) = default;
181 |
182 | /**
183 | * @brief Get the results from all the futures stored in this `multi_future`, rethrowing any stored exceptions.
184 | *
185 | * @return If the futures return `void`, this function returns `void` as well. Otherwise, it returns a vector containing the results.
186 | */
187 | [[nodiscard]] std::conditional_t,void,std::vector> get()
188 | {
189 | if constexpr(std::is_void_v)
190 | {
191 | for(std::future& future : *this)
192 | future.get();
193 | return;
194 | }
195 | else
196 | {
197 | std::vector results;
198 | results.reserve(this->size());
199 | for(std::future& future : *this)
200 | results.push_back(future.get());
201 | return results;
202 | }
203 | }
204 |
205 | /**
206 | * @brief Check how many of the futures stored in this `multi_future` are ready.
207 | *
208 | * @return The number of ready futures.
209 | */
210 | [[nodiscard]] size_t ready_count() const
211 | {
212 | size_t count = 0;
213 | for(const std::future& future : *this)
214 | {
215 | if(future.wait_for(std::chrono::duration::zero()) == std::future_status::ready)
216 | ++count;
217 | }
218 | return count;
219 | }
220 |
221 | /**
222 | * @brief Check if all the futures stored in this `multi_future` are valid.
223 | *
224 | * @return `true` if all futures are valid, `false` if at least one of the futures is not valid.
225 | */
226 | [[nodiscard]] bool valid() const
227 | {
228 | bool is_valid = true;
229 | for(const std::future& future : *this)
230 | is_valid = is_valid && future.valid();
231 | return is_valid;
232 | }
233 |
234 | /**
235 | * @brief Wait for all the futures stored in this `multi_future`.
236 | */
237 | void wait() const
238 | {
239 | for(const std::future& future : *this)
240 | future.wait();
241 | }
242 |
243 | /**
244 | * @brief Wait for all the futures stored in this `multi_future`, but stop waiting after the specified duration has passed. This function first waits for the first future for the desired duration. If that future is ready before the duration expires, this function waits for the second future for whatever remains of the duration. It continues similarly until the duration expires.
245 | *
246 | * @tparam R An arithmetic type representing the number of ticks to wait.
247 | * @tparam P An `std::ratio` representing the length of each tick in seconds.
248 | * @param duration The amount of time to wait.
249 | * @return `true` if all futures have been waited for before the duration expired, `false` otherwise.
250 | */
251 | template
252 | bool wait_for(const std::chrono::duration& duration) const
253 | {
254 | const std::chrono::time_point start_time = std::chrono::steady_clock::now();
255 | for(const std::future& future : *this)
256 | {
257 | future.wait_for(duration - (std::chrono::steady_clock::now() - start_time));
258 | if(duration < std::chrono::steady_clock::now() - start_time)
259 | return false;
260 | }
261 | return true;
262 | }
263 |
264 | /**
265 | * @brief Wait for all the futures stored in this `multi_future`, but stop waiting after the specified time point has been reached. This function first waits for the first future until the desired time point. If that future is ready before the time point is reached, this function waits for the second future until the desired time point. It continues similarly until the time point is reached.
266 | *
267 | * @tparam C The type of the clock used to measure time.
268 | * @tparam D An `std::chrono::duration` type used to indicate the time point.
269 | * @param timeout_time The time point at which to stop waiting.
270 | * @return `true` if all futures have been waited for before the time point was reached, `false` otherwise.
271 | */
272 | template
273 | bool wait_until(const std::chrono::time_point& timeout_time) const
274 | {
275 | for(const std::future& future : *this)
276 | {
277 | future.wait_until(timeout_time);
278 | if(timeout_time < std::chrono::steady_clock::now())
279 | return false;
280 | }
281 | return true;
282 | }
283 | }; // class multi_future
284 |
285 | /**
286 | * @brief A fast, lightweight, and easy-to-use C++17 thread pool class.
287 | */
288 | class [[nodiscard]] thread_pool
289 | {
290 | public:
291 | // ============================
292 | // Constructors and destructors
293 | // ============================
294 |
295 | /**
296 | * @brief Construct a new thread pool. The number of threads will be the total number of hardware threads available, as reported by the implementation. This is usually determined by the number of cores in the CPU. If a core is hyperthreaded, it will count as two threads.
297 | */
298 | thread_pool(): thread_pool(0,[] {}) {}
299 |
300 | /**
301 | * @brief Construct a new thread pool with the specified number of threads.
302 | *
303 | * @param num_threads The number of threads to use.
304 | */
305 | explicit thread_pool(const concurrency_t num_threads): thread_pool(num_threads,[] {}) {}
306 |
307 | /**
308 | * @brief Construct a new thread pool with the specified initialization function.
309 | *
310 | * @param init_task An initialization function to run in each thread before it starts to execute any submitted tasks. The function must take no arguments and have no return value. It will only be executed exactly once, when the thread is first constructed.
311 | */
312 | explicit thread_pool(const std::function& init_task): thread_pool(0,init_task) {}
313 |
314 | /**
315 | * @brief Construct a new thread pool with the specified number of threads and initialization function.
316 | *
317 | * @param num_threads The number of threads to use.
318 | * @param init_task An initialization function to run in each thread before it starts to execute any submitted tasks. The function must take no arguments and have no return value. It will only be executed exactly once, when the thread is first constructed.
319 | */
320 | thread_pool(const concurrency_t num_threads,const std::function& init_task): thread_count(determine_thread_count(num_threads)),threads(std::make_unique(determine_thread_count(num_threads)))
321 | {
322 | create_threads(init_task);
323 | }
324 |
325 | // The copy and move constructors and assignment operators are deleted. The thread pool uses a mutex, which cannot be copied or moved.
326 | thread_pool(const thread_pool&) = delete;
327 | thread_pool(thread_pool&&) = delete;
328 | thread_pool& operator=(const thread_pool&) = delete;
329 | thread_pool& operator=(thread_pool&&) = delete;
330 |
331 | /**
332 | * @brief Destruct the thread pool. Waits for all tasks to complete, then destroys all threads. Note that if the pool is paused, then any tasks still in the queue will never be executed.
333 | */
334 | ~thread_pool()
335 | {
336 | wait();
337 | destroy_threads();
338 | }
339 |
340 | // =======================
341 | // Public member functions
342 | // =======================
343 |
344 | #ifdef BS_THREAD_POOL_ENABLE_NATIVE_HANDLES
345 | /**
346 | * @brief Get a vector containing the underlying implementation-defined thread handles for each of the pool's threads, as obtained by `std::thread::native_handle()`. Only enabled if `BS_THREAD_POOL_ENABLE_NATIVE_HANDLES` is defined.
347 | *
348 | * @return The native thread handles.
349 | */
350 | [[nodiscard]] std::vector get_native_handles() const
351 | {
352 | std::vector native_handles(thread_count);
353 | for(concurrency_t i = 0; i < thread_count; ++i)
354 | {
355 | native_handles[i] = threads[i].native_handle();
356 | }
357 | return native_handles;
358 | }
359 | #endif
360 |
361 | /**
362 | * @brief Get the number of tasks currently waiting in the queue to be executed by the threads.
363 | *
364 | * @return The number of queued tasks.
365 | */
366 | [[nodiscard]] size_t get_tasks_queued() const
367 | {
368 | const std::scoped_lock tasks_lock(tasks_mutex);
369 | return tasks.size();
370 | }
371 |
372 | /**
373 | * @brief Get the number of tasks currently being executed by the threads.
374 | *
375 | * @return The number of running tasks.
376 | */
377 | [[nodiscard]] size_t get_tasks_running() const
378 | {
379 | const std::scoped_lock tasks_lock(tasks_mutex);
380 | return tasks_running;
381 | }
382 |
383 | /**
384 | * @brief Get the total number of unfinished tasks: either still waiting in the queue, or running in a thread. Note that `get_tasks_total() == get_tasks_queued() + get_tasks_running()`.
385 | *
386 | * @return The total number of tasks.
387 | */
388 | [[nodiscard]] size_t get_tasks_total() const
389 | {
390 | const std::scoped_lock tasks_lock(tasks_mutex);
391 | return tasks_running + tasks.size();
392 | }
393 |
394 | /**
395 | * @brief Get the number of threads in the pool.
396 | *
397 | * @return The number of threads.
398 | */
399 | [[nodiscard]] concurrency_t get_thread_count() const
400 | {
401 | return thread_count;
402 | }
403 |
404 | /**
405 | * @brief Get a vector containing the unique identifiers for each of the pool's threads, as obtained by `std::thread::get_id()`.
406 | *
407 | * @return The unique thread identifiers.
408 | */
409 | [[nodiscard]] std::vector get_thread_ids() const
410 | {
411 | std::vector thread_ids(thread_count);
412 | for(concurrency_t i = 0; i < thread_count; ++i)
413 | {
414 | thread_ids[i] = threads[i].get_id();
415 | }
416 | return thread_ids;
417 | }
418 |
419 | #ifdef BS_THREAD_POOL_ENABLE_PAUSE
420 | /**
421 | * @brief Check whether the pool is currently paused. Only enabled if `BS_THREAD_POOL_ENABLE_PAUSE` is defined.
422 | *
423 | * @return `true` if the pool is paused, `false` if it is not paused.
424 | */
425 | [[nodiscard]] bool is_paused() const
426 | {
427 | const std::scoped_lock tasks_lock(tasks_mutex);
428 | return paused;
429 | }
430 |
431 | /**
432 | * @brief Pause the pool. The workers will temporarily stop retrieving new tasks out of the queue, although any tasks already executed will keep running until they are finished. Only enabled if `BS_THREAD_POOL_ENABLE_PAUSE` is defined.
433 | */
434 | void pause()
435 | {
436 | const std::scoped_lock tasks_lock(tasks_mutex);
437 | paused = true;
438 | }
439 | #endif
440 |
441 | /**
442 | * @brief Purge all the tasks waiting in the queue. Tasks that are currently running will not be affected, but any tasks still waiting in the queue will be discarded, and will never be executed by the threads. Please note that there is no way to restore the purged tasks.
443 | */
444 | void purge()
445 | {
446 | const std::scoped_lock tasks_lock(tasks_mutex);
447 | while(!tasks.empty())
448 | tasks.pop();
449 | }
450 |
451 | /**
452 | * @brief Submit a function with no arguments and no return value into the task queue, with the specified priority. To push a function with arguments, enclose it in a lambda expression. Does not return a future, so the user must use `wait()` or some other method to ensure that the task finishes executing, otherwise bad things will happen.
453 | *
454 | * @tparam F The type of the function.
455 | * @param task The function to push.
456 | * @param priority The priority of the task. Should be between -32,768 and 32,767 (a signed 16-bit integer). The default is 0. Only enabled if `BS_THREAD_POOL_ENABLE_PRIORITY` is defined.
457 | */
458 | template
459 | void detach_task(F&& task BS_THREAD_POOL_PRIORITY_INPUT)
460 | {
461 | {
462 | const std::scoped_lock tasks_lock(tasks_mutex);
463 | tasks.emplace(std::forward(task) BS_THREAD_POOL_PRIORITY_OUTPUT);
464 | }
465 | task_available_cv.notify_one();
466 | }
467 |
468 | /**
469 | * @brief Parallelize a loop by automatically splitting it into blocks and submitting each block separately to the queue, with the specified priority. The block function takes two arguments, the start and end of the block, so that it is only called only once per block, but it is up to the user make sure the block function correctly deals with all the indices in each block. Does not return a `multi_future`, so the user must use `wait()` or some other method to ensure that the loop finishes executing, otherwise bad things will happen.
470 | *
471 | * @tparam T The type of the indices. Should be a signed or unsigned integer.
472 | * @tparam F The type of the function to loop through.
473 | * @param first_index The first index in the loop.
474 | * @param index_after_last The index after the last index in the loop. The loop will iterate from `first_index` to `(index_after_last - 1)` inclusive. In other words, it will be equivalent to `for (T i = first_index; i < index_after_last; ++i)`. Note that if `index_after_last <= first_index`, no blocks will be submitted.
475 | * @param block A function that will be called once per block. Should take exactly two arguments: the first index in the block and the index after the last index in the block. `block(start, end)` should typically involve a loop of the form `for (T i = start; i < end; ++i)`.
476 | * @param num_blocks The maximum number of blocks to split the loop into. The default is 0, which means the number of blocks will be equal to the number of threads in the pool.
477 | * @param priority The priority of the tasks. Should be between -32,768 and 32,767 (a signed 16-bit integer). The default is 0. Only enabled if `BS_THREAD_POOL_ENABLE_PRIORITY` is defined.
478 | */
479 | template
480 | void detach_blocks(const T first_index,const T index_after_last,F&& block,const size_t num_blocks = 0 BS_THREAD_POOL_PRIORITY_INPUT)
481 | {
482 | if(index_after_last > first_index)
483 | {
484 | const blocks blks(first_index,index_after_last,num_blocks ? num_blocks : thread_count);
485 | for(size_t blk = 0; blk < blks.get_num_blocks(); ++blk)
486 | detach_task(
487 | [block = std::forward(block),start = blks.start(blk),end = blks.end(blk)]
488 | {
489 | block(start,end);
490 | } BS_THREAD_POOL_PRIORITY_OUTPUT);
491 | }
492 | }
493 |
494 | /**
495 | * @brief Parallelize a loop by automatically splitting it into blocks and submitting each block separately to the queue, with the specified priority. The loop function takes one argument, the loop index, so that it is called many times per block. Does not return a `multi_future`, so the user must use `wait()` or some other method to ensure that the loop finishes executing, otherwise bad things will happen.
496 | *
497 | * @tparam T The type of the indices. Should be a signed or unsigned integer.
498 | * @tparam F The type of the function to loop through.
499 | * @param first_index The first index in the loop.
500 | * @param index_after_last The index after the last index in the loop. The loop will iterate from `first_index` to `(index_after_last - 1)` inclusive. In other words, it will be equivalent to `for (T i = first_index; i < index_after_last; ++i)`. Note that if `index_after_last <= first_index`, no blocks will be submitted.
501 | * @param loop The function to loop through. Will be called once per index, many times per block. Should take exactly one argument: the loop index.
502 | * @param num_blocks The maximum number of blocks to split the loop into. The default is 0, which means the number of blocks will be equal to the number of threads in the pool.
503 | * @param priority The priority of the tasks. Should be between -32,768 and 32,767 (a signed 16-bit integer). The default is 0. Only enabled if `BS_THREAD_POOL_ENABLE_PRIORITY` is defined.
504 | */
505 | template
506 | void detach_loop(const T first_index,const T index_after_last,F&& loop,const size_t num_blocks = 0 BS_THREAD_POOL_PRIORITY_INPUT)
507 | {
508 | if(index_after_last > first_index)
509 | {
510 | const blocks blks(first_index,index_after_last,num_blocks ? num_blocks : thread_count);
511 | for(size_t blk = 0; blk < blks.get_num_blocks(); ++blk)
512 | detach_task(
513 | [loop = std::forward(loop),start = blks.start(blk),end = blks.end(blk)]
514 | {
515 | for(T i = start; i < end; ++i)
516 | loop(i);
517 | } BS_THREAD_POOL_PRIORITY_OUTPUT);
518 | }
519 | }
520 |
521 | /**
522 | * @brief Submit a sequence of tasks enumerated by indices to the queue, with the specified priority. Does not return a `multi_future`, so the user must use `wait()` or some other method to ensure that the sequence finishes executing, otherwise bad things will happen.
523 | *
524 | * @tparam T The type of the indices. Should be a signed or unsigned integer.
525 | * @tparam F The type of the function used to define the sequence.
526 | * @param first_index The first index in the sequence.
527 | * @param index_after_last The index after the last index in the sequence. The sequence will iterate from `first_index` to `(index_after_last - 1)` inclusive. In other words, it will be equivalent to `for (T i = first_index; i < index_after_last; ++i)`. Note that if `index_after_last <= first_index`, no tasks will be submitted.
528 | * @param sequence The function used to define the sequence. Will be called once per index. Should take exactly one argument, the index.
529 | * @param priority The priority of the tasks. Should be between -32,768 and 32,767 (a signed 16-bit integer). The default is 0. Only enabled if `BS_THREAD_POOL_ENABLE_PRIORITY` is defined.
530 | */
531 | template
532 | void detach_sequence(const T first_index,const T index_after_last,F&& sequence BS_THREAD_POOL_PRIORITY_INPUT)
533 | {
534 | for(T i = first_index; i < index_after_last; ++i)
535 | detach_task(
536 | [sequence = std::forward(sequence),i]
537 | {
538 | sequence(i);
539 | } BS_THREAD_POOL_PRIORITY_OUTPUT);
540 | }
541 |
542 | /**
543 | * @brief Reset the pool with the total number of hardware threads available, as reported by the implementation. Waits for all currently running tasks to be completed, then destroys all threads in the pool and creates a new thread pool with the new number of threads. Any tasks that were waiting in the queue before the pool was reset will then be executed by the new threads. If the pool was paused before resetting it, the new pool will be paused as well.
544 | */
545 | void reset()
546 | {
547 | reset(0,[] {});
548 | }
549 |
550 | /**
551 | * @brief Reset the pool with a new number of threads. Waits for all currently running tasks to be completed, then destroys all threads in the pool and creates a new thread pool with the new number of threads. Any tasks that were waiting in the queue before the pool was reset will then be executed by the new threads. If the pool was paused before resetting it, the new pool will be paused as well.
552 | *
553 | * @param num_threads The number of threads to use.
554 | */
555 | void reset(const concurrency_t num_threads)
556 | {
557 | reset(num_threads,[] {});
558 | }
559 |
560 | /**
561 | * @brief Reset the pool with the total number of hardware threads available, as reported by the implementation, and a new initialization function. Waits for all currently running tasks to be completed, then destroys all threads in the pool and creates a new thread pool with the new number of threads and initialization function. Any tasks that were waiting in the queue before the pool was reset will then be executed by the new threads. If the pool was paused before resetting it, the new pool will be paused as well.
562 | *
563 | * @param init_task An initialization function to run in each thread before it starts to execute any submitted tasks. The function must take no arguments and have no return value. It will only be executed exactly once, when the thread is first constructed.
564 | */
565 | void reset(const std::function& init_task)
566 | {
567 | reset(0,init_task);
568 | }
569 |
570 | /**
571 | * @brief Reset the pool with a new number of threads and a new initialization function. Waits for all currently running tasks to be completed, then destroys all threads in the pool and creates a new thread pool with the new number of threads and initialization function. Any tasks that were waiting in the queue before the pool was reset will then be executed by the new threads. If the pool was paused before resetting it, the new pool will be paused as well.
572 | *
573 | * @param num_threads The number of threads to use.
574 | * @param init_task An initialization function to run in each thread before it starts to execute any submitted tasks. The function must take no arguments and have no return value. It will only be executed exactly once, when the thread is first constructed.
575 | */
576 | void reset(const concurrency_t num_threads,const std::function& init_task)
577 | {
578 | #ifdef BS_THREAD_POOL_ENABLE_PAUSE
579 | std::unique_lock tasks_lock(tasks_mutex);
580 | const bool was_paused = paused;
581 | paused = true;
582 | tasks_lock.unlock();
583 | #endif
584 | wait();
585 | destroy_threads();
586 | thread_count = determine_thread_count(num_threads);
587 | threads = std::make_unique(thread_count);
588 | create_threads(init_task);
589 | #ifdef BS_THREAD_POOL_ENABLE_PAUSE
590 | tasks_lock.lock();
591 | paused = was_paused;
592 | #endif
593 | }
594 |
595 | /**
596 | * @brief Submit a function with no arguments into the task queue, with the specified priority. To submit a function with arguments, enclose it in a lambda expression. If the function has a return value, get a future for the eventual returned value. If the function has no return value, get an `std::future` which can be used to wait until the task finishes.
597 | *
598 | * @tparam F The type of the function.
599 | * @tparam R The return type of the function (can be `void`).
600 | * @param task The function to submit.
601 | * @param priority The priority of the task. Should be between -32,768 and 32,767 (a signed 16-bit integer). The default is 0. Only enabled if `BS_THREAD_POOL_ENABLE_PRIORITY` is defined.
602 | * @return A future to be used later to wait for the function to finish executing and/or obtain its returned value if it has one.
603 | */
604 | template >>
605 | [[nodiscard]] std::future submit_task(F&& task BS_THREAD_POOL_PRIORITY_INPUT)
606 | {
607 | const std::shared_ptr> task_promise = std::make_shared>();
608 | detach_task(
609 | [task = std::forward(task),task_promise]
610 | {
611 | #ifndef BS_THREAD_POOL_DISABLE_EXCEPTION_HANDLING
612 | try
613 | {
614 | #endif
615 | if constexpr(std::is_void_v)
616 | {
617 | task();
618 | task_promise->set_value();
619 | }
620 | else
621 | {
622 | task_promise->set_value(task());
623 | }
624 | #ifndef BS_THREAD_POOL_DISABLE_EXCEPTION_HANDLING
625 | }
626 | catch(...)
627 | {
628 | try
629 | {
630 | task_promise->set_exception(std::current_exception());
631 | }
632 | catch(...)
633 | {
634 | }
635 | }
636 | #endif
637 | } BS_THREAD_POOL_PRIORITY_OUTPUT);
638 | return task_promise->get_future();
639 | }
640 |
641 | /**
642 | * @brief Parallelize a loop by automatically splitting it into blocks and submitting each block separately to the queue, with the specified priority. The block function takes two arguments, the start and end of the block, so that it is only called only once per block, but it is up to the user make sure the block function correctly deals with all the indices in each block. Returns a `multi_future` that contains the futures for all of the blocks.
643 | *
644 | * @tparam T The type of the indices. Should be a signed or unsigned integer.
645 | * @tparam F The type of the function to loop through.
646 | * @tparam R The return type of the function to loop through (can be `void`).
647 | * @param first_index The first index in the loop.
648 | * @param index_after_last The index after the last index in the loop. The loop will iterate from `first_index` to `(index_after_last - 1)` inclusive. In other words, it will be equivalent to `for (T i = first_index; i < index_after_last; ++i)`. Note that if `index_after_last <= first_index`, no blocks will be submitted, and an empty `multi_future` will be returned.
649 | * @param block A function that will be called once per block. Should take exactly two arguments: the first index in the block and the index after the last index in the block. `block(start, end)` should typically involve a loop of the form `for (T i = start; i < end; ++i)`.
650 | * @param num_blocks The maximum number of blocks to split the loop into. The default is 0, which means the number of blocks will be equal to the number of threads in the pool.
651 | * @param priority The priority of the tasks. Should be between -32,768 and 32,767 (a signed 16-bit integer). The default is 0. Only enabled if `BS_THREAD_POOL_ENABLE_PRIORITY` is defined.
652 | * @return A `multi_future` that can be used to wait for all the blocks to finish. If the block function returns a value, the `multi_future` can also be used to obtain the values returned by each block.
653 | */
654 | template ,T,T>>
655 | [[nodiscard]] multi_future submit_blocks(const T first_index,const T index_after_last,F&& block,const size_t num_blocks = 0 BS_THREAD_POOL_PRIORITY_INPUT)
656 | {
657 | if(index_after_last > first_index)
658 | {
659 | const blocks blks(first_index,index_after_last,num_blocks ? num_blocks : thread_count);
660 | multi_future future;
661 | future.reserve(blks.get_num_blocks());
662 | for(size_t blk = 0; blk < blks.get_num_blocks(); ++blk)
663 | future.push_back(submit_task(
664 | [block = std::forward(block),start = blks.start(blk),end = blks.end(blk)]
665 | {
666 | return block(start,end);
667 | } BS_THREAD_POOL_PRIORITY_OUTPUT));
668 | return future;
669 | }
670 | return {};
671 | }
672 |
673 | /**
674 | * @brief Parallelize a loop by automatically splitting it into blocks and submitting each block separately to the queue, with the specified priority. The loop function takes one argument, the loop index, so that it is called many times per block. It must have no return value. Returns a `multi_future` that contains the futures for all of the blocks.
675 | *
676 | * @tparam T The type of the indices. Should be a signed or unsigned integer.
677 | * @tparam F The type of the function to loop through.
678 | * @param first_index The first index in the loop.
679 | * @param index_after_last The index after the last index in the loop. The loop will iterate from `first_index` to `(index_after_last - 1)` inclusive. In other words, it will be equivalent to `for (T i = first_index; i < index_after_last; ++i)`. Note that if `index_after_last <= first_index`, no tasks will be submitted, and an empty `multi_future` will be returned.
680 | * @param loop The function to loop through. Will be called once per index, many times per block. Should take exactly one argument: the loop index. It cannot have a return value.
681 | * @param num_blocks The maximum number of blocks to split the loop into. The default is 0, which means the number of blocks will be equal to the number of threads in the pool.
682 | * @param priority The priority of the tasks. Should be between -32,768 and 32,767 (a signed 16-bit integer). The default is 0. Only enabled if `BS_THREAD_POOL_ENABLE_PRIORITY` is defined.
683 | * @return A `multi_future` that can be used to wait for all the blocks to finish.
684 | */
685 | template
686 | [[nodiscard]] multi_future submit_loop(const T first_index,const T index_after_last,F&& loop,const size_t num_blocks = 0 BS_THREAD_POOL_PRIORITY_INPUT)
687 | {
688 | if(index_after_last > first_index)
689 | {
690 | const blocks blks(first_index,index_after_last,num_blocks ? num_blocks : thread_count);
691 | multi_future future;
692 | future.reserve(blks.get_num_blocks());
693 | for(size_t blk = 0; blk < blks.get_num_blocks(); ++blk)
694 | future.push_back(submit_task(
695 | [loop = std::forward(loop),start = blks.start(blk),end = blks.end(blk)]
696 | {
697 | for(T i = start; i < end; ++i)
698 | loop(i);
699 | } BS_THREAD_POOL_PRIORITY_OUTPUT));
700 | return future;
701 | }
702 | return {};
703 | }
704 |
705 | /**
706 | * @brief Submit a sequence of tasks enumerated by indices to the queue, with the specified priority. Returns a `multi_future` that contains the futures for all of the tasks.
707 | *
708 | * @tparam T The type of the indices. Should be a signed or unsigned integer.
709 | * @tparam F The type of the function used to define the sequence.
710 | * @tparam R The return type of the function used to define the sequence (can be `void`).
711 | * @param first_index The first index in the sequence.
712 | * @param index_after_last The index after the last index in the sequence. The sequence will iterate from `first_index` to `(index_after_last - 1)` inclusive. In other words, it will be equivalent to `for (T i = first_index; i < index_after_last; ++i)`. Note that if `index_after_last <= first_index`, no tasks will be submitted, and an empty `multi_future` will be returned.
713 | * @param sequence The function used to define the sequence. Will be called once per index. Should take exactly one argument, the index.
714 | * @param priority The priority of the tasks. Should be between -32,768 and 32,767 (a signed 16-bit integer). The default is 0. Only enabled if `BS_THREAD_POOL_ENABLE_PRIORITY` is defined.
715 | * @return A `multi_future` that can be used to wait for all the tasks to finish. If the sequence function returns a value, the `multi_future` can also be used to obtain the values returned by each task.
716 | */
717 | template ,T>>
718 | [[nodiscard]] multi_future submit_sequence(const T first_index,const T index_after_last,F&& sequence BS_THREAD_POOL_PRIORITY_INPUT)
719 | {
720 | if(index_after_last > first_index)
721 | {
722 | multi_future future;
723 | future.reserve(static_cast(index_after_last - first_index));
724 | for(T i = first_index; i < index_after_last; ++i)
725 | future.push_back(submit_task(
726 | [sequence = std::forward(sequence),i]
727 | {
728 | return sequence(i);
729 | } BS_THREAD_POOL_PRIORITY_OUTPUT));
730 | return future;
731 | }
732 | return {};
733 | }
734 |
735 | #ifdef BS_THREAD_POOL_ENABLE_PAUSE
736 | /**
737 | * @brief Unpause the pool. The workers will resume retrieving new tasks out of the queue. Only enabled if `BS_THREAD_POOL_ENABLE_PAUSE` is defined.
738 | */
739 | void unpause()
740 | {
741 | {
742 | const std::scoped_lock tasks_lock(tasks_mutex);
743 | paused = false;
744 | }
745 | task_available_cv.notify_all();
746 | }
747 | #endif
748 |
749 | // Macros used internally to enable or disable pausing in the waiting and worker functions.
750 | #ifdef BS_THREAD_POOL_ENABLE_PAUSE
751 | #define BS_THREAD_POOL_PAUSED_OR_EMPTY (paused || tasks.empty())
752 | #else
753 | #define BS_THREAD_POOL_PAUSED_OR_EMPTY tasks.empty()
754 | #endif
755 |
756 | /**
757 | * @brief Wait for tasks to be completed. Normally, this function waits for all tasks, both those that are currently running in the threads and those that are still waiting in the queue. However, if the pool is paused, this function only waits for the currently running tasks (otherwise it would wait forever). Note: To wait for just one specific task, use `submit_task()` instead, and call the `wait()` member function of the generated future.
758 | *
759 | * @throws `wait_deadlock` if called from within a thread of the same pool, which would result in a deadlock. Only enabled if `BS_THREAD_POOL_ENABLE_WAIT_DEADLOCK_CHECK` is defined.
760 | */
761 | void wait()
762 | {
763 | #ifdef BS_THREAD_POOL_ENABLE_WAIT_DEADLOCK_CHECK
764 | if(this_thread::get_pool() == this)
765 | throw wait_deadlock();
766 | #endif
767 | std::unique_lock tasks_lock(tasks_mutex);
768 | waiting = true;
769 | tasks_done_cv.wait(tasks_lock,
770 | [this]
771 | {
772 | return (tasks_running == 0) && BS_THREAD_POOL_PAUSED_OR_EMPTY;
773 | });
774 | waiting = false;
775 | }
776 |
777 | /**
778 | * @brief Wait for tasks to be completed, but stop waiting after the specified duration has passed.
779 | *
780 | * @tparam R An arithmetic type representing the number of ticks to wait.
781 | * @tparam P An `std::ratio` representing the length of each tick in seconds.
782 | * @param duration The amount of time to wait.
783 | * @return `true` if all tasks finished running, `false` if the duration expired but some tasks are still running.
784 | *
785 | * @throws `wait_deadlock` if called from within a thread of the same pool, which would result in a deadlock. Only enabled if `BS_THREAD_POOL_ENABLE_WAIT_DEADLOCK_CHECK` is defined.
786 | */
787 | template
788 | bool wait_for(const std::chrono::duration& duration)
789 | {
790 | #ifdef BS_THREAD_POOL_ENABLE_WAIT_DEADLOCK_CHECK
791 | if(this_thread::get_pool() == this)
792 | throw wait_deadlock();
793 | #endif
794 | std::unique_lock tasks_lock(tasks_mutex);
795 | waiting = true;
796 | const bool status = tasks_done_cv.wait_for(tasks_lock,duration,
797 | [this]
798 | {
799 | return (tasks_running == 0) && BS_THREAD_POOL_PAUSED_OR_EMPTY;
800 | });
801 | waiting = false;
802 | return status;
803 | }
804 |
805 | /**
806 | * @brief Wait for tasks to be completed, but stop waiting after the specified time point has been reached.
807 | *
808 | * @tparam C The type of the clock used to measure time.
809 | * @tparam D An `std::chrono::duration` type used to indicate the time point.
810 | * @param timeout_time The time point at which to stop waiting.
811 | * @return `true` if all tasks finished running, `false` if the time point was reached but some tasks are still running.
812 | *
813 | * @throws `wait_deadlock` if called from within a thread of the same pool, which would result in a deadlock. Only enabled if `BS_THREAD_POOL_ENABLE_WAIT_DEADLOCK_CHECK` is defined.
814 | */
815 | template
816 | bool wait_until(const std::chrono::time_point& timeout_time)
817 | {
818 | #ifdef BS_THREAD_POOL_ENABLE_WAIT_DEADLOCK_CHECK
819 | if(this_thread::get_pool() == this)
820 | throw wait_deadlock();
821 | #endif
822 | std::unique_lock tasks_lock(tasks_mutex);
823 | waiting = true;
824 | const bool status = tasks_done_cv.wait_until(tasks_lock,timeout_time,
825 | [this]
826 | {
827 | return (tasks_running == 0) && BS_THREAD_POOL_PAUSED_OR_EMPTY;
828 | });
829 | waiting = false;
830 | return status;
831 | }
832 |
833 | #ifdef BS_THREAD_POOL_ENABLE_WAIT_DEADLOCK_CHECK
834 | // ==============
835 | // Public classes
836 | // ==============
837 |
838 | /**
839 | * @brief An exception that will be thrown by `wait()`, `wait_for()`, and `wait_until()` if the user tries to call them from within a thread of the same pool, which would result in a deadlock.
840 | */
841 | struct wait_deadlock: public std::runtime_error
842 | {
843 | wait_deadlock(): std::runtime_error("BS::thread_pool::wait_deadlock"){};
844 | };
845 | #endif
846 |
847 | private:
848 | // ========================
849 | // Private member functions
850 | // ========================
851 |
852 | /**
853 | * @brief Create the threads in the pool and assign a worker to each thread.
854 | *
855 | * @param init_task An initialization function to run in each thread before it starts to execute any submitted tasks.
856 | */
857 | void create_threads(const std::function& init_task)
858 | {
859 | {
860 | const std::scoped_lock tasks_lock(tasks_mutex);
861 | tasks_running = thread_count;
862 | workers_running = true;
863 | }
864 | for(concurrency_t i = 0; i < thread_count; ++i)
865 | {
866 | threads[i] = std::thread(&thread_pool::worker,this,i,init_task);
867 | }
868 | }
869 |
870 | /**
871 | * @brief Destroy the threads in the pool.
872 | */
873 | void destroy_threads()
874 | {
875 | {
876 | const std::scoped_lock tasks_lock(tasks_mutex);
877 | workers_running = false;
878 | }
879 | task_available_cv.notify_all();
880 | for(concurrency_t i = 0; i < thread_count; ++i)
881 | {
882 | threads[i].join();
883 | }
884 | }
885 |
886 | /**
887 | * @brief Determine how many threads the pool should have, based on the parameter passed to the constructor or reset().
888 | *
889 | * @param num_threads The parameter passed to the constructor or `reset()`. If the parameter is a positive number, then the pool will be created with this number of threads. If the parameter is non-positive, or a parameter was not supplied (in which case it will have the default value of 0), then the pool will be created with the total number of hardware threads available, as obtained from `std::thread::hardware_concurrency()`. If the latter returns zero for some reason, then the pool will be created with just one thread.
890 | * @return The number of threads to use for constructing the pool.
891 | */
892 | [[nodiscard]] static concurrency_t determine_thread_count(const concurrency_t num_threads)
893 | {
894 | if(num_threads > 0)
895 | return num_threads;
896 | if(std::thread::hardware_concurrency() > 0)
897 | return std::thread::hardware_concurrency();
898 | return 1;
899 | }
900 |
901 | /**
902 | * @brief A worker function to be assigned to each thread in the pool. Waits until it is notified by `detach_task()` that a task is available, and then retrieves the task from the queue and executes it. Once the task finishes, the worker notifies `wait()` in case it is waiting.
903 | *
904 | * @param idx The index of this thread.
905 | * @param init_task An initialization function to run in this thread before it starts to execute any submitted tasks.
906 | */
907 | void worker(const concurrency_t idx,const std::function& init_task)
908 | {
909 | this_thread::get_index.index = idx;
910 | this_thread::get_pool.pool = this;
911 | init_task();
912 | std::unique_lock tasks_lock(tasks_mutex);
913 | while(true)
914 | {
915 | --tasks_running;
916 | tasks_lock.unlock();
917 | if(waiting && (tasks_running == 0) && BS_THREAD_POOL_PAUSED_OR_EMPTY)
918 | tasks_done_cv.notify_all();
919 | tasks_lock.lock();
920 | task_available_cv.wait(tasks_lock,
921 | [this]
922 | {
923 | return !BS_THREAD_POOL_PAUSED_OR_EMPTY || !workers_running;
924 | });
925 | if(!workers_running)
926 | break;
927 | {
928 | #ifdef BS_THREAD_POOL_ENABLE_PRIORITY
929 | const std::function task = std::move(std::remove_const_t(tasks.top()).task);
930 | tasks.pop();
931 | #else
932 | const std::function task = std::move(tasks.front());
933 | tasks.pop();
934 | #endif
935 | ++tasks_running;
936 | tasks_lock.unlock();
937 | task();
938 | }
939 | tasks_lock.lock();
940 | }
941 | this_thread::get_index.index = std::nullopt;
942 | this_thread::get_pool.pool = std::nullopt;
943 | }
944 |
945 | // ===============
946 | // Private classes
947 | // ===============
948 |
949 | /**
950 | * @brief A helper class to divide a range into blocks. Used by `detach_blocks()`, `submit_blocks()`, `detach_loop()`, and `submit_loop()`.
951 | *
952 | * @tparam T The type of the indices. Should be a signed or unsigned integer.
953 | */
954 | template
955 | class [[nodiscard]] blocks
956 | {
957 | public:
958 | /**
959 | * @brief Construct a `blocks` object with the given specifications.
960 | *
961 | * @param first_index_ The first index in the range.
962 | * @param index_after_last_ The index after the last index in the range.
963 | * @param num_blocks_ The desired number of blocks to divide the range into.
964 | */
965 | blocks(const T first_index_,const T index_after_last_,const size_t num_blocks_): first_index(first_index_),index_after_last(index_after_last_),num_blocks(num_blocks_)
966 | {
967 | if(index_after_last > first_index)
968 | {
969 | const size_t total_size = static_cast(index_after_last - first_index);
970 | if(num_blocks > total_size)
971 | num_blocks = total_size;
972 | block_size = total_size / num_blocks;
973 | remainder = total_size % num_blocks;
974 | if(block_size == 0)
975 | {
976 | block_size = 1;
977 | num_blocks = (total_size > 1) ? total_size : 1;
978 | }
979 | }
980 | else
981 | {
982 | num_blocks = 0;
983 | }
984 | }
985 |
986 | /**
987 | * @brief Get the first index of a block.
988 | *
989 | * @param block The block number.
990 | * @return The first index.
991 | */
992 | [[nodiscard]] T start(const size_t block) const
993 | {
994 | return first_index + static_cast(block * block_size) + static_cast(block < remainder ? block : remainder);
995 | }
996 |
997 | /**
998 | * @brief Get the index after the last index of a block.
999 | *
1000 | * @param block The block number.
1001 | * @return The index after the last index.
1002 | */
1003 | [[nodiscard]] T end(const size_t block) const
1004 | {
1005 | return (block == num_blocks - 1) ? index_after_last : start(block + 1);
1006 | }
1007 |
1008 | /**
1009 | * @brief Get the number of blocks. Note that this may be different than the desired number of blocks that was passed to the constructor.
1010 | *
1011 | * @return The number of blocks.
1012 | */
1013 | [[nodiscard]] size_t get_num_blocks() const
1014 | {
1015 | return num_blocks;
1016 | }
1017 |
1018 | private:
1019 | /**
1020 | * @brief The size of each block (except possibly the last block).
1021 | */
1022 | size_t block_size = 0;
1023 |
1024 | /**
1025 | * @brief The first index in the range.
1026 | */
1027 | T first_index = 0;
1028 |
1029 | /**
1030 | * @brief The index after the last index in the range.
1031 | */
1032 | T index_after_last = 0;
1033 |
1034 | /**
1035 | * @brief The number of blocks.
1036 | */
1037 | size_t num_blocks = 0;
1038 |
1039 | /**
1040 | * @brief The remainder obtained after dividing the total size by the number of blocks.
1041 | */
1042 | size_t remainder = 0;
1043 | }; // class blocks
1044 |
1045 | #ifdef BS_THREAD_POOL_ENABLE_PRIORITY
1046 | /**
1047 | * @brief A helper class to store a task with an assigned priority.
1048 | */
1049 | class [[nodiscard]] pr_task
1050 | {
1051 | friend class thread_pool;
1052 |
1053 | public:
1054 | /**
1055 | * @brief Construct a new task with an assigned priority by copying the task.
1056 | *
1057 | * @param task_ The task.
1058 | * @param priority_ The desired priority.
1059 | */
1060 | explicit pr_task(const std::function& task_,const priority_t priority_ = 0): task(task_),priority(priority_) {}
1061 |
1062 | /**
1063 | * @brief Construct a new task with an assigned priority by moving the task.
1064 | *
1065 | * @param task_ The task.
1066 | * @param priority_ The desired priority.
1067 | */
1068 | explicit pr_task(std::function&& task_,const priority_t priority_ = 0): task(std::move(task_)),priority(priority_) {}
1069 |
1070 | /**
1071 | * @brief Compare the priority of two tasks.
1072 | *
1073 | * @param lhs The first task.
1074 | * @param rhs The second task.
1075 | * @return `true` if the first task has a lower priority than the second task, `false` otherwise.
1076 | */
1077 | [[nodiscard]] friend bool operator<(const pr_task& lhs,const pr_task& rhs)
1078 | {
1079 | return lhs.priority < rhs.priority;
1080 | }
1081 |
1082 | private:
1083 | /**
1084 | * @brief The task.
1085 | */
1086 | std::function task ={};
1087 |
1088 | /**
1089 | * @brief The priority of the task.
1090 | */
1091 | priority_t priority = 0;
1092 | }; // class pr_task
1093 | #endif
1094 |
1095 | // ============
1096 | // Private data
1097 | // ============
1098 |
1099 | #ifdef BS_THREAD_POOL_ENABLE_PAUSE
1100 | /**
1101 | * @brief A flag indicating whether the workers should pause. When set to `true`, the workers temporarily stop retrieving new tasks out of the queue, although any tasks already executed will keep running until they are finished. When set to `false` again, the workers resume retrieving tasks.
1102 | */
1103 | bool paused = false;
1104 | #endif
1105 |
1106 | /**
1107 | * @brief A condition variable to notify `worker()` that a new task has become available.
1108 | */
1109 | std::condition_variable task_available_cv ={};
1110 |
1111 | /**
1112 | * @brief A condition variable to notify `wait()` that the tasks are done.
1113 | */
1114 | std::condition_variable tasks_done_cv ={};
1115 |
1116 | /**
1117 | * @brief A queue of tasks to be executed by the threads.
1118 | */
1119 | #ifdef BS_THREAD_POOL_ENABLE_PRIORITY
1120 | std::priority_queue tasks ={};
1121 | #else
1122 | std::queue> tasks ={};
1123 | #endif
1124 |
1125 | /**
1126 | * @brief A counter for the total number of currently running tasks.
1127 | */
1128 | size_t tasks_running = 0;
1129 |
1130 | /**
1131 | * @brief A mutex to synchronize access to the task queue by different threads.
1132 | */
1133 | mutable std::mutex tasks_mutex ={};
1134 |
1135 | /**
1136 | * @brief The number of threads in the pool.
1137 | */
1138 | concurrency_t thread_count = 0;
1139 |
1140 | /**
1141 | * @brief A smart pointer to manage the memory allocated for the threads.
1142 | */
1143 | std::unique_ptr threads = nullptr;
1144 |
1145 | /**
1146 | * @brief A flag indicating that `wait()` is active and expects to be notified whenever a task is done.
1147 | */
1148 | bool waiting = false;
1149 |
1150 | /**
1151 | * @brief A flag indicating to the workers to keep running. When set to `false`, the workers terminate permanently.
1152 | */
1153 | bool workers_running = false;
1154 | }; // class thread_pool
1155 | } // namespace BS
1156 | #endif
1157 |
--------------------------------------------------------------------------------