OrcLib: CommandAgent: synchronize with child process creation

Use CreateChildProcess, ResumeChildProcess and new notifications CommandNotification::Created and CommandNotification::Started to ensure WolfTask will handle correctly process creation even if quickly terminated.
DFIR-ORC · Nov 15, 2023 · 617faba · 617faba
1 parent e21f681
commit 617faba
Show file tree

Hide file tree

Showing 5 changed files with 56 additions and 11 deletions.
diff --git a/src/OrcCommand/Command/WolfLauncher/WolfTask.cpp b/src/OrcCommand/Command/WolfLauncher/WolfTask.cpp
@@ -49,6 +49,9 @@ HRESULT WolfTask::ApplyNotification(
             m_executableSha1 = notification->GetExecutableSha1();
             m_isSelfOrcExecutable = notification->IsSelfOrcExecutable();
             m_orcTool = notification->GetOrcTool();
+
+            actions.push_back(CommandMessage::MakeStartMessage(
+                notification->GetKeyword(), static_cast<DWORD>(static_cast<DWORD64>(notification->GetProcessID()))));
         }
         break;
         case CommandNotification::Started:

diff --git a/src/OrcLib/CommandAgent.cpp b/src/OrcLib/CommandAgent.cpp
@@ -714,6 +714,42 @@ std::shared_ptr<CommandExecute> CommandAgent::PrepareCommandExecute(const std::s
     return retval;
 }
 
+void CommandAgent::StartCommandExecute(const std::shared_ptr<CommandMessage>& message)
+{
+    Concurrency::critical_section::scoped_lock s(m_cs);
+
+    std::shared_ptr<CommandExecute> command = nullptr;
+    for (const auto& runningCommand : m_RunningCommands)
+    {
+        if (runningCommand == nullptr)
+        {
+            continue;
+        }
+
+        if (runningCommand->ProcessID() == message->ProcessID())
+        {
+            command = runningCommand;
+        }
+    }
+
+    if (command == nullptr)
+    {
+        Log::Critical(L"Command '{}' resume rejected, command not found", message->Keyword());
+        return;
+    }
+
+    HRESULT hr = command->ResumeChildProcess();
+    if (FAILED(hr))
+    {
+        Log::Critical(L"Command '{}' resume failed [{}]", message->Keyword(), SystemError(hr));
+        TerminateProcess(command->ProcessHandle(), -1);
+        return;
+    }
+
+    Concurrency::send<CommandNotification::Notification>(
+        m_target, CommandNotification::NotifyStarted(command->GetKeyword(), command->ProcessID()));
+}
+
 typedef struct _CompletionBlock
 {
     CommandAgent* pAgent;
@@ -804,6 +840,7 @@ HRESULT CommandAgent::ExecuteNextCommand()
             timer->start();
         }
 
+        // Register a callback that will handle process termination (release semaphore, notify, etc...)
         HANDLE hWaitObject = INVALID_HANDLE_VALUE;
         CompletionBlock* pBlockPtr = (CompletionBlock*)Concurrency::Alloc(sizeof(CompletionBlock));
         CompletionBlock* pBlock = new (pBlockPtr) CompletionBlock;
@@ -834,17 +871,6 @@ HRESULT CommandAgent::ExecuteNextCommand()
         notification->SetProcessCommandLine(command->m_commandLine);
 
         SendResult(notification);
-
-        hr = command->ResumeChildProcess();
-        if (FAILED(hr))
-        {
-            m_MaximumRunningSemaphore.Release();
-            command->CompleteExecution();
-            return S_OK;
-        }
-
-        Concurrency::send<CommandNotification::Notification>(
-            m_target, CommandNotification::NotifyStarted(command->GetKeyword(), command->ProcessID()));
     }
 
     return S_OK;
@@ -1230,6 +1256,11 @@ void CommandAgent::run()
                 }
             }
             break;
+            case CommandMessage::Start: {
+                Log::Debug(L"CommandAgent: Start command '{}'", request->Keyword());
+                StartCommandExecute(request);
+            }
+            break;
             case CommandMessage::RefreshRunningList: {
                 Concurrency::critical_section::scoped_lock s(m_cs);
                 auto new_end = std::remove_if(

diff --git a/src/OrcLib/CommandAgent.h b/src/OrcLib/CommandAgent.h
@@ -156,6 +156,7 @@ class CommandAgent : public Concurrency::agent
     std::shared_ptr<ProcessRedirect>
     PrepareRedirection(const std::shared_ptr<CommandExecute>& cmd, const CommandParameter& output);
     std::shared_ptr<CommandExecute> PrepareCommandExecute(const std::shared_ptr<CommandMessage>& message);
+    void StartCommandExecute(const std::shared_ptr<CommandMessage>& message);
 
     HRESULT ExecuteNextCommand();
 

diff --git a/src/OrcLib/CommandMessage.cpp b/src/OrcLib/CommandMessage.cpp
@@ -40,6 +40,14 @@ CommandMessage::Message CommandMessage::MakeCancelAnyPendingAndStopMessage()
     return retval;
 }
 
+CommandMessage::Message CommandMessage::MakeStartMessage(const std::wstring& keyword, DWORD dwProcessID)
+{
+    auto retval = std::make_shared<::CommandMessageT>(CommandMessage::Start);
+    retval->m_dwPid = dwProcessID;
+    retval->m_Keyword = keyword;
+    return retval;
+}
+
 CommandMessage::Message CommandMessage::MakeTerminateMessage(DWORD dwProcessID)
 {
     auto retval = std::make_shared<::CommandMessageT>(CommandMessage::Terminate);

diff --git a/src/OrcLib/CommandMessage.h b/src/OrcLib/CommandMessage.h
@@ -75,6 +75,7 @@ class CommandMessage
     typedef enum _Request
     {
         Execute = 0,
+        Start,
         Abort,
         Terminate,
         QueryRunningList,
@@ -112,6 +113,7 @@ class CommandMessage
 
     static Message MakeCancelMessage();
     static Message MakeAbortMessage(const std::wstring& keyword, DWORD processId, HANDLE hProcess);
+    static Message MakeStartMessage(const std::wstring& keyword, DWORD dwProcessID);
     static Message MakeTerminateMessage(DWORD dwProcessID);
     static Message MakeCancelAnyPendingAndStopMessage();
     static Message MakeTerminateAllMessage();