bug fix; code style fix; added test for ExecInAnyThread

DarthTon · Jun 20, 2021 · a672509 · a672509
1 parent e36af76
commit a672509
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 12 deletions.
diff --git a/src/BlackBone/Process/RPC/RemoteExec.cpp b/src/BlackBone/Process/RPC/RemoteExec.cpp
@@ -225,7 +225,7 @@ NTSTATUS RemoteExec::ExecInAnyThread( PVOID pCode, size_t size, uint64_t& callRe
             (*a)->mov( asmjit::Mem( asmjit::host::rsp, i * sizeof( uint64_t ) ), regs[i] );
 
         a->GenCall( _userCode[_currentBufferIdx].ptr(), { _userData[_currentBufferIdx].ptr() } );
-        AddReturnWithEvent( *a );
+        AddReturnWithEvent( *a, mt_mod64, rt_int32, INTRET_OFFSET );
 
         // Restore registers
         for (int i = 0; i < count; i++)
@@ -279,11 +279,7 @@ NTSTATUS RemoteExec::ExecInAnyThread( PVOID pCode, size_t size, uint64_t& callRe
     if (NT_SUCCESS( status ))
     {
         WaitForSingleObject( _hWaitEvent, 20 * 1000/*INFINITE*/ );
-
-        if (!_process.core().isWow64())
-            status = _userData[_currentBufferIdx].Read( RET_OFFSET, callResult );
-        else
-            status = _userData[_currentBufferIdx].Read( INTRET_OFFSET, callResult );
+        status = _userData[_currentBufferIdx].Read( INTRET_OFFSET, callResult );
     }
 
     SwitchActiveBuffer();

diff --git a/src/BlackBone/Process/RPC/RemoteFunction.hpp b/src/BlackBone/Process/RPC/RemoteFunction.hpp
@@ -55,7 +55,7 @@ class RemoteFunctionBase
     RemoteFunctionBase( Process& proc, ptr_t ptr, ThreadPtr boundThread = nullptr )
         : _process( proc )
         , _ptr( ptr )
-		, _boundThread(boundThread)
+        , _boundThread( boundThread )
     {
         static_assert(
             (... && !std::is_reference_v<Args>),
@@ -71,7 +71,7 @@ class RemoteFunctionBase
         auto a = AsmFactory::GetAssembler( _process.core().isWow64() );
 
         if (!contextThread)
-        	contextThread = _boundThread;
+            contextThread = _boundThread;
 
         // Ensure RPC environment exists
         status = _process.remote().CreateRPCEnvironment( Worker_None, contextThread != nullptr );
@@ -133,19 +133,19 @@ class RemoteFunctionBase
         CallArguments a( args, std::index_sequence_for<Args...>() ); 
         return Call( a, contextThread ); 
     } 
-        
+
     call_result_t<ReturnType> Call( const std::initializer_list<AsmVariant>& args, ThreadPtr contextThread = nullptr ) 
     { 
         CallArguments a( args ); 
         return Call( a, contextThread ); 
     } 
-        
+
     call_result_t<ReturnType> operator()( const Args&... args ) 
     { 
         CallArguments a( args... ); 
         return Call( a ); 
     } 
-        
+
     auto MakeArguments( const Args&... args ) 
     { 
         return CallArguments( args... ); 
@@ -155,7 +155,7 @@ class RemoteFunctionBase
     { 
         return CallArguments( args ); 
     } 
-        
+
     auto MakeArguments( const std::tuple<Args...>& args ) 
     { 
         return CallArguments( args, std::index_sequence_for<Args...>() ); 

diff --git a/src/BlackBoneTest/TestRemoteCall.cpp b/src/BlackBoneTest/TestRemoteCall.cpp
@@ -142,6 +142,47 @@ namespace Testing
             }
         }
 
+        TEST_METHOD( BoundThread )
+        {
+            Process process;
+            AssertEx::NtSuccess( process.Attach( GetCurrentProcessId() ) );
+
+            DWORD id = 0;
+            auto code = []( void* ) -> DWORD
+            {
+                for (;;)
+                    Sleep( 1 );
+
+                return ERROR_SUCCESS;
+            };
+
+            HANDLE hThread = CreateThread( nullptr, 0, code, nullptr, 0, &id );
+
+            AssertEx::IsNotNull( hThread );
+            AssertEx::IsNotZero( id );
+
+            auto thread = process.threads().get( id );
+            AssertEx::IsNotNull( thread.get() );
+
+            auto pFN = MakeRemoteFunction<decltype(&TestFn)>( process, &TestFn, thread );
+            double d = 0.0;
+
+            _input.ival = 0xDEAD;
+            _input.fval = 1337.0f;
+            _input.uval = 0xDEADC0DEA4DBEEFull;
+
+            for (auto i = 0; i < 100; i++)
+            {
+                auto [status, result] = pFN.Call( { 1, 2.0f, 3.0, &d, 5ll, _cbuf, _wbuf, &_output, _input } );
+                AssertEx::NtSuccess( status );
+                AssertEx::IsTrue( result.has_value() );
+                AssertEx::AreEqual( 1 + 5, result.value() );
+            }
+
+            TerminateThread( hThread, ERROR_SUCCESS );
+            CloseHandle( hThread );
+        }
+
         TEST_METHOD( NtQueryVirtualMemory )
         {
             auto path = GetTestHelperHost();