Merge pull request #904 from isanae/lock-job

Use a Job to capture all child processes

8 files changed, 344 insertions, 103 deletions

diff --git a/src/envmodule.cpp b/src/envmodule.cpp
index 09593e61..13831631 100644
--- a/src/envmodule.cpp
+++ b/src/envmodule.cpp
@@ -365,16 +365,13 @@ const QString& Process::name() const
 
 HandlePtr Process::openHandleForWait() const
 {
-  HandlePtr h(OpenProcess(
-    PROCESS_QUERY_LIMITED_INFORMATION | SYNCHRONIZE, FALSE, m_pid));
-
-  if (!h) {
-    const auto e = GetLastError();
-    log::error("can't get name of process {}, {}", m_pid, formatSystemMessage(e));
-    return {};
-  }
+  const auto rights =
+    PROCESS_QUERY_LIMITED_INFORMATION |    // exit code, image name, etc.
+    SYNCHRONIZE |                          // wait functions
+    PROCESS_SET_QUOTA | PROCESS_TERMINATE; // add to job
 
-  return h;
+  // don't log errors, failure can happen if the process doesn't exist
+  return HandlePtr(OpenProcess(rights, FALSE, m_pid));
 }
 
 // whether this process can be accessed; fails if the current process doesn't
@@ -405,6 +402,11 @@ std::vector<Process>& Process::children()
   return m_children;
 }
 
+const std::vector<Process>& Process::children() const
+{
+  return m_children;
+}
+
 
 std::vector<Module> getLoadedModules()
 {
@@ -531,9 +533,10 @@ void findChildren(Process& parent, const std::vector<Process>& processes)
   }
 }
 
-Process getProcessTree(HANDLE parent)
+
+Process getProcessTreeFromProcess(HANDLE h)
 {
-  const auto parentPID = ::GetProcessId(parent);
+  const auto parentPID = ::GetProcessId(h);
   const auto v = getRunningProcesses();
 
   Process root;
@@ -553,6 +556,161 @@ Process getProcessTree(HANDLE parent)
   return root;
 }
 
+
+std::vector<DWORD> processesInJob(HANDLE h)
+{
+  for (int tries=0; tries<5; ++tries) {
+    DWORD maxIds = 100;
+
+    const DWORD idsSize = sizeof(ULONG_PTR) * maxIds;
+    const DWORD bufferSize = sizeof(JOBOBJECT_BASIC_PROCESS_ID_LIST) + idsSize;
+
+    MallocPtr<void> buffer(std::malloc(bufferSize));
+    auto* ids = static_cast<JOBOBJECT_BASIC_PROCESS_ID_LIST*>(buffer.get());
+
+    const auto r = QueryInformationJobObject(
+      h, JobObjectBasicProcessIdList, ids, bufferSize, nullptr);
+
+    if (!r) {
+      const auto e = GetLastError();
+      log::error("failed to get process ids in job, {}", formatSystemMessage(e));
+      return {};
+    }
+
+    if (ids->NumberOfProcessIdsInList >= ids->NumberOfAssignedProcesses) {
+      std::vector<DWORD> v;
+      for (DWORD i=0; i<ids->NumberOfProcessIdsInList; ++i) {
+        v.push_back(ids->ProcessIdList[i]);
+      }
+
+      return v;
+    }
+
+    // try again with a larger buffer
+    maxIds *= 2;
+  }
+
+  log::error("failed to get processes in job, can't get a buffer large enough");
+  return {};
+}
+
+
+void findChildProcesses(Process& parent, std::vector<Process>& processes)
+{
+  // find all processes that are direct children of `parent`
+  auto itor = processes.begin();
+
+  while (itor != processes.end()) {
+    if (itor->ppid() == parent.pid()) {
+      parent.addChild(*itor);
+      itor = processes.erase(itor);
+    } else {
+      ++itor;
+    }
+  }
+
+  // find all processes that are direct children of `parent`'s children
+  for (auto&& c : parent.children()) {
+    findChildProcesses(c, processes);
+  }
+}
+
+Process getProcessTreeFromJob(HANDLE h)
+{
+  const auto ids = processesInJob(h);
+  if (ids.empty()) {
+    return {};
+  }
+
+  std::vector<Process> ps;
+
+  forEachRunningProcess([&](auto&& entry) {
+    for (auto&& id : ids) {
+      if (entry.th32ProcessID == id) {
+        ps.push_back(Process(
+          entry.th32ProcessID,
+          entry.th32ParentProcessID,
+          QString::fromStdWString(entry.szExeFile)));
+
+        break;
+      }
+    }
+
+    return true;
+  });
+
+  Process root;
+
+  {
+    // getting processes whose parent is not in the list
+    for (auto&& possibleRoot : ps) {
+      const auto ppid = possibleRoot.ppid();
+      bool found = false;
+
+      for (auto&& p : ps) {
+        if (p.pid() == ppid) {
+          found = true;
+          break;
+        }
+      }
+
+      if (!found) {
+        // this is a root process
+        root.addChild(possibleRoot);
+      }
+    }
+
+    // removing root processes from the list
+    auto newEnd = std::remove_if(ps.begin(), ps.end(), [&](auto&& p) {
+      for (auto&& rp : root.children()) {
+        if (rp.pid() == p.pid()) {
+          return true;
+        }
+      }
+
+      return false;
+    });
+
+    ps.erase(newEnd, ps.end());
+  }
+
+  // at this point, `processes` should only contain processes that are direct
+  // or indirect children of the ones in `root`
+
+  if (ps.empty()) {
+    // and that's all there is
+    return root;
+  }
+
+  {
+    // recursively find children
+    for (auto&& r : root.children()) {
+      findChildProcesses(r, ps);
+    }
+  }
+
+  return root;
+}
+
+bool isJobHandle(HANDLE h)
+{
+  JOBOBJECT_BASIC_ACCOUNTING_INFORMATION info = {};
+
+  const auto r = ::QueryInformationJobObject(
+    h, JobObjectBasicAccountingInformation, &info, sizeof(info), nullptr);
+
+  return r;
+}
+
+Process getProcessTree(HANDLE h)
+{
+  if (isJobHandle(h)) {
+    return getProcessTreeFromJob(h);
+  } else {
+    return getProcessTreeFromProcess(h);
+  }
+}
+
 QString getProcessName(DWORD pid)
 {
   HandlePtr h(OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, pid));
diff --git a/src/envmodule.h b/src/envmodule.h
index d152b840..c2829b32 100644
--- a/src/envmodule.h
+++ b/src/envmodule.h
@@ -136,6 +136,7 @@ public:
 
   void addChild(Process p);
   std::vector<Process>& children();
+  const std::vector<Process>& children() const;
 
 private:
   DWORD m_pid;
@@ -148,7 +149,9 @@ private:
 std::vector<Process> getRunningProcesses();
 std::vector<Module> getLoadedModules();
 
-Process getProcessTree(HANDLE parent);
+// works for both jobs and processes
+//
+Process getProcessTree(HANDLE h);
 
 QString getProcessName(DWORD pid);
 QString getProcessName(HANDLE process);
diff --git a/src/loot.cpp b/src/loot.cpp
index 47a70596..a86cdcff 100644
--- a/src/loot.cpp
+++ b/src/loot.cpp
@@ -385,8 +385,8 @@ QString Loot::Message::toMarkdown() const
 }
 
 
-Loot::Loot()
-  : m_thread(nullptr), m_cancel(false), m_result(false)
+Loot::Loot(OrganizerCore& core)
+  : m_core(core), m_thread(nullptr), m_cancel(false), m_result(false)
 {
 }
 
@@ -408,7 +408,7 @@ Loot::~Loot()
   }
 }
 
-bool Loot::start(QWidget* parent, OrganizerCore& core, bool didUpdateMasterList)
+bool Loot::start(QWidget* parent, bool didUpdateMasterList)
 {
   log::debug("starting loot");
 
@@ -420,10 +420,10 @@ bool Loot::start(QWidget* parent, OrganizerCore& core, bool didUpdateMasterList)
   }
 
   // vfs
-  core.prepareVFS();
+  m_core.prepareVFS();
 
   // spawning
-  if (!spawnLootcli(parent, core, didUpdateMasterList, std::move(stdoutHandle))) {
+  if (!spawnLootcli(parent, didUpdateMasterList, std::move(stdoutHandle))) {
     return false;
   }
 
@@ -436,19 +436,29 @@ bool Loot::start(QWidget* parent, OrganizerCore& core, bool didUpdateMasterList)
 }
 
 bool Loot::spawnLootcli(
-  QWidget* parent, OrganizerCore& core, bool didUpdateMasterList,
-  env::HandlePtr stdoutHandle)
+  QWidget* parent, bool didUpdateMasterList, env::HandlePtr stdoutHandle)
 {
-  const auto logLevel = core.settings().diagnostics().lootLogLevel();
+  const auto logLevel = m_core.settings().diagnostics().lootLogLevel();
 
   QStringList parameters;
   parameters
-    << "--game" << core.managedGame()->gameShortName()
-    << "--gamePath" << QString("\"%1\"").arg(core.managedGame()->gameDirectory().absolutePath())
-    << "--pluginListPath" << QString("\"%1/loadorder.txt\"").arg(core.profilePath())
-    << "--logLevel" << QString::fromStdString(lootcli::logLevelToString(logLevel))
-    << "--out" << QString("\"%1\"").arg(LootReportPath)
-    << "--language" << core.settings().interface().language();
+    << "--game"
+    << m_core.managedGame()->gameShortName()
+
+    << "--gamePath"
+    << QString("\"%1\"").arg(m_core.managedGame()->gameDirectory().absolutePath())
+
+    << "--pluginListPath"
+    << QString("\"%1/loadorder.txt\"").arg(m_core.profilePath())
+
+    << "--logLevel"
+    << QString::fromStdString(lootcli::logLevelToString(logLevel))
+
+    << "--out"
+    << QString("\"%1\"").arg(LootReportPath)
+
+    << "--language"
+    << m_core.settings().interface().language();
 
   if (didUpdateMasterList) {
     parameters << "--skipUpdateMasterlist";
@@ -696,6 +706,12 @@ Loot::Plugin Loot::reportPlugin(const QJsonObject& plugin) const
     return {};
   }
 
+  // ignore disabled plugins; lootcli doesn't know if a plugin is enabled or not
+  // and will report information on any plugin that's in the filesystem
+  if (!m_core.pluginList()->isEnabled(p.name)) {
+    return {};
+  }
+
   if (plugin.contains("incompatibilities")) {
     p.incompatibilities = reportFiles(getOpt<QJsonArray>(plugin, "incompatibilities"));
   }
@@ -836,10 +852,10 @@ bool runLoot(QWidget* parent, OrganizerCore& core, bool didUpdateMasterList)
   core.savePluginList();
 
   try {
-    Loot loot;
+    Loot loot(core);
     LootDialog dialog(parent, core, loot);
 
-    if (!loot.start(parent, core, didUpdateMasterList)) {
+    if (!loot.start(parent, didUpdateMasterList)) {
       return false;
     }
 
diff --git a/src/loot.h b/src/loot.h
index 7e5e01bd..4ec06d6f 100644
--- a/src/loot.h
+++ b/src/loot.h
@@ -79,10 +79,10 @@ public:
   };
 
 
-  Loot();
+  Loot(OrganizerCore& core);
   ~Loot();
 
-  bool start(QWidget* parent, OrganizerCore& core, bool didUpdateMasterList);
+  bool start(QWidget* parent, bool didUpdateMasterList);
   void cancel();
   bool result() const;
   const QString& outPath() const;
@@ -95,6 +95,7 @@ signals:
   void finished();
 
 private:
+  OrganizerCore& m_core;
   std::unique_ptr<QThread> m_thread;
   std::atomic<bool> m_cancel;
   std::atomic<bool> m_result;
@@ -104,8 +105,7 @@ private:
   Report m_report;
 
   bool spawnLootcli(
-    QWidget* parent, OrganizerCore& core, bool didUpdateMasterList,
-    env::HandlePtr stdoutHandle);
+    QWidget* parent, bool didUpdateMasterList, env::HandlePtr stdoutHandle);
 
   void lootThread();
   bool waitForCompletion();
diff --git a/src/processrunner.cpp b/src/processrunner.cpp
index cfddbb63..ce1ac681 100644
--- a/src/processrunner.cpp
+++ b/src/processrunner.cpp
@@ -106,21 +106,46 @@ QString toString(Interest i)
 }
 
 
+struct InterestingProcess
+{
+  env::Process p;
+  Interest interest = Interest::None;
+  env::HandlePtr handle;
+};
+
+
+InterestingProcess findRandomProcess(const env::Process& root)
+{
+  for (auto&& c : root.children()) {
+    env::HandlePtr h = c.openHandleForWait();
+    if (h) {
+      return {c, Interest::Weak, std::move(h)};
+    }
+
+    auto r = findRandomProcess(c);
+    if (r.handle) {
+      return r;
+    }
+  }
+
+  return {};
+}
+
 // returns a process that's in the hidden list, or the top-level process if
 // they're all hidden; returns an invalid process if the list is empty
 //
-std::pair<env::Process, Interest> findInterestingProcessInTrees(
-  std::vector<env::Process>& processes)
+InterestingProcess findInterestingProcessInTrees(const env::Process& root)
 {
-  if (processes.empty()) {
-    return {{}, Interest::None};
-  }
-
   // Certain process names we wish to "hide" for aesthetic reason:
-  const std::vector<QString> hiddenList = {
-    QFileInfo(QCoreApplication::applicationFilePath()).fileName()
+  static const std::vector<QString> hiddenList = {
+    QFileInfo(QCoreApplication::applicationFilePath()).fileName(),
+    "conhost.exe"
   };
 
+  if (root.children().empty()) {
+    return {};
+  }
+
   auto isHidden = [&](auto&& p) {
     for (auto h : hiddenList) {
       if (p.name().contains(h, Qt::CaseInsensitive)) {
@@ -132,54 +157,61 @@ std::pair<env::Process, Interest> findInterestingProcessInTrees(
   };
 
 
-  for (auto&& root : processes) {
-    if (!isHidden(root)) {
-      return {root, Interest::Strong};
+  for (auto&& p : root.children()) {
+    if (!isHidden(p)) {
+      env::HandlePtr h = p.openHandleForWait();
+      if (h) {
+        return {p, Interest::Strong, std::move(h)};
+      }
     }
 
-    for (auto&& child : root.children()) {
-      if (!isHidden(child)) {
-        return {child, Interest::Strong};
-      }
+    auto r = findInterestingProcessInTrees(p);
+    if (r.interest == Interest::Strong) {
+      return r;
     }
   }
 
-  // everything is hidden, just pick the first one
-  return {processes[0], Interest::Weak};
+  // everything is hidden, just pick the first one that can be used
+  return findRandomProcess(root);
 }
 
-// gets the most interesting process in the list
-//
-std::pair<env::Process, Interest> getInterestingProcess(
-  const std::vector<HANDLE>& initialProcesses)
+void dump(const env::Process& p, int indent)
 {
-  if (initialProcesses.empty()) {
-    log::debug("nothing to wait for");
-    return {{}, Interest::None};
+  log::debug(
+    "{}{}, pid={}, ppid={}",
+    std::string(indent * 4, ' '), p.name(), p.pid(), p.ppid());
+
+  for (auto&& c : p.children()) {
+    dump(c, indent + 1);
   }
+}
 
-  std::vector<env::Process> processes;
+void dump(const env::Process& root)
+{
+  log::debug("process tree:");
 
-  // getting process trees for all processes
-  for (auto&& h : initialProcesses) {
-    auto tree = env::getProcessTree(h);
-    if (tree.isValid()) {
-      processes.push_back(tree);
-    }
+  for (auto&& p : root.children()) {
+    dump(p, 1);
   }
+}
 
-  if (processes.empty()) {
-    // if the initial list wasn't empty but this one is, it means all the
-    // processes were already completed
-    log::debug("processes are already completed");
-    return {{}, Interest::None};
+// gets the most interesting process in the list
+//
+InterestingProcess getInterestingProcess(HANDLE job)
+{
+  env::Process root = env::getProcessTree(job);
+  if (root.children().empty()) {
+    log::debug("nothing to wait for");
+    return {};
   }
 
-  const auto interest = findInterestingProcessInTrees(processes);
-  if (!interest.first.isValid()) {
-    // this shouldn't happen
+  dump(root);
+
+  auto interest = findInterestingProcessInTrees(root);
+  if (!interest.handle) {
+    // this can happen if none of the processes can be opened
     log::debug("no interesting process to wait for");
-    return {{}, Interest::None};
+    return {};
   }
 
   return interest;
@@ -255,56 +287,52 @@ std::optional<ProcessRunner::Results> timedWait(
 }
 
 ProcessRunner::Results waitForProcessesThreadImpl(
-  const std::vector<HANDLE>& initialProcesses, UILocker::Session& ls)
+  HANDLE job, UILocker::Session& ls)
 {
   using namespace std::chrono;
 
-  if (initialProcesses.empty()) {
-    // shouldn't happen
-    return ProcessRunner::Completed;
-  }
-
   DWORD currentPID = 0;
 
   // if the interesting process that was found is weak (such as ModOrganizer.exe
   // when starting a program from within the Data directory), start with a short
   // wait and check for more interesting children
-  milliseconds wait(50);
+  const milliseconds defaultWait(50);
+  auto wait = defaultWait;
 
   for (;;) {
-    auto [p, interest] = getInterestingProcess(initialProcesses);
-    if (!p.isValid()) {
+    auto ip = getInterestingProcess(job);
+    if (!ip.handle) {
       // nothing to wait on
       return ProcessRunner::Completed;
     }
 
     // update the lock widget
-    ls.setInfo(p.pid(), p.name());
+    ls.setInfo(ip.p.pid(), ip.p.name());
 
-    // open the process
-    auto interestingHandle = p.openHandleForWait();
-    if (!interestingHandle) {
-      return ProcessRunner::Error;
-    }
-
-    if (p.pid() != currentPID) {
+    if (ip.p.pid() != currentPID) {
       // log any change in the process being waited for
-      currentPID = p.pid();
+      currentPID = ip.p.pid();
 
       log::debug(
         "waiting for completion on {} ({}), {} interest",
-        p.name(), p.pid(), toString(interest));
+        ip.p.name(), ip.p.pid(), toString(ip.interest));
     }
 
-    if (interest == Interest::Strong) {
+    if (ip.interest == Interest::Strong) {
       // don't bother with short wait, this is a good process to wait for
       wait = Infinite;
     }
 
-    const auto r = timedWait(interestingHandle.get(), p.pid(), ls, wait);
+    const auto r = timedWait(ip.handle.get(), ip.p.pid(), ls, wait);
     if (r) {
-      // the process has completed or returned an error
-      return *r;
+      if (*r == ProcessRunner::Results::Completed) {
+        // process completed, check another one, reset the wait time to find
+        // interesting processes
+        wait = defaultWait;
+      } else if (*r != ProcessRunner::Results::Running) {
+        // something's wrong, or the user unlocked the ui
+        return *r;
+      }
     }
 
     // exponentially increase the wait time between checks for interesting
@@ -314,20 +342,44 @@ ProcessRunner::Results waitForProcessesThreadImpl(
 }
 
 void waitForProcessesThread(
-  ProcessRunner::Results& result,
-  const std::vector<HANDLE>& initialProcesses, UILocker::Session& ls)
+  ProcessRunner::Results& result, HANDLE job, UILocker::Session& ls)
 {
-  result = waitForProcessesThreadImpl(initialProcesses, ls);
+  result = waitForProcessesThreadImpl(job, ls);
   ls.unlock();
 }
 
 ProcessRunner::Results waitForProcesses(
   const std::vector<HANDLE>& initialProcesses, UILocker::Session& ls)
 {
+  // using a job so any child process started by any of those processes can also
+  // be captured and monitored
+  env::HandlePtr job(CreateJobObjectW(nullptr, nullptr));
+  if (!job) {
+    const auto e = GetLastError();
+
+    log::error(
+      "failed to create job to wait for processes, {}",
+      formatSystemMessage(e));
+
+    return ProcessRunner::Error;
+  }
+
+  for (auto&& h : initialProcesses) {
+    if (!::AssignProcessToJobObject(job.get(), h)) {
+      const auto e = GetLastError();
+
+      log::error(
+        "can't assign process to job to wait for processes, {}",
+        formatSystemMessage(e));
+
+      // keep going
+    }
+  }
+
   auto results = ProcessRunner::Running;
 
   auto* t = QThread::create(
-    waitForProcessesThread, std::ref(results), initialProcesses, std::ref(ls));
+    waitForProcessesThread, std::ref(results), job.get(), std::ref(ls));
 
   QEventLoop events;
   QObject::connect(t, &QThread::finished, [&]{
diff --git a/src/uilocker.cpp b/src/uilocker.cpp
index 01794d6b..8c4e0c2c 100644
--- a/src/uilocker.cpp
+++ b/src/uilocker.cpp
@@ -377,6 +377,8 @@ UILocker::~UILocker()
       unlock(s.get());
     }
   }
+
+  g_instance = nullptr;
 }
 
 UILocker& UILocker::instance()
@@ -449,6 +451,12 @@ void UILocker::unlockCurrent()
 
 void UILocker::updateLabel()
 {
+  if (!m_ui) {
+    // this can happen if the lock overlay was destroyed while a cross-thread
+    // call for updateLabel() was in flight
+    return;
+  }
+
   QStringList labels;
 
   for (auto itor=m_sessions.rbegin(); itor!=m_sessions.rend(); ++itor) {
diff --git a/src/usvfsconnector.cpp b/src/usvfsconnector.cpp
index 3dba3efc..3c8c355b 100644
--- a/src/usvfsconnector.cpp
+++ b/src/usvfsconnector.cpp
@@ -280,13 +280,17 @@ std::vector<HANDLE> getRunningUSVFSProcesses()
   const auto thisPid = GetCurrentProcessId();
   std::vector<HANDLE> v;
 
+  const auto rights =
+    PROCESS_QUERY_LIMITED_INFORMATION |    // exit code, image name, etc.
+    SYNCHRONIZE |                          // wait functions
+    PROCESS_SET_QUOTA | PROCESS_TERMINATE; // add to job
+
   for (auto&& pid : pids) {
     if (pid == thisPid) {
       continue; // obviously don't wait for MO process
     }
 
-    HANDLE handle = ::OpenProcess(
-      PROCESS_QUERY_LIMITED_INFORMATION | SYNCHRONIZE, FALSE, pid);
+    HANDLE handle = ::OpenProcess(rights, FALSE, pid);
 
     if (handle == INVALID_HANDLE_VALUE) {
       const auto e = GetLastError();
diff --git a/src/version.rc b/src/version.rc
index 7dac7b3a..963d1f11 100644
--- a/src/version.rc
+++ b/src/version.rc
@@ -3,8 +3,8 @@
 // If VS_FF_PRERELEASE is not set, MO labels the build as a release and uses VER_FILEVERSION to determine version number.
 // Otherwise, if letters are used in VER_FILEVERSION_STR, uses the full MOBase::VersionInfo parser
 // Otherwise, uses the numbers from VER_FILEVERSION and sets the release type as pre-alpha
-#define VER_FILEVERSION     2,2,2,7
-#define VER_FILEVERSION_STR "2.2.2alpha7.1\0"
+#define VER_FILEVERSION     2,2,2,8
+#define VER_FILEVERSION_STR "2.2.2alpha8\0"
 
 VS_VERSION_INFO VERSIONINFO
 FILEVERSION     VER_FILEVERSION