Skip to content

Commit 1c50ebf

Browse files
DarthTonDarthTon
committed
DarthTon#341: HandleGuard refactoring
1 parent 04bcad5 commit 1c50ebf

File tree

9 files changed

+65
-55
lines changed

9 files changed

+65
-55
lines changed

‎src/BlackBone/DriverControl/DriverControl.cpp

Lines changed: 5 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -793,50 +793,31 @@ NTSTATUS DriverControl::UnloadDriver( const std::wstring& svcName )
793793
/// <returns>Status code</returns>
794794
LSTATUS DriverControl::PrepareDriverRegEntry( const std::wstring& svcName, const std::wstring& path )
795795
{
796-
HKEY key1, key2;
796+
RegHandle svcRoot, svcKey;
797797
DWORD dwType = 1;
798798
LSTATUS status = 0;
799799
WCHAR wszLocalPath[MAX_PATH] = { 0 };
800800

801801
swprintf_s( wszLocalPath, ARRAYSIZE( wszLocalPath ), L"\\??\\%s", path.c_str() );
802802

803-
status = RegOpenKeyW( HKEY_LOCAL_MACHINE, L"system\\CurrentControlSet\\Services", &key1 );
803+
status = RegOpenKeyW( HKEY_LOCAL_MACHINE, L"system\\CurrentControlSet\\Services", &svcRoot );
804804
if (status)
805805
return status;
806806

807-
status = RegCreateKeyW( key1, svcName.c_str(), &key2 );
807+
status = RegCreateKeyW( svcRoot, svcName.c_str(), &svcKey );
808808
if (status)
809-
{
810-
RegCloseKey( key1 );
811809
return status;
812-
}
813810

814811
status = RegSetValueExW(
815-
key2, L"ImagePath", 0, REG_SZ,
812+
svcKey, L"ImagePath", 0, REG_SZ,
816813
reinterpret_cast<const BYTE*>(wszLocalPath),
817814
static_cast<DWORD>(sizeof( WCHAR )* (wcslen( wszLocalPath ) + 1))
818815
);
819816

820817
if (status)
821-
{
822-
RegCloseKey( key2 );
823-
RegCloseKey( key1 );
824-
return status;
825-
}
826-
827-
status = RegSetValueExW( key2, L"Type", 0, REG_DWORD, reinterpret_cast<const BYTE*>(&dwType), sizeof( dwType ) );
828-
if (status)
829-
{
830-
RegCloseKey( key2 );
831-
RegCloseKey( key1 );
832818
return status;
833-
}
834819

835-
RegCloseKey( key2 );
836-
RegCloseKey( key1 );
837-
838-
return status;
820+
return RegSetValueExW( svcKey, L"Type", 0, REG_DWORD, reinterpret_cast<const BYTE*>(&dwType), sizeof( dwType ) );
839821
}
840822

841-
842823
}

‎src/BlackBone/DriverControl/DriverControl.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -282,7 +282,7 @@ class DriverControl
282282
/// Check if driver is loaded
283283
/// </summary>
284284
/// <returns></returns>
285-
BLACKBONE_API inline bool loaded() const { return _hDriver != INVALID_HANDLE_VALUE; }
285+
BLACKBONE_API inline bool loaded() const { return _hDriver.valid(); }
286286
BLACKBONE_API inline NTSTATUS status() const { return _loadStatus; }
287287

288288
private:
@@ -312,7 +312,7 @@ class DriverControl
312312
/// <returns>Status code</returns>
313313
LSTATUS PrepareDriverRegEntry( const std::wstring& svcName, const std::wstring& path );
314314
private:
315-
FileHandle _hDriver;
315+
Handle _hDriver;
316316
NTSTATUS _loadStatus = STATUS_NOT_FOUND;
317317
};
318318

‎src/BlackBone/Include/HandleGuard.h

Lines changed: 50 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,45 @@
44
namespace blackbone
55
{
66

7+
template<typename T>
8+
struct non_negative
9+
{
10+
static bool call( T handle ) noexcept
11+
{
12+
return intptr_t( handle ) > 0;
13+
}
14+
};
15+
16+
template<template<typename> typename wrapped_t, typename T>
17+
struct with_pseudo_t
18+
{
19+
static bool call( T handle ) noexcept
20+
{
21+
if (wrapped_t<T>::call( handle ))
22+
return true;
23+
24+
// Check if it's a pseudo handle
25+
auto h = (HANDLE)(uintptr_t)handle;
26+
return h == GetCurrentProcess() || h == GetCurrentThread();
27+
}
28+
};
29+
30+
template<template<typename> typename wrapped_t>
31+
struct with_pseudo
32+
{
33+
template<typename T>
34+
using type = with_pseudo_t<wrapped_t, T>;
35+
};
36+
737
/// <summary>
838
/// Strong exception guarantee
939
/// </summary>
10-
template<typename handle_t, auto close_fn, handle_t zero_handle>
40+
template<typename handle_t, auto close, template<typename> typename is_valid = non_negative>
1141
class HandleGuard
1242
{
43+
public:
44+
static constexpr handle_t zero_handle = handle_t( 0 );
45+
1346
public:
1447
explicit HandleGuard( handle_t handle = zero_handle ) noexcept
1548
: _handle( handle )
@@ -24,8 +57,8 @@ class HandleGuard
2457

2558
~HandleGuard()
2659
{
27-
if (_handle != zero_handle)
28-
close_fn( _handle );
60+
if (non_negative<handle_t>::call( _handle ))
61+
close( _handle );
2962
}
3063

3164
HandleGuard( const HandleGuard& ) = delete;
@@ -53,8 +86,8 @@ class HandleGuard
5386
if (handle == _handle)
5487
return;
5588

56-
if (_handle != zero_handle)
57-
close_fn( _handle );
89+
if (non_negative<handle_t>::call( _handle ))
90+
close( _handle );
5891

5992
_handle = handle;
6093
}
@@ -66,25 +99,25 @@ class HandleGuard
6699
return tmp;
67100
}
68101

69-
inline handle_t get() const noexcept { return _handle; }
102+
handle_t get() const noexcept { return _handle; }
103+
bool valid() const noexcept { return is_valid<handle_t>::call( _handle ); }
70104

71-
inline operator handle_t() const noexcept { return _handle; }
72-
inline explicit operator bool() const noexcept { return _handle != zero_handle; }
105+
operator handle_t() const noexcept { return _handle; }
106+
explicit operator bool() const noexcept { return valid(); }
73107

74-
inline handle_t* operator &() noexcept { return &_handle; }
108+
handle_t* operator &() noexcept { return &_handle; }
75109

76-
inline bool operator ==( const HandleGuard& rhs ) const noexcept { return _handle == rhs._handle; }
77-
inline bool operator <( const HandleGuard& rhs ) const noexcept { return _handle < rhs._handle; }
110+
bool operator ==( const HandleGuard& rhs ) const noexcept { return _handle == rhs._handle; }
111+
bool operator <( const HandleGuard& rhs ) const noexcept { return _handle < rhs._handle; }
78112

79113
private:
80114
handle_t _handle;
81115
};
82116

83-
using Handle = HandleGuard<HANDLE, &CloseHandle, nullptr>;
84-
using FileHandle = HandleGuard<HANDLE, &CloseHandle, INVALID_HANDLE_VALUE>;
85-
using ACtxHandle = HandleGuard<HANDLE, &ReleaseActCtx, INVALID_HANDLE_VALUE>;
86-
using FileMapHandle = HandleGuard<void*, &UnmapViewOfFile, nullptr>;
87-
using SnapHandle = FileHandle;
88-
using RegHandle = HandleGuard<HKEY, &RegCloseKey, nullptr>;
117+
118+
using Handle = HandleGuard<HANDLE, &CloseHandle>;
119+
using ProcessHandle = HandleGuard<HANDLE, &CloseHandle, with_pseudo<non_negative>::type>;
120+
using ACtxHandle = HandleGuard<HANDLE, &ReleaseActCtx>;
121+
using RegHandle = HandleGuard<HKEY, &RegCloseKey>;
89122

90123
}

‎src/BlackBone/Misc/NameResolve.cpp

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -167,7 +167,7 @@ NTSTATUS NameResolve::ResolvePath(
167167
// Perform search accordingly to Windows Image loader search order
168168
// 1. KnownDlls
169169
//
170-
HKEY hKey = NULL;
170+
RegHandle hKey;
171171
LRESULT res = 0;
172172
res = RegOpenKeyW( HKEY_LOCAL_MACHINE, L"SYSTEM\\CurrentControlSet\\Control\\Session Manager\\KnownDLLs", &hKey );
173173

@@ -197,14 +197,10 @@ NTSTATUS NameResolve::ResolvePath(
197197
if (res == ERROR_SUCCESS)
198198
{
199199
path = std::wstring( sys_path ) + L"\\" + value_data;
200-
201-
RegCloseKey( hKey );
202200
return STATUS_SUCCESS;
203201
}
204202
}
205203
}
206-
207-
RegCloseKey( hKey );
208204
}
209205

210206

@@ -359,7 +355,7 @@ NTSTATUS NameResolve::ProbeSxSRedirect( std::wstring& path, Process& proc, HANDL
359355
/// <returns>Process executable directory</returns>
360356
std::wstring NameResolve::GetProcessDirectory( DWORD pid )
361357
{
362-
SnapHandle snapshot;
358+
Handle snapshot;
363359
MODULEENTRY32W mod = { sizeof(MODULEENTRY32W), 0 };
364360
std::wstring path = L"";
365361

‎src/BlackBone/PE/PEImage.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -497,7 +497,7 @@ NTSTATUS PEImage::PrepareACTX( const wchar_t* filepath /*= nullptr*/ )
497497
if (GetTempFileNameW( tempDir, L"ImageManifest", 0, tempPath ) == 0)
498498
return STATUS_SXS_CANT_GEN_ACTCTX;
499499

500-
auto hTmpFile = FileHandle( CreateFileW( tempPath, FILE_GENERIC_WRITE, 0, NULL, OPEN_ALWAYS, 0, NULL ) );
500+
auto hTmpFile = Handle( CreateFileW( tempPath, FILE_GENERIC_WRITE, 0, NULL, OPEN_ALWAYS, 0, NULL ) );
501501
if (hTmpFile != INVALID_HANDLE_VALUE)
502502
{
503503
DWORD bytes = 0;

‎src/BlackBone/PE/PEImage.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -310,7 +310,7 @@ class PEImage
310310
void* GetManifest( uint32_t& size, int32_t& manifestID );
311311

312312
private:
313-
FileHandle _hFile; // Target file HANDLE
313+
Handle _hFile; // Target file HANDLE
314314
Handle _hMapping; // Memory mapping object
315315
void* _pFileBase = nullptr; // Mapping base
316316
bool _isPlainData = false; // File mapped as plain data file
@@ -324,7 +324,7 @@ class PEImage
324324
uint32_t _imgSize = 0; // Image size
325325
uint32_t _epRVA = 0; // Entry point RVA
326326
uint32_t _hdrSize = 0; // Size of headers
327-
ACtxHandle _hctx; // Activation context
327+
ACtxHandle _hctx; // Activation context
328328
int32_t _manifestIdx = 0; // Manifest resource ID
329329
uint32_t _subsystem = 0; // Image subsystem
330330
int32_t _ILFlagOffset = 0; // Offset of pure IL flag

‎src/BlackBone/Process/Process.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -323,7 +323,7 @@ call_result_t<std::vector<HandleInfo>> Process::EnumHandles()
323323
std::vector<DWORD> Process::EnumByName( const std::wstring& name )
324324
{
325325
std::vector<DWORD> found;
326-
auto hProcSnap = SnapHandle( CreateToolhelp32Snapshot( TH32CS_SNAPPROCESS, 0 ) );
326+
auto hProcSnap = Handle( CreateToolhelp32Snapshot( TH32CS_SNAPPROCESS, 0 ) );
327327
if (!hProcSnap)
328328
return found;
329329

‎src/BlackBone/Process/Threads/Threads.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ call_result_t<ThreadPtr> ProcessThreads::CreateNew( ptr_t threadProc, ptr_t arg,
5353
std::vector<ThreadPtr> ProcessThreads::getAll() const
5454
{
5555
std::vector<ThreadPtr> result;
56-
auto hThreadSnapshot = SnapHandle( CreateToolhelp32Snapshot( TH32CS_SNAPTHREAD, 0 ) );
56+
auto hThreadSnapshot = Handle( CreateToolhelp32Snapshot( TH32CS_SNAPTHREAD, 0 ) );
5757
if (!hThreadSnapshot)
5858
return result;
5959

‎src/BlackBoneTest/TestRemoteCall.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -195,7 +195,7 @@ namespace Testing
195195
AssertEx::AreEqual( static_cast<DWORD>(sizeof( writeBuf )), bytes );
196196

197197
// Check locally
198-
FileHandle hFile = FileHandle( CreateFileW( filePath, GENERIC_READ, 0x7, nullptr, OPEN_EXISTING, FILE_DELETE_ON_CLOSE, nullptr ) );
198+
auto hFile = Handle( CreateFileW( filePath, GENERIC_READ, 0x7, nullptr, OPEN_EXISTING, FILE_DELETE_ON_CLOSE, nullptr ) );
199199
AssertEx::IsTrue( hFile );
200200

201201
uint8_t readBuf[sizeof( writeBuf )] = { };

0 commit comments

Comments
 (0)