Fix COM reference counting

Author: mamoreau-devolutions

dll/AxHost/RdpAxHostWnd.cpp

@@ -120,10 +120,10 @@ class CRdpAxHostWnd : public IUnknown
         if (riid == IID_IUnknown) {
             *ppv = this;
         }
-        else if (riid == IID_IOleClientSite) {
+        else if (riid == IID_IOleClientSite && m_pOleClientSite) {
             *ppv = (void*)m_pOleClientSite;
         }
-        else if (riid == IID_IOleInPlaceSiteEx) {
+        else if (riid == IID_IOleInPlaceSiteEx && m_pOleInPlaceSiteEx) {
             *ppv = (void*)m_pOleInPlaceSiteEx;
         }
 
@@ -687,7 +687,7 @@ class CRdpAxHostWnd : public IUnknown
     HANDLE m_stopEvent = NULL;
 
 private:
-    LONG m_refCount = 0;
+    LONG m_refCount = 1;
     HWND m_hWnd = NULL;
     HWND m_hWndParent = NULL;
     DWORD m_dwAdviseCookie = 0;

dll/AxHost/RdpEventSink.cpp

@@ -164,7 +164,7 @@ STDMETHODIMP CRdpEventSink::OnConfirmClose(VARIANT_BOOL* pfAllowClose)
 
 CRdpEventSink::CRdpEventSink(HWND hWndParent)
 {
-    m_refCount = 0;
+    m_refCount = 1;
     m_hWndParent = hWndParent;
 }
 

dll/AxHost/RdpOleSite.cpp

@@ -3,7 +3,7 @@
 
 CRdpOleClientSite::CRdpOleClientSite(IUnknown* pUnkOuter)
 {
-    m_refCount = 0;
+    m_refCount = 1;
     m_pUnkOuter = pUnkOuter;
     m_pUnkOuter->AddRef();
 }
@@ -91,7 +91,7 @@ STDMETHODIMP CRdpOleClientSite::RequestNewObjectLayout()
 }
 
 CRdpOleInPlaceSiteEx::CRdpOleInPlaceSiteEx(IUnknown* pUnkOuter)
-    : m_refCount(0), m_hWnd(0)
+    : m_refCount(1), m_hWnd(0)
 {
     m_pUnkOuter = pUnkOuter;
     m_pUnkOuter->AddRef();

dll/MsRdpClient.cpp

@@ -150,8 +150,9 @@ class CMsRdpClient : public IMsRdpClient10
 public:
     CMsRdpClient(IUnknown* pUnknown)
     {
-        m_refCount = 0;
+        m_refCount = 1;
         m_pUnknown = pUnknown;
+        pUnknown->AddRef();
 
         pUnknown->QueryInterface(IID_IDispatch, (LPVOID*)&m_pDispatch);
         pUnknown->QueryInterface(IID_IMsTscAx, (LPVOID*)&m_pMsTscAx);
@@ -168,13 +169,10 @@ class CMsRdpClient : public IMsRdpClient10
 
         m_pMsRdpExInstance = CMsRdpExInstance_New(this);
         IMsRdpExInstance* pMsRdpExInstance = (IMsRdpExInstance*)m_pMsRdpExInstance;
-        pMsRdpExInstance->AddRef();
         pMsRdpExInstance->GetSessionId(&m_sessionId);
         MsRdpEx_InstanceManager_Add(m_pMsRdpExInstance);
 
         m_pMsRdpExtendedSettings = CMsRdpExtendedSettings_New(pUnknown, (IUnknown*)m_pMsTscAx, &m_sessionId);
-        IMsRdpExtendedSettings* pMsRdpExtendedSettings = (IMsRdpExtendedSettings*)m_pMsRdpExtendedSettings;
-        pMsRdpExtendedSettings->AddRef();
         pMsRdpExInstance->AttachExtendedSettings(m_pMsRdpExtendedSettings);
 
         void* pCorePropsRaw = NULL;
@@ -205,10 +203,9 @@ class CMsRdpClient : public IMsRdpClient10
         }
 
         if (m_pMsRdpExInstance) {
-            IMsRdpExInstance* pMsRdpExInstance = (IMsRdpExInstance*)m_pMsRdpExInstance;
             MsRdpEx_InstanceManager_Remove(m_pMsRdpExInstance);
-            pMsRdpExInstance->Release();
-            pMsRdpExInstance = NULL;
+            ((IMsRdpExInstance*)m_pMsRdpExInstance)->Release();
+            m_pMsRdpExInstance = NULL;
         }
     }
 
@@ -715,7 +712,7 @@ class CMsRdpClient : public IMsRdpClient10
 
 private:
     GUID m_sessionId;
-    ULONG m_refCount = 0;
+    ULONG m_refCount;
     IUnknown* m_pUnknown = NULL;
     IDispatch* m_pDispatch = NULL;
     IMsTscAx* m_pMsTscAx = NULL;
@@ -762,12 +759,12 @@ class CClassFactory : IClassFactory
 
         if (riid == IID_IUnknown) {
             *ppvObject = (LPVOID)((IUnknown*)this);
-            m_refCount++;
+            InterlockedIncrement(&m_refCount);
             return S_OK;
         }
         if (riid == IID_IClassFactory) {
             *ppvObject = (LPVOID)((IClassFactory*)this);
-            m_refCount++;
+            InterlockedIncrement(&m_refCount);
             return S_OK;
         }
 
@@ -776,19 +773,21 @@ class CClassFactory : IClassFactory
 
     ULONG STDMETHODCALLTYPE AddRef()
     {
-        MsRdpEx_LogPrint(DEBUG, "CClassFactory::AddRef");
-        return ++m_refCount;
+        ULONG refCount = InterlockedIncrement(&m_refCount);
+        MsRdpEx_LogPrint(DEBUG, "CClassFactory::AddRef = %d", refCount);
+        return refCount;
     }
 
     ULONG STDMETHODCALLTYPE Release()
     {
-        MsRdpEx_LogPrint(DEBUG, "CClassFactory::Release");
-        if (--m_refCount == 0)
+        ULONG refCount = InterlockedDecrement(&m_refCount);
+        MsRdpEx_LogPrint(DEBUG, "CClassFactory::Release = %d", refCount);
+        if (refCount == 0)
         {
             delete this;
             return 0;
         }
-        return m_refCount;
+        return refCount;
     }
 
     // IClassFactory interface
@@ -824,6 +823,7 @@ class CClassFactory : IClassFactory
             {
                 CMsRdpClient* pMsRdpClient = new CMsRdpClient((IUnknown*)*ppvObject);
                 hr = pMsRdpClient->QueryInterface(riid, ppvObject);
+                pMsRdpClient->Release();
             }
         }
 

dll/RdpCoreApi.cpp

@@ -15,7 +15,7 @@ class CMsRdpExCoreApi : public IMsRdpExCoreApi
 public:
     CMsRdpExCoreApi()
     {
-        m_refCount = 0;
+        m_refCount = 1;
         MsRdpEx_PathCchDetect(m_MsRdpExDllPath, MSRDPEX_MAX_PATH, MSRDPEX_CURRENT_LIBRARY_PATH);
     }
 
@@ -127,22 +127,21 @@ class CMsRdpExCoreApi : public IMsRdpExCoreApi
 
     bool __stdcall QueryInstanceByWindowHandle(HWND hWnd, LPVOID* ppvObject)
     {
-        HRESULT hr;
         IMsRdpExInstance* instance = NULL;
 
         instance = (IMsRdpExInstance*) MsRdpEx_InstanceManager_FindByOutputPresenterHwnd(hWnd);
 
         if (!instance)
             return false;
 
-        hr = instance->QueryInterface(IID_IMsRdpExInstance, ppvObject);
+        instance->AddRef();
+        *ppvObject = (LPVOID)instance;
 
-        return (hr == S_OK) ? true : false;
+        return true;
     }
 
     bool __stdcall OpenInstanceForWindowHandle(HWND hWnd, LPVOID* ppvObject)
     {
-        HRESULT hr;
         IMsRdpExInstance* instance = NULL;
 
         instance = (IMsRdpExInstance*) MsRdpEx_InstanceManager_FindByOutputPresenterHwnd(hWnd);
@@ -153,10 +152,14 @@ class CMsRdpExCoreApi : public IMsRdpExCoreApi
             instance->AttachOutputWindow(hWnd, NULL);
             MsRdpEx_InstanceManager_Add((CMsRdpExInstance*) instance);
         }
+        else
+        {
+            instance->AddRef();
+        }
 
-        hr = instance->QueryInterface(IID_IMsRdpExInstance, ppvObject);
+        *ppvObject = (LPVOID)instance;
 
-        return (hr == S_OK) ? true : false;
+        return true;
     }
 
 private:

dll/RdpDvcClient.cpp

@@ -93,7 +93,7 @@ void CRdpDvcClient::SetListener(IWTSListener* pListener)
 
 CRdpDvcClient::CRdpDvcClient(void)
 {
-
+    m_refCount = 1;
 }
 
 CRdpDvcClient::~CRdpDvcClient()
@@ -184,7 +184,7 @@ HRESULT STDMETHODCALLTYPE CRdpDvcListener::OnNewChannelConnection(IWTSVirtualCha
 
 CRdpDvcListener::CRdpDvcListener(void)
 {
-
+    m_refCount = 1;
 }
 
 CRdpDvcListener::~CRdpDvcListener()
@@ -288,7 +288,7 @@ HRESULT STDMETHODCALLTYPE CRdpDvcPlugin::Terminated(void)
 
 CRdpDvcPlugin::CRdpDvcPlugin(void)
 {
-
+    m_refCount = 1;
 }
 
 CRdpDvcPlugin::~CRdpDvcPlugin()
@@ -324,12 +324,12 @@ class CDvcPluginClassFactory : IClassFactory
 
         if (riid == IID_IUnknown) {
             *ppvObject = (LPVOID)((IUnknown*)this);
-            m_refCount++;
+            InterlockedIncrement(&m_refCount);
             return S_OK;
         }
         if (riid == IID_IClassFactory) {
             *ppvObject = (LPVOID)((IClassFactory*)this);
-            m_refCount++;
+            InterlockedIncrement(&m_refCount);
             return S_OK;
         }
 

dll/RdpDvcClient.h

@@ -26,7 +26,7 @@ class CRdpDvcClient :
     virtual ~CRdpDvcClient();
 
 private:
-    ULONG m_refCount = 0;
+    ULONG m_refCount;
     IWTSVirtualChannel* m_pChannel = NULL;
     IWTSListener* m_pListener = NULL;
 };
@@ -48,7 +48,7 @@ class CRdpDvcListener :
     CRdpDvcListener(void);
     virtual ~CRdpDvcListener();
 private:
-    ULONG m_refCount = 0;
+    ULONG m_refCount;
 };
 
 class CRdpDvcPlugin :
@@ -71,7 +71,7 @@ class CRdpDvcPlugin :
     virtual ~CRdpDvcPlugin();
 
 private:
-    ULONG m_refCount = 0;
+    ULONG m_refCount;
     IWTSVirtualChannel* m_pChannel = NULL;
 };
 

dll/RdpInstance.cpp

@@ -15,7 +15,7 @@ class CMsRdpExInstance : public IMsRdpExInstance
 public:
     CMsRdpExInstance(CMsRdpClient* pMsRdpClient)
     {
-        m_refCount = 0;
+        m_refCount = 1;
         m_pMsRdpClient = pMsRdpClient;
         MsRdpEx_GuidGenerate(&m_sessionId);
 
@@ -48,6 +48,12 @@ class CMsRdpExInstance : public IMsRdpExInstance
         char iid[MSRDPEX_GUID_STRING_SIZE];
         MsRdpEx_GuidBinToStr((GUID*)&riid, iid, 0);
 
+        if (!ppvObject) {
+            return E_INVALIDARG;
+        }
+
+        *ppvObject = NULL;  // Initialize to NULL for safety
+
         if (riid == IID_IUnknown)
         {
             *ppvObject = (LPVOID)((IUnknown*)this);
@@ -56,7 +62,7 @@ class CMsRdpExInstance : public IMsRdpExInstance
         }
         else if (riid == IID_IMsRdpExInstance)
         {
-            *ppvObject = (LPVOID)((IUnknown*)this);
+            *ppvObject = (LPVOID)((IMsRdpExInstance*)this);
             refCount = InterlockedIncrement(&m_refCount);
             hr = S_OK;
         }
@@ -300,7 +306,7 @@ class CMsRdpExInstance : public IMsRdpExInstance
 
 public:
     GUID m_sessionId;
-    ULONG m_refCount = NULL;
+    ULONG m_refCount;
     CMsRdpClient* m_pMsRdpClient = NULL;
     HWND m_hInputCaptureWnd = NULL;
     HWND m_hOutputPresenterWnd = NULL;

dll/RdpProcess.cpp

@@ -86,7 +86,7 @@ class CMsRdpExProcess : public IMsRdpExProcess
 public:
     CMsRdpExProcess()
     {
-        m_refCount = 0;
+        m_refCount = 1;
         m_exitCode = 0;
         m_hasExited = false;
         m_filename = NULL;