Fixed COM driver revokes

Author: meruune

driver_00Amethyst/DriverService.cpp

@@ -3,26 +3,14 @@
 #include <ranges>
 #include <RpcProxy.h>
 #include <shellapi.h>
-#include <wil/cppwinrt_helpers.h>
 
 #include "constants.hpp"
 #include "Logging.h"
 #include "util/color.hpp"
 
-// Wide String to UTF8 String
-inline std::string WStringToString(const std::wstring& w_str)
-{
-    const int count = WideCharToMultiByte(CP_UTF8, 0, w_str.c_str(), w_str.length(), nullptr, 0, nullptr, nullptr);
-    std::string str(count, 0);
-    WideCharToMultiByte(CP_UTF8, 0, w_str.c_str(), -1, str.data(), count, nullptr, nullptr);
-    return str;
-}
-
 DWORD DriverService::proxy_stub_registration_cookie_ = 0;
 
-DriverService::DriverService() : register_cookie_(0)
-{
-}
+DriverService::DriverService() = default;
 
 HRESULT DriverService::GetVersion(DWORD* apiVersion) noexcept
 {
@@ -36,42 +24,46 @@ HRESULT DriverService::GetVersion(DWORD* apiVersion) noexcept
 
 HRESULT DriverService::SetTrackerState(dTrackerBase tracker)
 {
+    if (tracker_vector_ == nullptr) return E_FAIL;
+
     // HMD pose override
     if (tracker.Role == TrackerHead)
         return EnableOverride(0, tracker.ConnectionState);
 
     // Normal case
-    if (tracker_vector_.contains(static_cast<ITrackerType>(tracker.Role)))
+    if (tracker_vector_->contains(static_cast<ITrackerType>(tracker.Role)))
     {
         // Create a handle to the updated (native) tracker
-        const auto p_tracker = &tracker_vector_[static_cast<ITrackerType>(tracker.Role)];
+        const auto p_tracker = &tracker_vector_->at(static_cast<ITrackerType>(tracker.Role));
 
         // Check the state and attempts spawning the tracker
         if (!p_tracker->is_added() && !p_tracker->spawn())
         {
-            logMessage(std::format("Couldn't spawn tracker with ID {} due to an unknown native exception.",
+            logMessage(std::format("Couldn't spawn tracker  ID {} due to an unknown native exception.",
                                    static_cast<int>(tracker.Role)));
             return E_FAIL; // Failure
         }
 
         // Set the state of the native tracker
         p_tracker->set_state(tracker.ConnectionState);
-        logMessage(std::format("Unmanaged (native) tracker with ID {} state has been set to {}.",
+        logMessage(std::format("Tracker ID {} state set to {}.",
                                static_cast<int>(tracker.Role), tracker.ConnectionState == 1));
 
         // Call the VR update handler and compose the result
-        tracker_vector_[static_cast<ITrackerType>(tracker.Role)].update();
+        tracker_vector_->at(static_cast<ITrackerType>(tracker.Role)).update();
         return S_OK;
     }
 
-    logMessage(std::format("Couldn't spawn tracker with ID {}. The tracker index was out of bounds.",
+    logMessage(std::format("Couldn't spawn tracker ID {}. The tracker index was out of bounds.",
                            static_cast<int>(tracker.Role)));
 
     return ERROR_INVALID_INDEX; // Failure
 }
 
 HRESULT DriverService::UpdateTracker(dTrackerBase tracker)
 {
+    if (tracker_vector_ == nullptr) return E_FAIL;
+
     // HMD pose override
     if (tracker.Role == TrackerHead)
     {
@@ -84,12 +76,12 @@ HRESULT DriverService::UpdateTracker(dTrackerBase tracker)
     }
 
     // Normal case
-    if (tracker_vector_.contains(static_cast<ITrackerType>(tracker.Role)))
+    if (tracker_vector_->contains(static_cast<ITrackerType>(tracker.Role)))
     {
         // Update the pose of the passed tracker
-        if (!tracker_vector_[static_cast<ITrackerType>(tracker.Role)].set_pose(tracker))
+        if (!tracker_vector_->at(static_cast<ITrackerType>(tracker.Role)).set_pose(tracker))
         {
-            logMessage(std::format("Couldn't spawn tracker with ID {} due to an unknown native exception.",
+            logMessage(std::format("Couldn't spawn tracker ID {} due to an unknown native exception.",
                                    static_cast<int>(tracker.Role)));
             return E_FAIL; // Failure
         }
@@ -98,7 +90,7 @@ HRESULT DriverService::UpdateTracker(dTrackerBase tracker)
         return S_OK;
     }
 
-    logMessage(std::format("Couldn't spawn tracker with ID {}. The tracker index was out of bounds.",
+    logMessage(std::format("Couldn't spawn tracker ID {}. The tracker index was out of bounds.",
                            static_cast<int>(tracker.Role)));
 
     return ERROR_INVALID_INDEX; // Failure
@@ -171,9 +163,9 @@ HRESULT DriverService::UpdateInputBoolean(dTrackerType tracker, wchar_t* path, b
         return ERROR_EMPTY; // Compose the reply
     }
 
-    if (tracker_vector_.contains(static_cast<ITrackerType>(tracker)))
-        return tracker_vector_[static_cast<ITrackerType>(tracker)]
-               .update_input(WStringToString(path), static_cast<bool>(value))
+    if (tracker_vector_->contains(static_cast<ITrackerType>(tracker)))
+        return tracker_vector_->at(static_cast<ITrackerType>(tracker))
+                              .update_input(WStringToString(path), static_cast<bool>(value))
                    ? S_OK
                    : ERROR_INVALID_ACCESS;
 
@@ -188,9 +180,9 @@ HRESULT DriverService::UpdateInputScalar(dTrackerType tracker, wchar_t* path, fl
         return ERROR_EMPTY; // Compose the reply
     }
 
-    if (tracker_vector_.contains(static_cast<ITrackerType>(tracker)))
-        return tracker_vector_[static_cast<ITrackerType>(tracker)]
-               .update_input(WStringToString(path), value)
+    if (tracker_vector_->contains(static_cast<ITrackerType>(tracker)))
+        return tracker_vector_->at(static_cast<ITrackerType>(tracker))
+                              .update_input(WStringToString(path), value)
                    ? S_OK
                    : ERROR_INVALID_ACCESS;
 
@@ -228,45 +220,28 @@ void DriverService::UninstallProxyStub()
     }
 }
 
-std::optional<winrt::hresult_error> DriverService::SetupService(const _GUID clsid)
+void DriverService::TrackerVector(std::map<ITrackerType, BodyTracker>* const& vector)
 {
-    try
-    {
-        InstallProxyStub();
-
-        winrt::com_ptr<IUnknown> service;
-        winrt::check_hresult(RegisterActiveObject(
-            static_cast<IDriverService*>(this),
-            clsid, ACTIVEOBJECT_STRONG,
-            &register_cookie_));
-
-        winrt::check_hresult(GetActiveObject(
-            clsid, nullptr, service.put()));
-    }
-    catch (const winrt::hresult_error& e)
-    {
-        return e;
-    }
-    catch (...)
-    {
-        return winrt::hresult_error(-1);
-    }
-
-    return std::nullopt;
+    tracker_vector_ = vector;
 }
 
-void DriverService::UpdateTrackers()
+void DriverService::RebuildCallback(IRebuildCallback* callback)
 {
-    for (auto& tracker : tracker_vector_ | std::views::values)
-        tracker.update(); // Update all
+    rebuild_callback_ = callback;
 }
 
-void DriverService::AddTracker(const std::string& serial, const ITrackerType role)
+ULONG DriverService::Release() noexcept
 {
-    tracker_vector_[role] = BodyTracker(serial, role);
-}
+    const auto count = implements::Release();
+    logMessage(std::format("COM ref released, running total: {}", count));
 
-std::map<ITrackerType, BodyTracker> DriverService::TrackerVector()
-{
-    return tracker_vector_;
+    if (count == 1 && rebuild_callback_)
+    {
+        logMessage("Client disconnected");
+        logMessage("COM revocation detected, requesting rebuild!");
+        rebuild_callback_->OnRebuildRequested();
+        rebuild_callback_ = nullptr; // Clear the callback
+    }
+
+    return count;
 }

driver_00Amethyst/DriverService.h

@@ -14,12 +14,24 @@
 #include "driver_Amethyst.h"
 #include "wilx.hpp"
 #include <functional>
+#include "Logging.h"
+
+namespace winrt
+{
+    struct hresult_error;
+}
 
 extern "C" {
 _Check_return_ HRESULT STDAPICALLTYPE DLLGETCLASSOBJECT_ENTRY(
     _In_ REFCLSID rclsid, _In_ REFIID riid, _Outptr_ void** ppv);
 }
 
+struct IRebuildCallback
+{
+    virtual void OnRebuildRequested() = 0;
+    virtual ~IRebuildCallback() = default;
+};
+
 class DriverService : public winrt::implements<
         DriverService, IDriverService, IVersionedApi, winrt::non_agile>
 {
@@ -52,18 +64,17 @@ class DriverService : public winrt::implements<
     static void InstallProxyStub();
     static void UninstallProxyStub();
 
-    std::optional<winrt::hresult_error> SetupService(_GUID clsid);
+    void TrackerVector(std::map<ITrackerType, BodyTracker>* const& vector);
+    void RebuildCallback(IRebuildCallback* callback);
 
-    void UpdateTrackers();
-    void AddTracker(const std::string& serial, const ITrackerType role);
-    std::map<ITrackerType, BodyTracker> TrackerVector();
+    ULONG __stdcall Release() noexcept override;
 
     void RegisterDriverPoseHandler(const std::function<HRESULT(const uint32_t& id, dDriverPose pose)>& handler);
     void RegisterOverrideSetHandler(const std::function<HRESULT(const uint32_t& id, bool isEnabled)>& handler);
 
 private:
-    DWORD register_cookie_;
-    std::map<ITrackerType, BodyTracker> tracker_vector_;
+    IRebuildCallback* rebuild_callback_ = nullptr;
+    std::map<ITrackerType, BodyTracker>* tracker_vector_;
 
     std::function<HRESULT(const uint32_t& id, dDriverPose pose)> pose_update_handler_;
     std::function<HRESULT(const uint32_t& id, bool isEnabled)> override_set_handler_;

driver_00Amethyst/Logging.h

@@ -46,4 +46,13 @@ inline void logMessageVerbose(const std::string& message, ...)
     va_end(args);
 
     OutputDebugStringA(buffer);
-}
+}
+
+// Wide String to UTF8 String
+inline std::string WStringToString(const std::wstring& w_str)
+{
+    const int count = WideCharToMultiByte(CP_UTF8, 0, w_str.c_str(), w_str.length(), nullptr, 0, nullptr, nullptr);
+    std::string str(count, 0);
+    WideCharToMultiByte(CP_UTF8, 0, w_str.c_str(), -1, str.data(), count, nullptr, nullptr);
+    return str;
+}