From 2f93763e2407689cc905140dcba9c627432db077 Mon Sep 17 00:00:00 2001
From: Jacek Caban <jack@itma.pwr.wroc.pl>
Date: Tue, 2 Aug 2005 09:49:06 +0000
Subject: [PATCH] Added DllCanUnloadNow implementation.

---
 dlls/mshtml/htmldoc.c        |  8 +++++++-
 dlls/mshtml/main.c           | 21 ++++++++++++++++-----
 dlls/mshtml/mshtml_private.h |  4 ++++
 dlls/mshtml/protocol.c       | 32 +++++++++++++++++++++++---------
 4 files changed, 50 insertions(+), 15 deletions(-)

diff --git a/dlls/mshtml/htmldoc.c b/dlls/mshtml/htmldoc.c
index f27be357aea..142751bef12 100644
--- a/dlls/mshtml/htmldoc.c
+++ b/dlls/mshtml/htmldoc.c
@@ -140,6 +140,8 @@ static ULONG WINAPI HTMLDocument_Release(IHTMLDocument2 *iface)
         if(This->nscontainer)
             HTMLDocument_NSContainer_Destroy(This);
         HeapFree(GetProcessHeap(), 0, This);
+
+        UNLOCK_MODULE();
     }
 
     return ref;
@@ -974,8 +976,12 @@ HRESULT HTMLDocument_Create(IUnknown *pUnkOuter, REFIID riid, void** ppvObject)
     ret->ref = 0;
 
     hres = IHTMLDocument_QueryInterface(HTMLDOC(ret), riid, ppvObject);
-    if(FAILED(hres))
+    if(FAILED(hres)) {
         HeapFree(GetProcessHeap(), 0, ret);
+        return hres;
+    }
+
+    LOCK_MODULE();
 
     HTMLDocument_Persist_Init(ret);
     HTMLDocument_OleObj_Init(ret);
diff --git a/dlls/mshtml/main.c b/dlls/mshtml/main.c
index 88e1aa41106..052f70355d0 100644
--- a/dlls/mshtml/main.c
+++ b/dlls/mshtml/main.c
@@ -48,6 +48,7 @@
 WINE_DEFAULT_DEBUG_CHANNEL(mshtml);
 
 HINSTANCE hInst;
+LONG module_ref = 0;
 
 BOOL WINAPI DllMain(HINSTANCE hInstDLL, DWORD fdwReason, LPVOID lpv)
 {
@@ -100,8 +101,10 @@ static ULONG WINAPI ClassFactory_Release(IClassFactory *iface)
 
     TRACE("(%p) ref = %lu\n", This, ref);
 
-    if(!ref)
+    if(!ref) {
         HeapFree(GetProcessHeap(), 0, This);
+        UNLOCK_MODULE();
+    }
 
     return ref;
 }
@@ -115,7 +118,13 @@ static HRESULT WINAPI ClassFactory_CreateInstance(IClassFactory *iface, IUnknown
 
 static HRESULT WINAPI ClassFactory_LockServer(IClassFactory *iface, BOOL dolock)
 {
-    FIXME("(%p)->(%x) stub\n", iface, dolock);
+    TRACE("(%p)->(%x)\n", iface, dolock);
+
+    if(dolock)
+        LOCK_MODULE();
+    else
+        UNLOCK_MODULE();
+
     return S_OK;
 }
 
@@ -137,7 +146,9 @@ static HRESULT ClassFactory_Create(REFIID riid, void **ppv, CreateInstanceFunc f
     ret->fnCreateInstance = fnCreateInstance;
 
     hres = IClassFactory_QueryInterface((IClassFactory*)ret, riid, ppv);
-    if(FAILED(hres)) {
+    if(SUCCEEDED(hres)) {
+        LOCK_MODULE();
+    }else {
         HeapFree(GetProcessHeap(), 0, ret);
         *ppv = NULL;
     }
@@ -172,8 +183,8 @@ HRESULT WINAPI MSHTML_DllGetClassObject(REFCLSID rclsid, REFIID riid, LPVOID *pp
 
 HRESULT WINAPI MSHTML_DllCanUnloadNow(void)
 {
-    FIXME("()\n");
-    return S_FALSE;
+    TRACE("() ref=%ld\n", module_ref);
+    return module_ref ? S_FALSE : S_OK;
 }
 
 /***********************************************************************
diff --git a/dlls/mshtml/mshtml_private.h b/dlls/mshtml/mshtml_private.h
index eff57a28b34..d79ea4753e4 100644
--- a/dlls/mshtml/mshtml_private.h
+++ b/dlls/mshtml/mshtml_private.h
@@ -108,4 +108,8 @@ DEFINE_GUID(CLSID_MailtoProtocol, 0x3050F3DA, 0x98B5, 0x11CF, 0xBB,0x82, 0x00,0x
 DEFINE_GUID(CLSID_ResProtocol, 0x3050F3BC, 0x98B5, 0x11CF, 0xBB,0x82, 0x00,0xAA,0x00,0xBD,0xCE,0x0B);
 DEFINE_GUID(CLSID_SysimageProtocol, 0x76E67A63, 0x06E9, 0x11D2, 0xA8,0x40, 0x00,0x60,0x08,0x05,0x93,0x82);
 
+extern LONG module_ref;
+#define LOCK_MODULE()   InterlockedIncrement(&module_ref)
+#define UNLOCK_MODULE() InterlockedDecrement(&module_ref)
+
 extern HINSTANCE hInst;
diff --git a/dlls/mshtml/protocol.c b/dlls/mshtml/protocol.c
index 2a0312ea407..3434f9262f1 100644
--- a/dlls/mshtml/protocol.c
+++ b/dlls/mshtml/protocol.c
@@ -51,8 +51,7 @@ typedef struct {
     const IClassFactoryVtbl         *lpClassFactoryVtbl;
 } ProtocolFactory;
 
-#define PROTOCOLINFO_THIS(iface) \
-    (ProtocolFactory*)((char*)(iface)-offsetof(ProtocolFactory,lpInternetProtocolInfoVtbl))
+#define PROTOCOLINFO_THIS(iface) DEFINE_THIS(ProtocolFactory, InternetProtocolInfo, iface)
 
 static HRESULT WINAPI InternetProtocolInfo_QueryInterface(IInternetProtocolInfo *iface, REFIID riid, void **ppv)
 {
@@ -83,6 +82,7 @@ static ULONG WINAPI InternetProtocolInfo_AddRef(IInternetProtocolInfo *iface)
 {
     ProtocolFactory *This = PROTOCOLINFO_THIS(iface);
     TRACE("(%p)\n", This);
+    LOCK_MODULE();
     return 2;
 }
 
@@ -90,13 +90,13 @@ static ULONG WINAPI InternetProtocolInfo_Release(IInternetProtocolInfo *iface)
 {
     ProtocolFactory *This = PROTOCOLINFO_THIS(iface);
     TRACE("(%p)\n", This);
+    UNLOCK_MODULE();
     return 1;
 }
 
 #undef PROTOCOLINFO_THIS
 
-#define CLASSFACTORY_THIS(iface) \
-    (ProtocolFactory*)((char*)(iface)-offsetof(ProtocolFactory,lpClassFactoryVtbl))
+#define CLASSFACTORY_THIS(iface) DEFINE_THIS(ProtocolFactory, ClassFactory, iface)
 
 static HRESULT WINAPI ClassFactory_QueryInterface(IClassFactory *iface, REFIID riid, void **ppv)
 {
@@ -119,7 +119,14 @@ static ULONG WINAPI ClassFactory_Release(IClassFactory *iface)
 static HRESULT WINAPI ClassFactory_LockServer(IClassFactory *iface, BOOL dolock)
 {
     ProtocolFactory *This = CLASSFACTORY_THIS(iface);
-    FIXME("(%p)->(%x)\n", This, dolock);
+
+    TRACE("(%p)->(%x)\n", This, dolock);
+
+    if(dolock)
+        LOCK_MODULE();
+    else
+        UNLOCK_MODULE();
+
     return S_OK;
 }
 
@@ -176,8 +183,10 @@ static ULONG WINAPI AboutProtocol_Release(IInternetProtocol *iface)
 
     TRACE("(%p) ref=%lx\n", iface, ref);
 
-    if(!ref)
+    if(!ref) {
         HeapFree(GetProcessHeap(), 0, This);
+        UNLOCK_MODULE();
+    }
 
     return ref;
 }
@@ -287,7 +296,9 @@ static HRESULT WINAPI AboutProtocolFactory_CreateInstance(IClassFactory *iface,
 
     hres = IUnknown_QueryInterface((IUnknown*)ret, riid, ppv);
 
-    if(FAILED(hres))
+    if(SUCCEEDED(hres))
+        LOCK_MODULE();
+    else
         HeapFree(GetProcessHeap(), 0, ret);
 
     return hres;
@@ -408,6 +419,7 @@ static ULONG WINAPI ResProtocol_Release(IInternetProtocol *iface)
     if(!ref) {
         HeapFree(GetProcessHeap(), 0, This->data);
         HeapFree(GetProcessHeap(), 0, This);
+        UNLOCK_MODULE();
     }
 
     return ref;
@@ -612,7 +624,9 @@ static HRESULT WINAPI ResProtocolFactory_CreateInstance(IClassFactory *iface, IU
 
     hres = IUnknown_QueryInterface((IUnknown*)ret, riid, ppv);
 
-    if(FAILED(hres))
+    if(SUCCEEDED(hres))
+        LOCK_MODULE();
+    else
         HeapFree(GetProcessHeap(), 0, ret);
 
     return hres;
@@ -688,6 +702,6 @@ HRESULT ProtocolFactory_Create(REFCLSID rclsid, REFIID riid, void **ppv)
         FIXME("not implemented protocol %s\n", debugstr_guid(rclsid));
         return CLASS_E_CLASSNOTAVAILABLE;
     }
-
+ 
     return IUnknown_QueryInterface((IUnknown*)cf, riid, ppv);
 }