From c70c3ce7d8f4a345fc743f1fd1b59fc200d8633b Mon Sep 17 00:00:00 2001
From: hensm <m+github@matt.tf>
Date: Thu, 21 Mar 2019 21:07:37 +0000
Subject: [PATCH] Improve session message handling

---
 app/src/Media.ts                     |  4 ++-
 app/src/Session.ts                   |  4 ++-
 app/src/main.ts                      |  1 +
 ext/src/shim/cast/classes/Session.ts | 15 ++++++++--
 ext/src/shim/messageBridge.ts        | 41 ++++++++++++++++++++++++----
 5 files changed, 55 insertions(+), 10 deletions(-)

diff --git a/app/src/Media.ts b/app/src/Media.ts
index 3627ead..ee2fd41 100644
--- a/app/src/Media.ts
+++ b/app/src/Media.ts
@@ -1,5 +1,7 @@
 "use strict";
 
+import { Channel } from "castv2";
+
 import Session from "./Session";
 
 import { Message
@@ -27,7 +29,7 @@ export default class Media {
     private mediaSessionId: number;
     private referenceId: string;
     private session: Session;
-    private channel: any;
+    private channel: Channel;
     private sendMessageCallback: SendMessageCallback;
 
     constructor (
diff --git a/app/src/Session.ts b/app/src/Session.ts
index 5aa3534..76e1bc7 100644
--- a/app/src/Session.ts
+++ b/app/src/Session.ts
@@ -37,10 +37,12 @@ export default class Session {
           , port: number
           , appId: string
           , sessionId: number
+          , referenceId: string
           , sendMessageCallback: SendMessageCallback) {
 
-        this.sendMessageCallback = sendMessageCallback;
         this.sessionId = sessionId;
+        this.referenceId = referenceId;
+        this.sendMessageCallback = sendMessageCallback;
 
         this.client = new Client();
 
diff --git a/app/src/main.ts b/app/src/main.ts
index d10f0e8..45c4dc4 100755
--- a/app/src/main.ts
+++ b/app/src/main.ts
@@ -97,6 +97,7 @@ async function handleMessage (message: Message) {
                       , message.data.port
                       , message.data.appId
                       , message.data.sessionId
+                      , message._id
                       , sendMessage));
             }
         }
diff --git a/ext/src/shim/cast/classes/Session.ts b/ext/src/shim/cast/classes/Session.ts
index 75884e8..838123f 100755
--- a/ext/src/shim/cast/classes/Session.ts
+++ b/ext/src/shim/cast/classes/Session.ts
@@ -16,7 +16,9 @@ import { ErrorCode
        , SessionStatus
        , VolumeControlType } from "../enums";
 
-import { onMessage, sendMessageResponse } from "../../messageBridge";
+import { ListenerObject
+       , onMessage
+       , sendMessageResponse } from "../../messageBridge";
 
 import { Callbacks
        , CallbacksMap
@@ -29,6 +31,7 @@ import { Callbacks
 
 
 const _id = new WeakMap<Session, string>();
+const _listener = new WeakMap<Session, ListenerObject>();
 
 const _messageListeners = new WeakMap<
         Session, Map<string, Set<MessageListener>>>();
@@ -90,7 +93,7 @@ export default class Session {
             });
         }
 
-        onMessage(message => {
+        const listenerObject = onMessage(message => {
             // Filter other session messages
             if (message._id && message._id !== _id.get(this)) {
                 return;
@@ -98,6 +101,9 @@ export default class Session {
 
             switch (message.subject) {
                 case "shim:/session/stopped": {
+                    // Disconnect from extension messages
+                    _listener.get(this).disconnect();
+
                     this.status = SessionStatus.STOPPED;
 
                     for (const listener of _updateListeners.get(this)) {
@@ -215,6 +221,9 @@ export default class Session {
                     const [ successCallback, errorCallback ]
                             = _stopCallbacks.get(this).get(stopId);
 
+                    // Disconnect from extension messages
+                    _listener.get(this).disconnect();
+
                     if (error && errorCallback) {
                         errorCallback(new _Error(ErrorCode.SESSION_ERROR));
                     } else {
@@ -235,6 +244,8 @@ export default class Session {
                 }
             }
         });
+
+        _listener.set(this, listenerObject);
     }
 
 
diff --git a/ext/src/shim/messageBridge.ts b/ext/src/shim/messageBridge.ts
index fc7e957..7d32913 100644
--- a/ext/src/shim/messageBridge.ts
+++ b/ext/src/shim/messageBridge.ts
@@ -2,11 +2,16 @@
 
 import { Message } from "../types";
 
+
 type ListenerFunc = (message: Message) => void;
 
+export interface ListenerObject {
+    disconnect (): void;
+}
 
-export function onMessage (listener: ListenerFunc) {
-    document.addEventListener("__castMessage", (ev: CustomEvent) => {
+
+export function onMessage (listener: ListenerFunc): ListenerObject {
+    function on__castMessage (ev: CustomEvent) {
         listener(JSON.parse(ev.detail));
 
         /**
@@ -18,7 +23,19 @@ export function onMessage (listener: ListenerFunc) {
          * intercept the event, and cancel it.
          */
         ev.stopPropagation();
-    }, true);
+    }
+
+    document.addEventListener(
+            "__castMessage"
+          , on__castMessage, true);
+
+    return {
+        disconnect () {
+            document.removeEventListener(
+                    "__castMessage"
+                  , on__castMessage, true);
+        }
+    };
 }
 
 export function sendMessageResponse (message: Message) {
@@ -30,10 +47,22 @@ export function sendMessageResponse (message: Message) {
 }
 
 
-export function onMessageResponse (listener: ListenerFunc) {
-    document.addEventListener("__castMessageResponse", (ev: CustomEvent) => {
+export function onMessageResponse (listener: ListenerFunc): ListenerObject {
+    function on__castMessageResponse (ev: CustomEvent) {
         listener(JSON.parse(ev.detail));
-    }, true);
+    }
+
+    document.addEventListener(
+            "__castMessageResponse"
+          , on__castMessageResponse, true);
+
+    return {
+        disconnect () {
+            document.removeEventListener(
+                    "__castMessageResponse"
+                  , on__castMessageResponse, true);
+        }
+    };
 }
 
 export function sendMessage (message: Message) {