1 files changed, 55 insertions, 24 deletions

diff --git a/components/script/dom/messageevent.rs b/components/script/dom/messageevent.rs
index cb8ebaaebd1..3e9a7079605 100644
--- a/components/script/dom/messageevent.rs
+++ b/components/script/dom/messageevent.rs
@@ -5,9 +5,10 @@
 use crate::dom::bindings::codegen::Bindings::EventBinding::EventMethods;
 use crate::dom::bindings::codegen::Bindings::MessageEventBinding;
 use crate::dom::bindings::codegen::Bindings::MessageEventBinding::MessageEventMethods;
+use crate::dom::bindings::codegen::UnionTypes::WindowProxyOrMessagePortOrServiceWorker;
 use crate::dom::bindings::error::Fallible;
 use crate::dom::bindings::inheritance::Castable;
-use crate::dom::bindings::reflector::{reflect_dom_object, DomObject};
+use crate::dom::bindings::reflector::reflect_dom_object;
 use crate::dom::bindings::root::{Dom, DomRoot};
 use crate::dom::bindings::str::DOMString;
 use crate::dom::bindings::trace::RootedTraceableBox;
@@ -16,14 +17,39 @@ use crate::dom::event::Event;
 use crate::dom::eventtarget::EventTarget;
 use crate::dom::globalscope::GlobalScope;
 use crate::dom::messageport::MessagePort;
+use crate::dom::serviceworker::ServiceWorker;
 use crate::dom::windowproxy::WindowProxy;
 use crate::script_runtime::JSContext;
 use dom_struct::dom_struct;
-use js::jsapi::{Heap, JSObject};
+use js::jsapi::Heap;
 use js::jsval::JSVal;
 use js::rust::HandleValue;
 use servo_atoms::Atom;
-use std::ptr::NonNull;
+
+#[must_root]
+#[derive(JSTraceable, MallocSizeOf)]
+enum SrcObject {
+    WindowProxy(Dom<WindowProxy>),
+    MessagePort(Dom<MessagePort>),
+    ServiceWorker(Dom<ServiceWorker>),
+}
+
+impl From<&WindowProxyOrMessagePortOrServiceWorker> for SrcObject {
+    #[allow(unrooted_must_root)]
+    fn from(src_object: &WindowProxyOrMessagePortOrServiceWorker) -> SrcObject {
+        match src_object {
+            WindowProxyOrMessagePortOrServiceWorker::WindowProxy(blob) => {
+                SrcObject::WindowProxy(Dom::from_ref(&*blob))
+            },
+            WindowProxyOrMessagePortOrServiceWorker::MessagePort(stream) => {
+                SrcObject::MessagePort(Dom::from_ref(&*stream))
+            },
+            WindowProxyOrMessagePortOrServiceWorker::ServiceWorker(stream) => {
+                SrcObject::ServiceWorker(Dom::from_ref(&*stream))
+            },
+        }
+    }
+}
 
 #[dom_struct]
 pub struct MessageEvent {
@@ -31,7 +57,7 @@ pub struct MessageEvent {
     #[ignore_malloc_size_of = "mozjs"]
     data: Heap<JSVal>,
     origin: DOMString,
-    source: Option<Dom<WindowProxy>>,
+    source: Option<SrcObject>,
     lastEventId: DOMString,
     ports: Vec<DomRoot<MessagePort>>,
 }
@@ -52,14 +78,14 @@ impl MessageEvent {
         global: &GlobalScope,
         data: HandleValue,
         origin: DOMString,
-        source: Option<&WindowProxy>,
+        source: Option<&WindowProxyOrMessagePortOrServiceWorker>,
         lastEventId: DOMString,
         ports: Vec<DomRoot<MessagePort>>,
     ) -> DomRoot<MessageEvent> {
         let ev = Box::new(MessageEvent {
             event: Event::new_inherited(),
             data: Heap::default(),
-            source: source.map(Dom::from_ref),
+            source: source.map(|source| source.into()),
             origin,
             lastEventId,
             ports,
@@ -77,7 +103,7 @@ impl MessageEvent {
         cancelable: bool,
         data: HandleValue,
         origin: DOMString,
-        source: Option<&WindowProxy>,
+        source: Option<&WindowProxyOrMessagePortOrServiceWorker>,
         lastEventId: DOMString,
         ports: Vec<DomRoot<MessagePort>>,
     ) -> DomRoot<MessageEvent> {
@@ -94,10 +120,6 @@ impl MessageEvent {
         type_: DOMString,
         init: RootedTraceableBox<MessageEventBinding::MessageEventInit>,
     ) -> Fallible<DomRoot<MessageEvent>> {
-        let source = init
-            .source
-            .as_ref()
-            .and_then(|inner| inner.as_ref().map(|source| source.window_proxy()));
         let ev = MessageEvent::new(
             global,
             Atom::from(type_),
@@ -105,9 +127,9 @@ impl MessageEvent {
             init.parent.cancelable,
             init.data.handle(),
             init.origin.clone(),
-            source.as_ref().map(|source| &**source),
+            init.source.as_ref(),
             init.lastEventId.clone(),
-            init.ports.clone().unwrap_or(vec![]),
+            init.ports.clone(),
         );
         Ok(ev)
     }
@@ -129,7 +151,11 @@ impl MessageEvent {
             false,
             message,
             DOMString::from(origin.unwrap_or("")),
-            source,
+            source
+                .map(|source| {
+                    WindowProxyOrMessagePortOrServiceWorker::WindowProxy(DomRoot::from_ref(source))
+                })
+                .as_ref(),
             DOMString::new(),
             ports,
         );
@@ -138,10 +164,6 @@ impl MessageEvent {
 
     pub fn dispatch_error(target: &EventTarget, scope: &GlobalScope) {
         let init = MessageEventBinding::MessageEventInit::empty();
-        let source = init
-            .source
-            .as_ref()
-            .and_then(|inner| inner.as_ref().map(|source| source.window_proxy()));
         let messageevent = MessageEvent::new(
             scope,
             atom!("messageerror"),
@@ -149,9 +171,9 @@ impl MessageEvent {
             init.parent.cancelable,
             init.data.handle(),
             init.origin.clone(),
-            source.as_ref().map(|source| &**source),
+            init.source.as_ref(),
             init.lastEventId.clone(),
-            init.ports.clone().unwrap_or(vec![]),
+            init.ports.clone(),
         );
         messageevent.upcast::<Event>().fire(target);
     }
@@ -169,10 +191,19 @@ impl MessageEventMethods for MessageEvent {
     }
 
     // https://html.spec.whatwg.org/multipage/#dom-messageevent-source
-    fn GetSource(&self, _cx: JSContext) -> Option<NonNull<JSObject>> {
-        self.source
-            .as_ref()
-            .and_then(|source| NonNull::new(source.reflector().get_jsobject().get()))
+    fn GetSource(&self) -> Option<WindowProxyOrMessagePortOrServiceWorker> {
+        match &self.source {
+            Some(SrcObject::WindowProxy(i)) => Some(
+                WindowProxyOrMessagePortOrServiceWorker::WindowProxy(DomRoot::from_ref(&*i)),
+            ),
+            Some(SrcObject::MessagePort(i)) => Some(
+                WindowProxyOrMessagePortOrServiceWorker::MessagePort(DomRoot::from_ref(&*i)),
+            ),
+            Some(SrcObject::ServiceWorker(i)) => Some(
+                WindowProxyOrMessagePortOrServiceWorker::ServiceWorker(DomRoot::from_ref(&*i)),
+            ),
+            None => None,
+        }
     }
 
     /// <https://html.spec.whatwg.org/multipage/#dom-messageevent-lasteventid>