2 files changed, 18 insertions, 30 deletions
diff --git a/pages/test_ui_component.coffee b/pages/test_ui_component.coffee
index 886a3470..e140fb14 100644
--- a/pages/test_ui_component.coffee
+++ b/pages/test_ui_component.coffee
@@ -1,6 +1,8 @@
-UIComponentServer.addEventListener "message", (event) ->
+UIComponentServer.registerHandler (event) ->
   document.body.innerHTML = event.data
 
 window.addEventListener "keydown", (event) ->
   if KeyboardUtils.isEscape event
     UIComponentServer.postMessage "hide"
+  else
+    UIComponentServer.postMessage event.keyCode
diff --git a/pages/ui_component_server.coffee b/pages/ui_component_server.coffee
index ccbb9e2b..8b43095b 100644
--- a/pages/ui_component_server.coffee
+++ b/pages/ui_component_server.coffee
@@ -1,41 +1,27 @@
-# Register the port recieved from the parent window, and stop listening for messages on the window object.
-window.addEventListener "message", (event) ->
-  return unless event.source == window.parent
-  currentFunction = arguments.callee
 
-  # Check event.data against iframeMessageSecret so we can determine that this message hasn't been spoofed.
-  chrome.storage.local.get "iframeMessageSecret", ({iframeMessageSecret: secret}) ->
-    return unless event.data == secret
+# Fetch the Vimium secret, register the port recieved from the parent window, and stop listening for messages
+# on the window object. vimiumSecret is accessible only within the current instantion of Vimium.  So a
+# malicious host page trying to register its own port can do no better than guessing.
+registerPort = (event) ->
+  chrome.storage.local.get "vimiumSecret", ({vimiumSecret: secret}) ->
+    return unless event.source == window.parent and event.data == secret
     UIComponentServer.portOpen event.ports[0]
-    window.removeEventListener "message", currentFunction # Stop listening for message events.
+    window.removeEventListener "message", registerPort
+
+window.addEventListener "message", registerPort
 
 UIComponentServer =
   ownerPagePort: null
-  messageEventListeners: []
-  exitOnEsc: true
+  handleMessage: null
 
   portOpen: (@ownerPagePort) ->
-    @ownerPagePort.onmessage = (event) => @handleMessage event
-
-  postMessage: (message) -> @ownerPagePort.postMessage message
-
-  # Execute each event listener on the current event until we get a non-null falsy return value.
-  handleMessage: (event) ->
-    for listener in @messageEventListeners
-      retVal = listener.call this, event
-      retVal ?= true
-      return false unless retVal
-    true
+    @ownerPagePort.onmessage = (event) =>
+      @handleMessage event if @handleMessage
 
-  addEventListener: (type, listener) ->
-    if type == "message"
-      @messageEventListeners.push listener
-    undefined
+  registerHandler: (@handleMessage) ->
 
-  removeEventListener: (type, listener) ->
-    if type == "message"
-      @messageEventListeners = @messageEventListeners.filter (f) -> f != listener
-    undefined
+  postMessage: (message) ->
+    @ownerPagePort.postMessage message if @ownerPagePort
 
 root = exports ? window
 root.UIComponentServer = UIComponentServer