Closes mozilla-mobile#9550: Check port instance before disconnecting

Amejia481 · Jan 29, 2021 · a8a2833 · a8a2833
1 parent d3579e0
commit a8a2833
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 8 deletions.
diff --git a/...y/src/main/java/mozilla/components/browser/engine/gecko/webextension/GeckoWebExtension.kt b/...y/src/main/java/mozilla/components/browser/engine/gecko/webextension/GeckoWebExtension.kt
@@ -36,7 +36,7 @@ class GeckoWebExtension(
     val runtime: GeckoRuntime
 ) : WebExtension(nativeExtension.id, nativeExtension.location, true) {
 
-    private val connectedPorts: MutableMap<PortId, Port> = mutableMapOf()
+    private val connectedPorts: MutableMap<PortId, GeckoPort> = mutableMapOf()
     private val logger = Logger("GeckoWebExtension")
 
     /**
@@ -57,8 +57,11 @@ class GeckoWebExtension(
             }
 
             override fun onDisconnect(port: GeckoNativeWebExtension.Port) {
-                connectedPorts.remove(PortId(name))
-                messageHandler.onPortDisconnected(GeckoPort(port))
+                val connectedPort = connectedPorts[PortId(name)]
+                if (connectedPort != null && connectedPort.nativePort == port) {
+                    connectedPorts.remove(PortId(name))
+                    messageHandler.onPortDisconnected(GeckoPort(port))
+                }
             }
         }
 
@@ -96,9 +99,11 @@ class GeckoWebExtension(
             }
 
             override fun onDisconnect(port: GeckoNativeWebExtension.Port) {
-                val geckoPort = GeckoPort(port, session)
-                connectedPorts.remove(PortId(name, session))
-                messageHandler.onPortDisconnected(geckoPort)
+                val connectedPort = connectedPorts[PortId(name, session)]
+                if (connectedPort != null && connectedPort.nativePort == port) {
+                    connectedPorts.remove(PortId(name, session))
+                    messageHandler.onPortDisconnected(connectedPort)
+                }
             }
         }
 

diff --git a/...c/test/java/mozilla/components/browser/engine/gecko/webextension/GeckoWebExtensionTest.kt b/...c/test/java/mozilla/components/browser/engine/gecko/webextension/GeckoWebExtensionTest.kt
@@ -86,7 +86,10 @@ class GeckoWebExtensionTest {
         verify(messageHandler).onPortMessage(eq(portMessage), portCaptor.capture())
         assertSame(port, (portCaptor.value as GeckoPort).nativePort)
 
-        // Verify disconnected port is forwarded to message handler
+        // Verify disconnected port is forwarded to message handler if connected
+        portDelegate.onDisconnect(mock())
+        verify(messageHandler, never()).onPortDisconnected(portCaptor.capture())
+
         portDelegate.onDisconnect(port)
         verify(messageHandler).onPortDisconnected(portCaptor.capture())
         assertSame(port, (portCaptor.value as GeckoPort).nativePort)
@@ -145,7 +148,10 @@ class GeckoWebExtensionTest {
         assertSame(port, (portCaptor.value as GeckoPort).nativePort)
         assertSame(session, (portCaptor.value as GeckoPort).engineSession)
 
-        // Verify disconnected port is forwarded to message handler
+        // Verify disconnected port is forwarded to message handler if connected
+        portDelegate.onDisconnect(mock())
+        verify(messageHandler, never()).onPortDisconnected(portCaptor.capture())
+
         portDelegate.onDisconnect(port)
         verify(messageHandler).onPortDisconnected(portCaptor.capture())
         assertSame(port, (portCaptor.value as GeckoPort).nativePort)