Merge pull request #1957 from thc202/selenium-rm-webdrivers-provider

selenium: quit WebDrivers when removing provider
zaproxy · Mar 15, 2019 · 0ceff13 · 0ceff13
2 parents d210ead + 8f280ad
commit 0ceff13
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/src/org/zaproxy/zap/extension/selenium/ExtensionSelenium.java b/src/org/zaproxy/zap/extension/selenium/ExtensionSelenium.java
@@ -117,7 +117,7 @@ public class ExtensionSelenium extends ExtensionAdaptor {
      * A list containing all of the proxied WebDrivers opened, so that they can
      * be closed when ZAP is closed.
      */
-    private List<WebDriver> proxiedWebDrivers = new ArrayList<WebDriver>();
+    private Map<String, List<WebDriver>> proxiedWebDrivers = new HashMap<>();
 
     private List<WeakReference<ProvidedBrowsersComboBoxModel>> providedBrowserComboBoxModels =
             new ArrayList<WeakReference<ProvidedBrowsersComboBoxModel>>();
@@ -205,8 +205,16 @@ public boolean canUnload() {
     @Override
     public void stop() {
         super.stop();
-        for (WebDriver wd : this.proxiedWebDrivers) {
-            // Just to make sure
+        this.proxiedWebDrivers.values().forEach(ExtensionSelenium::quitWebDrivers);
+        this.proxiedWebDrivers.clear();
+    }
+
+    private static void quitWebDrivers(List<WebDriver> drivers) {
+        if (drivers == null || drivers.isEmpty()) {
+            return;
+        }
+
+        for (WebDriver wd : drivers) {
             try {
                 wd.quit();
             } catch (Exception ex) {
@@ -290,6 +298,7 @@ private static void validateWebDriverProvider(SingleWebDriverProvider webDriverP
     public void removeWebDriverProvider(SingleWebDriverProvider webDriverProvider) {
         validateWebDriverProvider(webDriverProvider);
 
+        quitWebDrivers(proxiedWebDrivers.remove(webDriverProvider.getId()));
         webDriverProviders.remove(webDriverProvider.getId());
         providedBrowsers.remove(webDriverProvider.getProvidedBrowser().getId());
         buildProvidedBrowserUIList();
@@ -567,7 +576,7 @@ public WebDriver getProxiedBrowser(final int requester, final String providedBro
             Model.getSingleton().getOptionsParam().getProxyParam().getProxyPort());
 
         if (webDriver != null) {
-            proxiedWebDrivers.add(webDriver);
+            proxiedWebDrivers.computeIfAbsent(providedBrowserId, k -> new ArrayList<>()).add(webDriver);
             if (url != null) {
                 webDriver.get(url);
             }

diff --git a/src/org/zaproxy/zap/extension/selenium/ZapAddOn.xml b/src/org/zaproxy/zap/extension/selenium/ZapAddOn.xml
@@ -10,6 +10,7 @@
 	<changes>
 	<![CDATA[
 	Remove support for Internet Explorer, does not support required capabilities.<br>
+	Quit corresponding WebDrivers when removing WebDriver provider.<br>
 	]]>
 	</changes>
 	<extensions>