Test aborting with exception before WS upgrade completes.

Signed-off-by: Lachlan Roberts <[email protected]>
jetty · Jul 15, 2020 · f9750c9 · f9750c9
1 parent 25a7da2
commit f9750c9
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 3 deletions.
diff --git a/...ocket-tests/src/test/java/org/eclipse/jetty/websocket/tests/client/ConnectFutureTest.java b/...ocket-tests/src/test/java/org/eclipse/jetty/websocket/tests/client/ConnectFutureTest.java
@@ -303,7 +303,49 @@ public void testFutureTimeout() throws Exception
     }
 
     @Test
-    public void testAbortWithException() throws Exception
+    public void testAbortWithExceptionBeforeUpgrade() throws Exception
+    {
+        CountDownLatch exitCreator = new CountDownLatch(1);
+        start(c ->
+        {
+            c.addMapping("/", (req, res) ->
+            {
+                try
+                {
+                    exitCreator.await();
+                    return new EchoSocket();
+                }
+                catch (InterruptedException e)
+                {
+                    throw new IllegalStateException(e);
+                }
+            });
+        });
+
+        // Complete the CompletableFuture with an exception the during the call to onOpened.
+        CloseTrackingEndpoint clientSocket = new CloseTrackingEndpoint();
+        Future<Session> connect = client.connect(clientSocket, WSURI.toWebsocket(server.getURI()));
+        CompletableFuture<Session> completableFuture = (CompletableFuture<Session>)connect;
+        assertTrue(completableFuture.completeExceptionally(new WebSocketException("custom exception")));
+        exitCreator.countDown();
+
+        // Exception from the future is correct.
+        ExecutionException futureError = assertThrows(ExecutionException.class, () -> connect.get(5, TimeUnit.SECONDS));
+        Throwable cause = futureError.getCause();
+        assertThat(cause, instanceOf(WebSocketException.class));
+        assertThat(cause.getMessage(), is("custom exception"));
+
+        // Exception from the endpoint is correct.
+        assertTrue(clientSocket.errorLatch.await(5, TimeUnit.SECONDS));
+        Throwable endpointError = clientSocket.error.get();
+        assertThat(endpointError, instanceOf(UpgradeException.class));
+        Throwable endpointErrorCause = endpointError.getCause();
+        assertThat(endpointError, instanceOf(WebSocketException.class));
+        assertThat(endpointErrorCause.getMessage(), is("custom exception"));
+    }
+
+    @Test
+    public void testAbortWithExceptionAfterUpgrade() throws Exception
     {
         start(c -> c.addMapping("/", EchoSocket.class));
         CountDownLatch exitOnOpen = new CountDownLatch(1);

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,7 +29,6 @@
 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;
@@ -425,7 +424,7 @@ public WebSocketUpgradeRequest(WebSocketClient wsClient, HttpClient httpClient,
         this.fut = new CompletableFuture<>();
         this.fut.whenComplete((session, throwable) ->
         {
-            if (throwable instanceof CancellationException)
+            if (throwable != null)
                 abort(throwable);
         });