From 7612405e34d282992b26a22d7bee921753020026 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B6rg=20Sautter?= <joerg.sautter@gmx.net>
Date: Sun, 18 Aug 2024 19:52:27 +0200
Subject: [PATCH] [grid] close the httpclient after connecting the websocket
 failed

---
 .../grid/node/ProxyNodeWebsockets.java        | 36 ++++++++++--------
 .../grid/router/ProxyWebsocketsIntoGrid.java  | 37 +++++++++++--------
 2 files changed, 42 insertions(+), 31 deletions(-)

diff --git a/java/src/org/openqa/selenium/grid/node/ProxyNodeWebsockets.java b/java/src/org/openqa/selenium/grid/node/ProxyNodeWebsockets.java
index 343828de34031..e3f656c069125 100644
--- a/java/src/org/openqa/selenium/grid/node/ProxyNodeWebsockets.java
+++ b/java/src/org/openqa/selenium/grid/node/ProxyNodeWebsockets.java
@@ -226,24 +226,30 @@ private Consumer<Message> createWsEndPoint(
     LOG.info("Establishing connection to " + uri);
 
     HttpClient client = clientFactory.createClient(ClientConfig.defaultConfig().baseUri(uri));
-    WebSocket upstream =
-        client.openSocket(
-            new HttpRequest(GET, uri.toString()),
-            new ForwardingListener(downstream, sessionConsumer, sessionId));
+    try {
+      WebSocket upstream =
+          client.openSocket(
+              new HttpRequest(GET, uri.toString()),
+              new ForwardingListener(downstream, sessionConsumer, sessionId));
 
-    return (msg) -> {
-      try {
-        upstream.send(msg);
-      } finally {
-        if (msg instanceof CloseMessage) {
-          try {
-            client.close();
-          } catch (Exception e) {
-            LOG.log(Level.WARNING, "Failed to shutdown the client of " + uri, e);
+      return (msg) -> {
+        try {
+          upstream.send(msg);
+        } finally {
+          if (msg instanceof CloseMessage) {
+            try {
+              client.close();
+            } catch (Exception e) {
+              LOG.log(Level.WARNING, "Failed to shutdown the client of " + uri, e);
+            }
           }
         }
-      }
-    };
+      };
+    } catch (Exception e) {
+      LOG.log(Level.WARNING, "Connecting to upstream websocket failed", e);
+      client.close();
+      throw e;
+    }
   }
 
   private static class ForwardingListener implements WebSocket.Listener {
diff --git a/java/src/org/openqa/selenium/grid/router/ProxyWebsocketsIntoGrid.java b/java/src/org/openqa/selenium/grid/router/ProxyWebsocketsIntoGrid.java
index af71d27611ad4..3cdf7784f02e1 100644
--- a/java/src/org/openqa/selenium/grid/router/ProxyWebsocketsIntoGrid.java
+++ b/java/src/org/openqa/selenium/grid/router/ProxyWebsocketsIntoGrid.java
@@ -68,24 +68,29 @@ public Optional<Consumer<Message>> apply(String uri, Consumer<Message> downstrea
 
       HttpClient client =
           clientFactory.createClient(ClientConfig.defaultConfig().baseUri(sessionUri));
-      WebSocket upstream =
-          client.openSocket(new HttpRequest(GET, uri), new ForwardingListener(downstream));
-
-      return Optional.of(
-          (msg) -> {
-            try {
-              upstream.send(msg);
-            } finally {
-              if (msg instanceof CloseMessage) {
-                try {
-                  client.close();
-                } catch (Exception e) {
-                  LOG.log(Level.WARNING, "Failed to shutdown the client of " + sessionUri, e);
+      try {
+        WebSocket upstream =
+            client.openSocket(new HttpRequest(GET, uri), new ForwardingListener(downstream));
+
+        return Optional.of(
+            (msg) -> {
+              try {
+                upstream.send(msg);
+              } finally {
+                if (msg instanceof CloseMessage) {
+                  try {
+                    client.close();
+                  } catch (Exception e) {
+                    LOG.log(Level.WARNING, "Failed to shutdown the client of " + sessionUri, e);
+                  }
                 }
               }
-            }
-          });
-
+            });
+      } catch (Exception e) {
+        LOG.log(Level.WARNING, "Connecting to upstream websocket failed", e);
+        client.close();
+        return Optional.empty();
+      }
     } catch (NoSuchSessionException e) {
       LOG.warning("Attempt to connect to non-existent session: " + uri);
       return Optional.empty();