Merge pull request zaproxy#1711 from thc202/websocket-rm-consumer-se

websocket: remove consumer on socket exception
thc202 · Jul 10, 2018 · e3019d4 · e3019d4
2 parents 6853447 + 6623620
commit e3019d4
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 11 deletions.
diff --git a/src/org/zaproxy/zap/extension/websocket/WebSocketAPI.java b/src/org/zaproxy/zap/extension/websocket/WebSocketAPI.java
@@ -20,6 +20,7 @@
 package org.zaproxy.zap.extension.websocket;
 
 import java.io.IOException;
+import java.net.SocketException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -241,24 +242,27 @@ public void onStateChange(State state, WebSocketProxy proxy) {
                         return;
                     }
 
-                    WebsocketEventConsumer consumer = evMap.remove(proxy.getChannelId());
-                    if (consumer == null) {
-                        return;
-                    }
-
-                    // TODO replace the loop with:
-                    // ZAP.getEventBus().unregisterConsumer(consumer);
-                    // once available in targeted ZAP version.
-                    for (String publisherName : consumer.getPublisherNames()) {
-                        ZAP.getEventBus().unregisterConsumer(consumer, publisherName);
-                    }
+                    removeEventConsumer(evMap.remove(proxy.getChannelId()));
                 }
             };
 
         }
         return observer;
     }
 
+    private void removeEventConsumer(WebsocketEventConsumer consumer) {
+        if (consumer == null) {
+            return;
+        }
+
+        // TODO replace the loop with:
+        // ZAP.getEventBus().unregisterConsumer(consumer);
+        // once available in targeted ZAP version.
+        for (String publisherName : consumer.getPublisherNames()) {
+            ZAP.getEventBus().unregisterConsumer(consumer, publisherName);
+        }
+    }
+
     private boolean sendWebSocketMessage(int channelId, boolean outgoing, String message) throws IOException {
         WebSocketProxy proxy = extension.getWebSocketProxy(channelId);
         if (proxy != null) {
@@ -493,6 +497,9 @@ public void eventReceived(Event ev) {
             }
             try {
                 sendWebSocketMessage(channelId, false, json.toString());
+            } catch (SocketException e) {
+                LOG.debug("Failed to dispatch event:", e);
+                removeEventConsumer(this);
             } catch (IOException e) {
                 LOG.error(e.getMessage(), e);
             }

diff --git a/src/org/zaproxy/zap/extension/websocket/ZapAddOn.xml b/src/org/zaproxy/zap/extension/websocket/ZapAddOn.xml
@@ -10,6 +10,7 @@
 	Register WebSocket Sender script type also in daemon mode.<br>
 	Fix an exception when dispatching events.<br>
 	Fix an exception while uninstalling the add-on with no GUI (Issue 4815).<br>
+	Remove event consumers when channel is no longer in expected state.<br>
 	]]>
 	</changes>
 	<classnames>

diff --git a/src/org/zaproxy/zap/extension/websocket/resources/help/contents/about.html b/src/org/zaproxy/zap/extension/websocket/resources/help/contents/about.html
@@ -21,6 +21,7 @@ <H3>Version 17 - TBD</H3>
 	<li>Register WebSocket Sender script type also in daemon mode.</li>
 	<li>Fix an exception when dispatching events.</li>
 	<li>Fix an exception while uninstalling the add-on with no GUI (Issue 4815).</li>
+	<li>Remove event consumers when channel is no longer in expected state.</li>
 </ul>
 
 <H3>Version 16 - 2018/06/05</H3>