Fix ws calls from an upgraded http domain

CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
  - Direct AJAX calls not getting upgraded to https [#530](https://github.com/zaproxy/zap-hud/issues/530)
+ - Fail to handle ws calls from an upgraded http domain [#525](https://github.com/zaproxy/zap-hud/issues/525) 
 
 ## [0.4.0] - 2019-06-07
 

src/main/java/org/zaproxy/zap/extension/hud/HttpUpgradeProxyListener.java

@@ -19,8 +19,12 @@
  */
 package org.zaproxy.zap.extension.hud;
 
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import org.apache.commons.httpclient.URI;
 import org.apache.commons.httpclient.URIException;
 import org.apache.log4j.Logger;
@@ -33,6 +37,11 @@
 
 public class HttpUpgradeProxyListener implements OverrideMessageProxyListener {
 
+    // Based on an example on
+    // https://stackoverflow.com/questions/24924072/website-url-validation-regex-in-java/37864510
+    private static final Pattern WSS_REGEX_PATTERN =
+            Pattern.compile("wss:\\/\\/(www\\.)?([\\w]+\\.)+[‌​\\w]{2,63}(:[0-9]+)?\\/?");
+
     private ExtensionHUD extHud;
 
     private Logger LOG = Logger.getLogger(this.getClass());
@@ -121,11 +130,36 @@ public boolean onHttpResponseReceived(HttpMessage msg) {
                                 .setBody(respBody.replace("http://" + domain, "https://" + domain));
                         msg.getResponseHeader().setContentLength(msg.getRequestBody().length());
                     }
+                    if (respBody.contains("ws://")) {
+                        // Need to replace hardcoded ws URLs with wss ones
+                        String body = respBody.replace("ws://", "wss://");
+                        // Now extract all of the wss urls to flag them as upgraded
+                        for (URI uri : extractWssUrls(body)) {
+                            LOG.debug("Adding upgraded ws URL: " + uri);
+                            extHud.addUpgradedHttpsDomain(uri);
+                        }
+                        msg.getResponseBody().setBody(body);
+                        msg.getResponseHeader().setContentLength(msg.getRequestBody().length());
+                    }
                 }
             } catch (URIException e) {
                 LOG.error(e.getMessage(), e);
             }
         }
         return false;
     }
+
+    protected static List<URI> extractWssUrls(String str) {
+        List<URI> list = new ArrayList<URI>();
+        Matcher m = WSS_REGEX_PATTERN.matcher(str);
+        while (m.find()) {
+            String wsUrl = m.group();
+            try {
+                list.add(new URI(wsUrl, false));
+            } catch (URIException e) {
+                // Not a valid url, ignore it
+            }
+        }
+        return list;
+    }
 }

src/test/java/org/zaproxy/zap/extension/hud/HttpUpgradeProxyListenerTest.java

@@ -20,7 +20,9 @@
 package org.zaproxy.zap.extension.hud;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import java.util.List;
 import org.apache.commons.httpclient.URI;
 import org.apache.commons.httpclient.URIException;
 import org.junit.jupiter.api.Test;
@@ -127,4 +129,38 @@ public void notUpgradedHttpUrlNotUpgradedInHtml() throws URIException {
         // Then
         assertEquals(body, msg.getResponseBody().toString());
     }
+
+    @Test
+    public void extractValidWssUrls() throws URIException, NullPointerException {
+        // Given
+        URI wssUri1 = new URI("wss://www.example.com", true);
+        URI wssUri2 = new URI("wss://www.example2.com", true);
+
+        String str = "  " + wssUri1 + " http://example.com  " + wssUri2;
+
+        // When
+        List<URI> list = HttpUpgradeProxyListener.extractWssUrls(str);
+
+        // Then
+        assertEquals(2, list.size());
+        assertTrue(list.contains(wssUri1));
+        assertTrue(list.contains(wssUri2));
+    }
+
+    @Test
+    public void extractValidWssUrlsWithPorts() throws URIException, NullPointerException {
+        // Given
+        URI wssUri1 = new URI("wss://www.example.com:80", true);
+        URI wssUri2 = new URI("wss://www.example2.com:9070/", true);
+
+        String str = "  " + wssUri1 + " ws://example.com  " + wssUri2;
+
+        // When
+        List<URI> list = HttpUpgradeProxyListener.extractWssUrls(str);
+
+        // Then
+        assertEquals(2, list.size());
+        assertTrue(list.contains(wssUri1));
+        assertTrue(list.contains(wssUri2));
+    }
 }