[java] JdkHttpClient - Handle a possible freeze when an incomplete http response is received (#12055)

joerg1985 · diemol · web-flow · commit 76701ad4b162 · 2023-05-18T12:47:13.000+02:00
Handle a possible freeze when an incomplete http response is received

Co-authored-by: Diego Molina &lt;diemol@users.noreply.github.com&gt;
diff --git a/java/src/org/openqa/selenium/remote/http/jdk/JdkHttpClient.java b/java/src/org/openqa/selenium/remote/http/jdk/JdkHttpClient.java
@@ -211,7 +211,31 @@ public void onError(java.net.http.WebSocket webSocket, Throwable error) {
                   }
                 });
 
-    java.net.http.WebSocket underlyingSocket = webSocketCompletableFuture.join();
+    java.net.http.WebSocket underlyingSocket;
+
+    try {
+      underlyingSocket = webSocketCompletableFuture.get(readTimeout.toMillis(), TimeUnit.MILLISECONDS);
+    } catch (CancellationException e) {
+      throw new WebDriverException(e.getMessage(), e);
+    } catch (ExecutionException e) {
+      Throwable cause = e.getCause();
+
+      if (cause instanceof HttpTimeoutException) {
+        throw new TimeoutException(cause);
+      } else if (cause instanceof IOException) {
+        throw new UncheckedIOException((IOException) cause);
+      } else if (cause instanceof RuntimeException) {
+        throw (RuntimeException) cause;
+      }
+
+      throw new WebDriverException((cause != null) ? cause : e);
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw new RuntimeException(e);
+    } catch (java.util.concurrent.TimeoutException e) {
+      webSocketCompletableFuture.cancel(true);
+      throw new TimeoutException(e);
+    }
 
     this.websocket =
         new WebSocket() {
@@ -340,7 +364,31 @@ public HttpResponse execute(HttpRequest req) throws UncheckedIOException {
       // - not run into https://bugs.openjdk.org/browse/JDK-8304701
       for (int i = 0; i < 100; i++) {
         java.net.http.HttpRequest request = messages.createRequest(req, rawUri);
-        java.net.http.HttpResponse<byte[]> response = client.send(request, byteHandler);
+        java.net.http.HttpResponse<byte[]> response;
+
+        // use sendAsync to not run into https://bugs.openjdk.org/browse/JDK-8258397
+        CompletableFuture<java.net.http.HttpResponse<byte[]>> future = client.sendAsync(request, byteHandler);
+
+        try {
+          response = future.get(readTimeout.toMillis(), TimeUnit.MILLISECONDS);
+        } catch (CancellationException e) {
+          throw new WebDriverException(e.getMessage(), e);
+        } catch (ExecutionException e) {
+          Throwable cause = e.getCause();
+
+          if (cause instanceof HttpTimeoutException) {
+            throw new TimeoutException(cause);
+          } else if (cause instanceof IOException) {
+            throw (IOException) cause;
+          } else if (cause instanceof RuntimeException) {
+            throw (RuntimeException) cause;
+          }
+
+          throw new WebDriverException((cause != null) ? cause : e);
+        } catch (java.util.concurrent.TimeoutException e) {
+          future.cancel(true);
+          throw new TimeoutException(e);
+        }
 
         switch (response.statusCode()) {
           case 301:
@@ -376,8 +424,6 @@ public HttpResponse execute(HttpRequest req) throws UncheckedIOException {
       }
 
       throw new ProtocolException("Too many redirects: 101");
-    } catch (HttpTimeoutException e) {
-      throw new TimeoutException(e);
     } catch (IOException e) {
       throw new UncheckedIOException(e);
     } catch (InterruptedException e) {
