Skip to content

Commit 13f9261

Browse files
DarthTonDarthTon
committed
addressed DarthTon#341: HandleGuard refactoring
1 parent 38bf089 commit 13f9261

File tree

13 files changed

+35
-62
lines changed

13 files changed

+35
-62
lines changed

‎src/BlackBone/DriverControl/DriverControl.cpp

Lines changed: 3 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -793,23 +793,20 @@ 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 key1, key2;
797797
DWORD dwType = 1;
798798
LSTATUS status = 0;
799799
WCHAR wszLocalPath[MAX_PATH] = { 0 };
800800

801-
swprintf_s( wszLocalPath, ARRAYSIZE( wszLocalPath ), L"\\??\\%s", path.c_str() );
801+
swprintf_s( wszLocalPath, std::size( wszLocalPath ), L"\\??\\%s", path.c_str() );
802802

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

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

814811
status = RegSetValueExW(
815812
key2, L"ImagePath", 0, REG_SZ,
@@ -818,24 +815,9 @@ LSTATUS DriverControl::PrepareDriverRegEntry( const std::wstring& svcName, const
818815
);
819816

820817
if (status)
821-
{
822-
RegCloseKey( key2 );
823-
RegCloseKey( key1 );
824818
return status;
825-
}
826819

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 );
832-
return status;
833-
}
834-
835-
RegCloseKey( key2 );
836-
RegCloseKey( key1 );
837-
838-
return status;
820+
return RegSetValueExW( key2, L"Type", 0, REG_DWORD, reinterpret_cast<const BYTE*>(&dwType), sizeof( dwType ) );
839821
}
840822

841823

‎src/BlackBone/DriverControl/DriverControl.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -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: 14 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ namespace blackbone
77
/// <summary>
88
/// Strong exception guarantee
99
/// </summary>
10-
template<typename handle_t, auto close_fn, handle_t zero_handle>
10+
template<typename handle_t, auto close_fn>
1111
class HandleGuard
1212
{
1313
public:
@@ -24,7 +24,7 @@ class HandleGuard
2424

2525
~HandleGuard()
2626
{
27-
if (_handle != zero_handle)
27+
if (valid())
2828
close_fn( _handle );
2929
}
3030

@@ -53,7 +53,7 @@ class HandleGuard
5353
if (handle == _handle)
5454
return;
5555

56-
if (_handle != zero_handle)
56+
if (valid())
5757
close_fn( _handle );
5858

5959
_handle = handle;
@@ -66,25 +66,24 @@ class HandleGuard
6666
return tmp;
6767
}
6868

69-
inline handle_t get() const noexcept { return _handle; }
69+
handle_t get() const noexcept { return _handle; }
70+
bool valid() const noexcept { return reinterpret_cast<intptr_t>(_handle) > reinterpret_cast<intptr_t>(zero_handle); }
7071

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

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

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; }
77+
bool operator ==( const HandleGuard& rhs ) const noexcept { return _handle == rhs._handle; }
78+
bool operator <( const HandleGuard& rhs ) const noexcept { return _handle < rhs._handle; }
7879

7980
private:
81+
static constexpr handle_t zero_handle = handle_t( 0 );
8082
handle_t _handle;
8183
};
8284

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>;
85+
using Handle = HandleGuard<HANDLE, &CloseHandle>;
86+
using ACtxHandle = HandleGuard<HANDLE, &ReleaseActCtx>;
87+
using RegHandle = HandleGuard<HKEY, &RegCloseKey>;
8988

9089
}

‎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: 2 additions & 3 deletions
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;
@@ -521,8 +521,7 @@ NTSTATUS PEImage::PrepareACTX( const wchar_t* filepath /*= nullptr*/ )
521521

522522
// Create ACTX
523523
_hctx = CreateActCtxW( &act );
524-
525-
if (_hctx != INVALID_HANDLE_VALUE)
524+
if (_hctx)
526525
return STATUS_SUCCESS;
527526

528527
// Return success if current process is protected

‎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/BlackBone/Syscalls/Syscall.h

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,7 @@
66

77
extern "C" void* syscall_stub();
88

9-
namespace blackbone
10-
{
11-
namespace syscall
9+
namespace blackbone::syscall
1210
{
1311
template<typename T>
1412
using to_int64 = std::conditional_t<sizeof( T ) < sizeof( int64_t ), int64_t, T>;
@@ -68,4 +66,3 @@ namespace syscall
6866

6967
#pragma warning(pop)
7068
}
71-
}

‎src/BlackBoneTest/TestAsmJit.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ namespace Testing
7878
AssertEx::IsTrue( b );
7979

8080
// Check locally
81-
FileHandle hFile = FileHandle( CreateFileW( filePath, GENERIC_READ, 0x7, nullptr, OPEN_EXISTING, FILE_DELETE_ON_CLOSE, nullptr ) );
81+
auto hFile = Handle( CreateFileW( filePath, GENERIC_READ, 0x7, nullptr, OPEN_EXISTING, FILE_DELETE_ON_CLOSE, nullptr ) );
8282
AssertEx::IsTrue( hFile );
8383

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

‎src/BlackBoneTest/TestManualMap.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,7 @@ namespace Testing
131131

132132
std::pair<std::unique_ptr<uint8_t[]>, uint32_t> GetFileData( const std::wstring& path )
133133
{
134-
auto hFile = FileHandle( CreateFileW( path.c_str(), FILE_GENERIC_READ, 0x7, nullptr, OPEN_EXISTING, 0, nullptr ) );
134+
auto hFile = Handle( CreateFileW( path.c_str(), FILE_GENERIC_READ, 0x7, nullptr, OPEN_EXISTING, 0, nullptr ) );
135135
if (hFile)
136136
{
137137
uint32_t size = GetFileSize( hFile, nullptr );

‎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 )] = { };

‎src/Samples/Main.cpp

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ int main( int /*argc*/, char* /*argv[]*/ )
4444
}
4545

4646
// Start new suspended process and attach immediately
47-
Process notepad;
47+
Process notepad;
4848
notepad.CreateAndAttach( L"C:\\windows\\system32\\notepad.exe", true );
4949
{
5050
// do stuff...
@@ -93,7 +93,7 @@ int main( int /*argc*/, char* /*argv[]*/ )
9393
memory.Read( mainMod->baseAddress, sizeof( dosHeader ), &dosHeader );
9494

9595
// Method 3
96-
auto[status, dosHeader2] = memory.Read<IMAGE_DOS_HEADER>( mainMod->baseAddress );
96+
auto [status, dosHeader2] = memory.Read<IMAGE_DOS_HEADER>( mainMod->baseAddress );
9797

9898
// Change memory protection
9999
if (NT_SUCCESS( memory.Protect( mainMod->baseAddress, sizeof( dosHeader ), PAGE_READWRITE ) ))
@@ -110,7 +110,7 @@ int main( int /*argc*/, char* /*argv[]*/ )
110110
}
111111

112112
// Allocate memory
113-
if (auto[status2, block] = memory.Allocate( 0x1000, PAGE_EXECUTE_READWRITE ); NT_SUCCESS( status2 ))
113+
if (auto [status2, block] = memory.Allocate( 0x1000, PAGE_EXECUTE_READWRITE ); NT_SUCCESS( status2 ))
114114
{
115115
// Write into memory block
116116
block->Write( 0x10, 12.0 );
@@ -164,7 +164,7 @@ int main( int /*argc*/, char* /*argv[]*/ )
164164
{
165165
auto& remote = notepad.remote();
166166
remote.CreateRPCEnvironment( Worker_None, true );
167-
167+
168168
auto GetModuleHandleWPtr = notepad.modules().GetExport( L"kernel32.dll", "GetModuleHandleW" );
169169
if (GetModuleHandleWPtr)
170170
{
@@ -256,7 +256,7 @@ int main( int /*argc*/, char* /*argv[]*/ )
256256
GetCurrentProcess(),
257257
GetModuleHandle( nullptr ),
258258
buf,
259-
sizeof(buf),
259+
sizeof( buf ),
260260
&bytes
261261
);
262262

0 commit comments

Comments
 (0)