[java] cancel the pending responses while disabling network #12576

SeleniumHQ · Aug 27, 2023 · fc5509b · fc5509b
1 parent bbe2ac5
commit fc5509b
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 11 deletions.
diff --git a/java/src/org/openqa/selenium/devtools/DevTools.java b/java/src/org/openqa/selenium/devtools/DevTools.java
@@ -62,6 +62,14 @@ public void close() {
 
   public void disconnectSession() {
     if (cdpSession != null) {
+      try {
+        // ensure network interception does cancel the wait for responses
+        getDomains().network().disable();
+      } catch (Exception e) {
+        // Exceptions should not prevent closing the connection and the web driver
+        LOG.log(Level.WARNING, "Exception while disabling network", e);
+      }
+
       SessionID id = cdpSession;
       cdpSession = null;
       try {
@@ -71,7 +79,7 @@ public void disconnectSession() {
             timeout);
       } catch (Exception e) {
         // Exceptions should not prevent closing the connection and the web driver
-        LOG.warning("Exception while detaching from target: " + e.getMessage());
+        LOG.log(Level.WARNING, "Exception while detaching from target", e);
       }
     }
   }

diff --git a/java/src/org/openqa/selenium/devtools/idealized/Network.java b/java/src/org/openqa/selenium/devtools/idealized/Network.java
@@ -30,6 +30,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
@@ -45,6 +46,7 @@
 import org.openqa.selenium.devtools.DevTools;
 import org.openqa.selenium.devtools.DevToolsException;
 import org.openqa.selenium.devtools.Event;
+import org.openqa.selenium.devtools.NetworkInterceptor;
 import org.openqa.selenium.internal.Either;
 import org.openqa.selenium.internal.Require;
 import org.openqa.selenium.remote.http.Contents;
@@ -57,21 +59,33 @@ public abstract class Network<AUTHREQUIRED, REQUESTPAUSED> {
 
   private static final Logger LOG = Logger.getLogger(Network.class.getName());
 
+  private static final HttpResponse STOP_PROCESSING =
+      new HttpResponse()
+          .addHeader("Selenium-Interceptor", "Stop")
+          .setContent(Contents.utf8String("Interception is stopped"));
+
   private final Map<Predicate<URI>, Supplier<Credentials>> authHandlers = new LinkedHashMap<>();
   private final Filter defaultFilter = next -> next::execute;
   private volatile Filter filter = defaultFilter;
   protected final DevTools devTools;
 
   private final AtomicBoolean fetchEnabled = new AtomicBoolean();
+  private final Map<String, CompletableFuture<HttpResponse>> pendingResponses =
+      new ConcurrentHashMap<>();
 
   public Network(DevTools devtools) {
     this.devTools = Require.nonNull("DevTools", devtools);
   }
 
   public void disable() {
     fetchEnabled.set(false);
-    devTools.send(disableFetch());
-    devTools.send(enableNetworkCaching());
+    try {
+      devTools.send(disableFetch());
+      devTools.send(enableNetworkCaching());
+    } finally {
+      // we stopped the fetch we will not receive any pending responses
+      pendingResponses.values().forEach(cf -> cf.cancel(false));
+    }
 
     synchronized (authHandlers) {
       authHandlers.clear();
@@ -183,8 +197,6 @@ public void prepareToInterceptTraffic() {
           devTools.send(cancelAuth(authRequired));
         });
 
-    Map<String, CompletableFuture<HttpResponse>> responses = new ConcurrentHashMap<>();
-
     devTools.addListener(
         requestPausedEvent(),
         pausedRequest -> {
@@ -194,7 +206,7 @@ public void prepareToInterceptTraffic() {
 
             if (message.isRight()) {
               HttpResponse res = message.right();
-              CompletableFuture<HttpResponse> future = responses.remove(id);
+              CompletableFuture<HttpResponse> future = pendingResponses.remove(id);
 
               if (future == null) {
                 devTools.send(continueWithoutModification(pausedRequest));
@@ -210,18 +222,22 @@ public void prepareToInterceptTraffic() {
                     .andFinally(
                         req -> {
                           // Convert the selenium request to a CDP one and fulfill.
-
-                          CompletableFuture<HttpResponse> res = new CompletableFuture<>();
-                          responses.put(id, res);
-
                           devTools.send(continueRequest(pausedRequest, req));
+                          CompletableFuture<HttpResponse> res = new CompletableFuture<>();
+                          // Save the future after the browser accepted the continueRequest
+                          pendingResponses.put(id, res);
 
                           // Wait for the CDP response and send that back.
                           try {
                             return res.get();
                           } catch (InterruptedException e) {
                             Thread.currentThread().interrupt();
                             throw new WebDriverException(e);
+                          } catch (CancellationException e) {
+                            // The interception was intentionally stopped, network().disable() has
+                            // been called
+                            pendingResponses.remove(id);
+                            return STOP_PROCESSING;
                           } catch (ExecutionException e) {
                             if (fetchEnabled.get()) {
                               LOG.log(WARNING, e, () -> "Unable to process request");
@@ -231,9 +247,12 @@ public void prepareToInterceptTraffic() {
                         })
                     .execute(message.left());
 
-            if ("Continue".equals(forBrowser.getHeader("Selenium-Interceptor"))) {
+            if (forBrowser == NetworkInterceptor.PROCEED_WITH_REQUEST) {
               devTools.send(continueWithoutModification(pausedRequest));
               return;
+            } else if (forBrowser == STOP_PROCESSING) {
+              // The interception was intentionally stopped, network().disable() has been called
+              return;
             }
 
             devTools.send(fulfillRequest(pausedRequest, forBrowser));