Bug 773962 - Fix up waivers after transplanting. r=mrbkap

dom/base/nsGlobalWindow.cpp

@@ -1976,7 +1976,7 @@ nsGlobalWindow::SetNewDocument(nsIDocument* aDocument,
 
       js::SetProxyExtra(mJSObject, 0, js::PrivateValue(NULL));
 
-      outerObject = JS_TransplantObject(cx, mJSObject, outerObject);
+      outerObject = xpc::TransplantObject(cx, mJSObject, outerObject);
       if (!outerObject) {
         NS_ERROR("unable to transplant wrappers, probably OOM");
         return NS_ERROR_FAILURE;

js/xpconnect/src/XPCWrappedNative.cpp

@@ -1631,15 +1631,15 @@ XPCWrappedNative::ReparentWrapperIfFound(XPCCallContext& ccx,
                 }
 
                 // Ok, now we do the special object-plus-wrapper transplant.
-                ww = js_TransplantObjectWithWrapper(ccx, flat, ww, newobj,
-                                                    newwrapper);
+                ww = xpc::TransplantObjectWithWrapper(ccx, flat, ww, newobj,
+                                                      newwrapper);
                 if (!ww)
                     return NS_ERROR_FAILURE;
 
                 flat = newobj;
                 wrapper->SetWrapper(ww);
             } else {
-                flat = JS_TransplantObject(ccx, flat, newobj);
+                flat = xpc::TransplantObject(ccx, flat, newobj);
                 if (!flat)
                     return NS_ERROR_FAILURE;
             }

js/xpconnect/src/xpcpublic.h

@@ -32,6 +32,16 @@ class nsScriptNameSpaceManager;
 #define BAD_TLS_INDEX ((PRUint32) -1)
 #endif
 
+namespace xpc {
+JSObject *
+TransplantObject(JSContext *cx, JSObject *origobj, JSObject *target);
+
+JSObject *
+TransplantObjectWithWrapper(JSContext *cx,
+                            JSObject *origobj, JSObject *origwrapper,
+                            JSObject *targetobj, JSObject *targetwrapper);
+} /* namespace xpc */
+
 nsresult
 xpc_CreateGlobalObject(JSContext *cx, JSClass *clasp,
                        nsIPrincipal *principal, nsISupports *ptr,

js/xpconnect/wrappers/WrapperFactory.cpp

@@ -590,4 +590,70 @@ WrapperFactory::WrapForSameCompartmentXray(JSContext *cx, JSObject *obj)
     return wrapperObj;
 }
 
+/*
+ * Calls to JS_TransplantObject* should go through these helpers here so that
+ * waivers get fixed up properly.
+ */
+
+static bool
+FixWaiverAfterTransplant(JSContext *cx, JSObject *oldWaiver, JSObject *newobj)
+{
+    MOZ_ASSERT(Wrapper::wrapperHandler(oldWaiver) == &XrayWaiver);
+    MOZ_ASSERT(!js::IsCrossCompartmentWrapper(newobj));
+
+    // Create a waiver in the new compartment. We know there's not one already
+    // because we _just_ transplanted, which means that |newobj| was either
+    // created from scratch, or was previously cross-compartment wrapper (which
+    // should have no waiver). CreateXrayWaiver asserts this.
+    JSObject *newWaiver = WrapperFactory::CreateXrayWaiver(cx, newobj);
+    if (!newWaiver)
+        return false;
+
+    // Update all the cross-compartment references to oldWaiver to point to
+    // newWaiver.
+    if (!js::RemapAllWrappersForObject(cx, oldWaiver, newWaiver))
+        return false;
+
+    // There should be no same-compartment references to oldWaiver, and we
+    // just remapped all cross-compartment references. It's dead, so we can
+    // remove it from the map.
+    CompartmentPrivate *priv = GetCompartmentPrivate(oldWaiver);
+    JSObject *key = Wrapper::wrappedObject(oldWaiver);
+    MOZ_ASSERT(priv->waiverWrapperMap->Find(key));
+    priv->waiverWrapperMap->Remove(key);
+    return true;
+}
+
+JSObject *
+TransplantObject(JSContext *cx, JSObject *origobj, JSObject *target)
+{
+    JSObject *oldWaiver = WrapperFactory::GetXrayWaiver(origobj);
+    JSObject *newIdentity = JS_TransplantObject(cx, origobj, target);
+    if (!newIdentity || !oldWaiver)
+       return newIdentity;
+
+    if (!FixWaiverAfterTransplant(cx, oldWaiver, newIdentity))
+        return NULL;
+    return newIdentity;
+}
+
+JSObject *
+TransplantObjectWithWrapper(JSContext *cx,
+                            JSObject *origobj, JSObject *origwrapper,
+                            JSObject *targetobj, JSObject *targetwrapper)
+{
+    JSObject *oldWaiver = WrapperFactory::GetXrayWaiver(origobj);
+    JSObject *newSameCompartmentWrapper =
+      js_TransplantObjectWithWrapper(cx, origobj, origwrapper, targetobj,
+                                     targetwrapper);
+    if (!newSameCompartmentWrapper || !oldWaiver)
+        return newSameCompartmentWrapper;
+
+    JSObject *newIdentity = Wrapper::wrappedObject(newSameCompartmentWrapper);
+    JS_ASSERT(js::IsWrapper(newIdentity));
+    if (!FixWaiverAfterTransplant(cx, oldWaiver, newIdentity))
+        return NULL;
+    return newSameCompartmentWrapper;
+}
+
 }