Issue jetty#5018 - cancellation of WebSocketClient.connect Future sho…

…uld fail upgrade Signed-off-by: Lachlan Roberts <lachlan@webtide.com>
leonchen83 · Aug 3, 2020 · 1155717 · 1155717
1 parent 980ad4c
commit 1155717
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 6 deletions.
diff --git a/...et/websocket-client/src/main/java/org/eclipse/jetty/websocket/client/WebSocketClient.java b/...et/websocket-client/src/main/java/org/eclipse/jetty/websocket/client/WebSocketClient.java
@@ -343,7 +343,7 @@ public Future<Session> connect(Object websocket, URI toUri, ClientUpgradeRequest
         }
 
         String scheme = toUri.getScheme().toLowerCase(Locale.ENGLISH);
-        if (("ws".equals(scheme) == false) && ("wss".equals(scheme) == false))
+        if ((!"ws".equals(scheme)) && (!"wss".equals(scheme)))
         {
             throw new IllegalArgumentException("WebSocket URI scheme only supports [ws] and [wss], not [" + scheme + "]");
         }

diff --git a/...cket-client/src/main/java/org/eclipse/jetty/websocket/client/WebSocketUpgradeRequest.java b/...cket-client/src/main/java/org/eclipse/jetty/websocket/client/WebSocketUpgradeRequest.java
@@ -29,6 +29,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.TreeMap;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ThreadLocalRandom;
@@ -421,12 +422,17 @@ public WebSocketUpgradeRequest(WebSocketClient wsClient, HttpClient httpClient,
         }
         this.localEndpoint = this.wsClient.getEventDriverFactory().wrap(localEndpoint);
 
-        this.fut = new CompletableFuture<Session>();
+        this.fut = new CompletableFuture<>();
+        this.fut.whenComplete((session, throwable) ->
+        {
+            if (throwable instanceof CancellationException)
+                abort(throwable);
+        });
 
         getConversation().setAttribute(HttpConnectionUpgrader.class.getName(), this);
     }
 
-    private final String genRandomKey()
+    private String genRandomKey()
     {
         byte[] bytes = new byte[16];
         ThreadLocalRandom.current().nextBytes(bytes);
@@ -580,7 +586,7 @@ public void upgrade(HttpResponse response, HttpConnectionOverHTTP oldConn)
         String expectedHash = AcceptHash.hashKey(reqKey);
         String respHash = response.getHeaders().get(HttpHeader.SEC_WEBSOCKET_ACCEPT);
 
-        if (expectedHash.equalsIgnoreCase(respHash) == false)
+        if (!expectedHash.equalsIgnoreCase(respHash))
         {
             throw new HttpResponseException("Invalid Sec-WebSocket-Accept hash", response);
         }

diff --git a/...t/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/WebSocketSession.java b/...t/websocket-common/src/main/java/org/eclipse/jetty/websocket/common/WebSocketSession.java
@@ -404,7 +404,7 @@ public void callApplicationOnError(Throwable cause)
         if (LOG.isDebugEnabled())
             LOG.debug("callApplicationOnError()", cause);
 
-        if (openFuture != null && !openFuture.isDone())
+        if (openFuture != null)
             openFuture.completeExceptionally(cause);
 
         // Only notify onError if onClose has not been called.
@@ -478,7 +478,7 @@ public void open()
             return;
         }
 
-        try (ThreadClassLoaderScope scope = new ThreadClassLoaderScope(classLoader))
+        try (ThreadClassLoaderScope ignored = new ThreadClassLoaderScope(classLoader))
         {
             // Upgrade success
             if (connection.opening())
@@ -535,6 +535,11 @@ public void setExtensionFactory(ExtensionFactory extensionFactory)
     public void setFuture(CompletableFuture<Session> fut)
     {
         this.openFuture = fut;
+        fut.whenComplete((s, t) ->
+        {
+            if (t != null)
+                close(t);
+        });
     }
 
     /**