Fix ownership of WTS plugin references

dll/RdpInstance.cpp

@@ -49,6 +49,10 @@ class CMsRdpExInstance : public IMsRdpExInstance
         if (m_pMsRdpExtendedSettings) {
             m_pMsRdpExtendedSettings->Release();
         }
+
+        if (m_WTSPlugin) {
+            m_WTSPlugin->Release();
+        }
     }
 
     // IUnknown interface
@@ -537,7 +541,11 @@ class CMsRdpExInstance : public IMsRdpExInstance
 
     HRESULT STDMETHODCALLTYPE SetWTSPluginObject(LPVOID pvObject)
     {
+        IUnknown* previousPlugin = m_WTSPlugin;
         m_WTSPlugin = (IUnknown*)pvObject;
+        if (previousPlugin) {
+            previousPlugin->Release();
+        }
         return S_OK;
     }
 

include/MsRdpEx/RdpInstance.h

@@ -37,6 +37,7 @@ struct __declspec(novtable)
     virtual void __stdcall UnlockShadowBitmap() = 0;
     virtual void __stdcall GetLastMousePosition(int32_t* posX, int32_t* posY) = 0;
     virtual void __stdcall SetLastMousePosition(int32_t posX, int32_t posY) = 0;
+    // Get returns a borrowed pointer; Set consumes one owned reference, even for the same pointer.
     virtual HRESULT __stdcall GetWTSPluginObject(LPVOID* ppvObject) = 0;
     virtual HRESULT __stdcall SetWTSPluginObject(LPVOID pvObject) = 0;
     virtual void __stdcall SetCursor(HCURSOR cursor) = 0;

tests/logging/CMakeLists.txt

@@ -33,7 +33,7 @@ endforeach()
 # a completion signal after DllMain returns. No test exports enter the shipped DLL.
 get_target_property(gateway_shutdown_sources MsRdpEx_Dll SOURCES)
 list(FILTER gateway_shutdown_sources INCLUDE REGEX "\\.(c|cpp)$")
-list(FILTER gateway_shutdown_sources EXCLUDE REGEX "(^|/)(GatewayIsolation|MsRdpEx)[.]cpp$")
+list(FILTER gateway_shutdown_sources EXCLUDE REGEX "(^|/)(GatewayIsolation|MsRdpEx|RdpInstance)[.]cpp$")
 set(gateway_shutdown_absolute_sources)
 foreach(source IN LISTS gateway_shutdown_sources)
     if(NOT IS_ABSOLUTE "${source}")
@@ -43,7 +43,8 @@ foreach(source IN LISTS gateway_shutdown_sources)
 endforeach()
 get_target_property(gateway_shutdown_libraries MsRdpEx_Dll LINK_LIBRARIES)
 add_library(MsRdpEx_GatewayShutdownTestDll SHARED ${gateway_shutdown_absolute_sources}
-    GatewayShutdownFixture.cpp GatewayShutdownDllMain.cpp GatewayShutdownFixture.def)
+    GatewayShutdownFixture.cpp GatewayShutdownDllMain.cpp GatewayShutdownFixture.def
+    PluginReferenceFixture.cpp)
 target_include_directories(MsRdpEx_GatewayShutdownTestDll PRIVATE "${PROJECT_SOURCE_DIR}/dll")
 target_link_libraries(MsRdpEx_GatewayShutdownTestDll PRIVATE ${gateway_shutdown_libraries})
 target_compile_options(MsRdpEx_GatewayShutdownTestDll PRIVATE /W4)
@@ -56,3 +57,15 @@ foreach(mode exit unload exit-held-lock)
         COMMAND MsRdpEx_GatewayShutdownTest "$<TARGET_FILE:MsRdpEx_GatewayShutdownTestDll>" ${mode})
     set_tests_properties(logging.gateway-shutdown.${mode} PROPERTIES TIMEOUT 45)
 endforeach()
+
+add_executable(MsRdpEx_PluginReferenceTest PluginReferenceTest.cpp)
+target_include_directories(MsRdpEx_PluginReferenceTest PRIVATE "${PROJECT_SOURCE_DIR}/dll")
+target_compile_features(MsRdpEx_PluginReferenceTest PRIVATE cxx_std_17)
+target_compile_options(MsRdpEx_PluginReferenceTest PRIVATE /W4 /EHsc)
+add_dependencies(MsRdpEx_PluginReferenceTest MsRdpEx_GatewayShutdownTestDll)
+foreach(scenario replacement same-pointer clearing destruction)
+    add_test(NAME logging.plugin-reference.${scenario}
+        COMMAND MsRdpEx_PluginReferenceTest
+            "$<TARGET_FILE:MsRdpEx_GatewayShutdownTestDll>" ${scenario})
+    set_tests_properties(logging.plugin-reference.${scenario} PROPERTIES TIMEOUT 45)
+endforeach()

tests/logging/GatewayShutdownFixture.def

@@ -1,2 +1,3 @@
 EXPORTS
     PrepareGatewayShutdown
+    CreatePluginReferenceInstance

tests/logging/PluginReferenceFixture.cpp

@@ -0,0 +1,6 @@
+#include "../../dll/RdpInstance.cpp"
+
+extern "C" IMsRdpExInstance* CreatePluginReferenceInstance()
+{
+    return CMsRdpExInstance_New(NULL);
+}

tests/logging/PluginReferenceTest.cpp

@@ -0,0 +1,95 @@
+#include <MsRdpEx/RdpInstance.h>
+
+#include <iostream>
+#include <stdexcept>
+#include <string>
+
+static void Check(bool condition, const char* message)
+{
+    if (!condition) throw std::runtime_error(message);
+}
+
+class CountedUnknown : public IUnknown
+{
+public:
+    HRESULT STDMETHODCALLTYPE QueryInterface(REFIID iid, void** object) override
+    {
+        if (!object) return E_POINTER;
+        *object = NULL;
+        if (iid != IID_IUnknown) return E_NOINTERFACE;
+        *object = static_cast<IUnknown*>(this);
+        AddRef();
+        return S_OK;
+    }
+
+    ULONG STDMETHODCALLTYPE AddRef() override { return ++refs; }
+    ULONG STDMETHODCALLTYPE Release() override { return --refs; }
+
+    ULONG refs = 1; // The test keeps its own reference; setters receive another.
+};
+
+int wmain(int argc, wchar_t** argv)
+{
+    try {
+        Check(argc == 3, "Expected test DLL path and scenario");
+        Check(SetEnvironmentVariableW(L"MSRDPEX_HOOK_ENABLED", L"0") != 0,
+            "Could not disable ActiveX hooks");
+        Check(SetEnvironmentVariableW(L"MSRDPEX_LOG_ENABLED", L"0") != 0,
+            "Could not disable logging");
+        HMODULE module = LoadLibraryW(argv[1]);
+        Check(module != NULL, "Could not load test DLL");
+        using Create = IMsRdpExInstance*(*)();
+        auto create = reinterpret_cast<Create>(
+            GetProcAddress(module, "CreatePluginReferenceInstance"));
+        Check(create != NULL, "Plugin instance fixture export missing");
+        IMsRdpExInstance* instance = create();
+        Check(instance != NULL, "Could not create plugin instance");
+
+        CountedUnknown first;
+        first.AddRef();
+        Check(SUCCEEDED(instance->SetWTSPluginObject(&first)), "Initial plugin setter failed");
+        void* borrowed = NULL;
+        Check(SUCCEEDED(instance->GetWTSPluginObject(&borrowed)) && borrowed == &first,
+            "Plugin getter did not return the registered pointer");
+
+        const std::wstring scenario = argv[2];
+        if (scenario == L"replacement") {
+            CountedUnknown second;
+            second.AddRef();
+            Check(SUCCEEDED(instance->SetWTSPluginObject(&second)), "Replacement setter failed");
+            Check(first.refs == 1, "Replacement did not release the previous plugin reference");
+            Check(second.refs == 2, "Replacement did not retain the transferred reference");
+            instance->Release();
+            Check(second.refs == 1, "Instance destruction did not release replacement");
+            Check(first.Release() == 0 && second.Release() == 0, "Test references not balanced");
+        } else if (scenario == L"same-pointer") {
+            first.AddRef();
+            Check(SUCCEEDED(instance->SetWTSPluginObject(&first)), "Repeated setter failed");
+            Check(first.refs == 2, "Repeating the pointer leaked a separate transferred reference");
+            instance->Release();
+            Check(first.refs == 1, "Instance destruction did not release the plugin");
+            Check(first.Release() == 0, "Test reference not balanced");
+        } else if (scenario == L"clearing") {
+            Check(SUCCEEDED(instance->SetWTSPluginObject(NULL)), "Plugin clearing failed");
+            Check(first.refs == 1, "Clearing did not release the plugin");
+            borrowed = &first;
+            Check(SUCCEEDED(instance->GetWTSPluginObject(&borrowed)) && !borrowed,
+                "Clearing did not remove the plugin");
+            instance->Release();
+            Check(first.Release() == 0, "Test reference not balanced");
+        } else if (scenario == L"destruction") {
+            instance->Release();
+            Check(first.refs == 1, "Instance destruction did not release the plugin");
+            Check(first.Release() == 0, "Test reference not balanced");
+        } else {
+            throw std::runtime_error("Unknown plugin reference scenario");
+        }
+
+        Check(FreeLibrary(module) != FALSE, "Could not unload test DLL");
+        std::wcout << L"PASS plugin reference: " << scenario << L'\n';
+        return 0;
+    } catch (const std::exception& error) {
+        std::cerr << error.what() << '\n';
+        return 1;
+    }
+}

tests/logging/README.md

@@ -1,6 +1,7 @@
-# Native logging regression tests
+# Native regression tests
 
-These tests exercise runtime diagnostic logging without making an RDP connection.
+These tests exercise runtime diagnostic logging, gateway behavior, and plugin
+reference ownership without making an RDP connection.
 They are opt-in and do not change normal builds or release packaging.
 
 ## Build and run
@@ -19,9 +20,10 @@ The test process must run on Windows with support for its target architecture.
 If multiple Visual Studio installations exist, select the one with ATL using
 `-DCMAKE_GENERATOR_INSTANCE="C:/Program Files/Microsoft Visual Studio/2022/Community"`.
 
-There are ten tests per architecture: five scenarios, each with ActiveX hooks
-disabled and enabled. CTest applies a 45-second timeout to each case. Scenarios
-retain their temporary log files and print the directory on success or failure.
+The ten logging tests cover five scenarios, each with ActiveX hooks disabled
+and enabled. CTest also runs gateway and plugin reference tests. Each case has
+a 45-second timeout. Logging scenarios retain their temporary log files and
+print the directory on success or failure.
 
 ## Reproduce against another DLL
 
@@ -71,3 +73,12 @@ not synchronous.
 
 The existing COM interface and managed API remain unchanged, and the setters
 still return `void`.
+
+## WTS plugin reference ownership
+
+The `logging.plugin-reference.*` cases use a test-only factory in the native
+fixture DLL to create a real `CMsRdpExInstance` without an RDP control. A counted
+`IUnknown` verifies that `SetWTSPluginObject` consumes an owned reference and
+releases it when replaced (including by another reference to the same object),
+cleared, or destroyed. The plugin getter returns a borrowed pointer. The
+fixture export is not included in the production DLL.