mshtml: Store client site in HTMLDocumentObj as soon as possible in SetClientSite.

dlls/mshtml/oleobj.c

@@ -96,8 +96,10 @@ static HRESULT WINAPI OleObject_SetClientSite(IOleObject *iface, IOleClientSite
     HTMLDocument *This = OLEOBJ_THIS(iface);
     IDocHostUIHandler *pDocHostUIHandler = NULL;
     IOleCommandTarget *cmdtrg = NULL;
+    IOleWindow *ole_window;
     BOOL hostui_setup;
     VARIANT silent;
+    HWND hwnd;
     HRESULT hres;
 
     TRACE("(%p)->(%p)\n", This, pClientSite);
@@ -121,6 +123,9 @@ static HRESULT WINAPI OleObject_SetClientSite(IOleObject *iface, IOleClientSite
     if(!pClientSite)
         return S_OK;
 
+    IOleClientSite_AddRef(pClientSite);
+    This->doc_obj->client = pClientSite;
+
     hostui_setup = This->doc_obj->hostui_setup;
 
     hres = IOleObject_QueryInterface(pClientSite, &IID_IDocHostUIHandler, (void**)&pDocHostUIHandler);
@@ -129,6 +134,8 @@ static HRESULT WINAPI OleObject_SetClientSite(IOleObject *iface, IOleClientSite
         LPOLESTR key_path = NULL, override_key_path = NULL;
         IDocHostUIHandler2 *pDocHostUIHandler2;
 
+        This->doc_obj->hostui = pDocHostUIHandler;
+
         memset(&hostinfo, 0, sizeof(DOCHOSTUIINFO));
         hostinfo.cbSize = sizeof(DOCHOSTUIINFO);
         hres = IDocHostUIHandler_GetHostInfo(pDocHostUIHandler, &hostinfo);
@@ -166,19 +173,16 @@ static HRESULT WINAPI OleObject_SetClientSite(IOleObject *iface, IOleClientSite
 
             This->doc_obj->hostui_setup = TRUE;
         }
+    }else {
+        This->doc_obj->hostui = NULL;
     }
 
     /* Native calls here GetWindow. What is it for?
      * We don't have anything to do with it here (yet). */
-    if(pClientSite) {
-        IOleWindow *pOleWindow = NULL;
-        HWND hwnd;
-
-        hres = IOleClientSite_QueryInterface(pClientSite, &IID_IOleWindow, (void**)&pOleWindow);
-        if(SUCCEEDED(hres)) {
-            IOleWindow_GetWindow(pOleWindow, &hwnd);
-            IOleWindow_Release(pOleWindow);
-        }
+    hres = IOleClientSite_QueryInterface(pClientSite, &IID_IOleWindow, (void**)&ole_window);
+    if(SUCCEEDED(hres)) {
+        IOleWindow_GetWindow(ole_window, &hwnd);
+        IOleWindow_Release(ole_window);
     }
 
     hres = IOleClientSite_QueryInterface(pClientSite, &IID_IOleCommandTarget, (void**)&cmdtrg);
@@ -204,10 +208,6 @@ static HRESULT WINAPI OleObject_SetClientSite(IOleObject *iface, IOleClientSite
         IOleCommandTarget_Release(cmdtrg);
     }
 
-    IOleClientSite_AddRef(pClientSite);
-    This->doc_obj->client = pClientSite;
-    This->doc_obj->hostui = pDocHostUIHandler;
-
     if(This->doc_obj->usermode == UNKNOWN_USERMODE)
         IOleControl_OnAmbientPropertyChange(CONTROL(This), DISPID_AMBIENT_USERMODE);
 

dlls/mshtml/tests/htmldoc.c

@@ -1368,6 +1368,74 @@ static const IOleContainerVtbl OleContainerVtbl = {
 
 static IOleContainer OleContainer = { &OleContainerVtbl };
 
+static HRESULT WINAPI ClientSite_QueryInterface(IOleClientSite *iface, REFIID riid, void **ppv)
+{
+    return QueryInterface(riid, ppv);
+}
+
+static ULONG WINAPI ClientSite_AddRef(IOleClientSite *iface)
+{
+    return 2;
+}
+
+static ULONG WINAPI ClientSite_Release(IOleClientSite *iface)
+{
+    return 1;
+}
+
+static HRESULT WINAPI ClientSite_SaveObject(IOleClientSite *iface)
+{
+    ok(0, "unexpected call\n");
+    return E_NOTIMPL;
+}
+
+static HRESULT WINAPI ClientSite_GetMoniker(IOleClientSite *iface, DWORD dwAssign, DWORD dwWhichMoniker,
+        IMoniker **ppmon)
+{
+    ok(0, "unexpected call\n");
+    return E_NOTIMPL;
+}
+
+static HRESULT WINAPI ClientSite_GetContainer(IOleClientSite *iface, IOleContainer **ppContainer)
+{
+    CHECK_EXPECT(GetContainer);
+    ok(ppContainer != NULL, "ppContainer = NULL\n");
+    *ppContainer = &OleContainer;
+    return S_OK;
+}
+
+static HRESULT WINAPI ClientSite_ShowObject(IOleClientSite *iface)
+{
+    ok(0, "unexpected call\n");
+    return E_NOTIMPL;
+}
+
+static HRESULT WINAPI ClientSite_OnShowWindow(IOleClientSite *iface, BOOL fShow)
+{
+    ok(0, "unexpected call\n");
+    return E_NOTIMPL;
+}
+
+static HRESULT WINAPI ClientSite_RequestNewObjectLayout(IOleClientSite *iface)
+{
+    ok(0, "unexpected call\n");
+    return E_NOTIMPL;
+}
+
+static const IOleClientSiteVtbl ClientSiteVtbl = {
+    ClientSite_QueryInterface,
+    ClientSite_AddRef,
+    ClientSite_Release,
+    ClientSite_SaveObject,
+    ClientSite_GetMoniker,
+    ClientSite_GetContainer,
+    ClientSite_ShowObject,
+    ClientSite_OnShowWindow,
+    ClientSite_RequestNewObjectLayout
+};
+
+static IOleClientSite ClientSite = { &ClientSiteVtbl };
+
 static HRESULT WINAPI InPlaceFrame_QueryInterface(IOleInPlaceFrame *iface, REFIID riid, void **ppv)
 {
     static const GUID undocumented_frame_iid = {0xfbece6c9,0x48d7,0x4a37,{0x8f,0xe3,0x6a,0xd4,0x27,0x2f,0xdd,0xac}};
@@ -1558,9 +1626,23 @@ static ULONG WINAPI InPlaceSiteWindowless_Release(IOleInPlaceSiteWindowless *ifa
 static HRESULT WINAPI InPlaceSiteWindowless_GetWindow(
         IOleInPlaceSiteWindowless *iface, HWND *phwnd)
 {
+    IOleClientSite *client_site;
+    IOleObject *ole_obj;
+    HRESULT hres;
+
     CHECK_EXPECT2(GetWindow);
     ok(phwnd != NULL, "phwnd = NULL\n");
     *phwnd = container_hwnd;
+
+    hres = IUnknown_QueryInterface(doc_unk, &IID_IOleObject, (void**)&ole_obj);
+    ok(hres == S_OK, "Could not get IOleObject: %08x\n", hres);
+
+    hres = IOleObject_GetClientSite(ole_obj, &client_site);
+    IOleObject_Release(ole_obj);
+    ok(hres == S_OK, "GetClientSite failed: %08x\n", hres);
+    ok(client_site == &ClientSite, "client_site != ClientSite\n");
+    IOleClientSite_Release(client_site);
+
     return S_OK;
 }
 
@@ -1821,74 +1903,6 @@ static const IOleInPlaceSiteWindowlessVtbl InPlaceSiteWindowlessVtbl = {
 
 static IOleInPlaceSiteWindowless InPlaceSiteWindowless = { &InPlaceSiteWindowlessVtbl };
 
-static HRESULT WINAPI ClientSite_QueryInterface(IOleClientSite *iface, REFIID riid, void **ppv)
-{
-    return QueryInterface(riid, ppv);
-}
-
-static ULONG WINAPI ClientSite_AddRef(IOleClientSite *iface)
-{
-    return 2;
-}
-
-static ULONG WINAPI ClientSite_Release(IOleClientSite *iface)
-{
-    return 1;
-}
-
-static HRESULT WINAPI ClientSite_SaveObject(IOleClientSite *iface)
-{
-    ok(0, "unexpected call\n");
-    return E_NOTIMPL;
-}
-
-static HRESULT WINAPI ClientSite_GetMoniker(IOleClientSite *iface, DWORD dwAssign, DWORD dwWhichMoniker,
-        IMoniker **ppmon)
-{
-    ok(0, "unexpected call\n");
-    return E_NOTIMPL;
-}
-
-static HRESULT WINAPI ClientSite_GetContainer(IOleClientSite *iface, IOleContainer **ppContainer)
-{
-    CHECK_EXPECT(GetContainer);
-    ok(ppContainer != NULL, "ppContainer = NULL\n");
-    *ppContainer = &OleContainer;
-    return S_OK;
-}
-
-static HRESULT WINAPI ClientSite_ShowObject(IOleClientSite *iface)
-{
-    ok(0, "unexpected call\n");
-    return E_NOTIMPL;
-}
-
-static HRESULT WINAPI ClientSite_OnShowWindow(IOleClientSite *iface, BOOL fShow)
-{
-    ok(0, "unexpected call\n");
-    return E_NOTIMPL;
-}
-
-static HRESULT WINAPI ClientSite_RequestNewObjectLayout(IOleClientSite *iface)
-{
-    ok(0, "unexpected call\n");
-    return E_NOTIMPL;
-}
-
-static const IOleClientSiteVtbl ClientSiteVtbl = {
-    ClientSite_QueryInterface,
-    ClientSite_AddRef,
-    ClientSite_Release,
-    ClientSite_SaveObject,
-    ClientSite_GetMoniker,
-    ClientSite_GetContainer,
-    ClientSite_ShowObject,
-    ClientSite_OnShowWindow,
-    ClientSite_RequestNewObjectLayout
-};
-
-static IOleClientSite ClientSite = { &ClientSiteVtbl };
-
 static HRESULT WINAPI DocumentSite_QueryInterface(IOleDocumentSite *iface, REFIID riid, void **ppv)
 {
     return QueryInterface(riid, ppv);