Merge pull request #48508 from emattheis/wsnext-callback-argument-fix

Allow use of Connection interface in callback arguments for WebSockets Next

Author: geoand

extensions/websockets-next/deployment/src/main/java/io/quarkus/websockets/next/deployment/ConnectionCallbackArgument.java

@@ -12,13 +12,15 @@ class ConnectionCallbackArgument implements CallbackArgument {
     public boolean matches(ParameterContext context) {
         DotName paramTypeName = context.parameter().type().name();
         if (context.callbackTarget() == Target.SERVER) {
-            if (WebSocketDotNames.WEB_SOCKET_CONNECTION.equals(paramTypeName)) {
+            if (WebSocketDotNames.WEB_SOCKET_CONNECTION.equals(paramTypeName)
+                    || WebSocketDotNames.CONNECTION.equals(paramTypeName)) {
                 return true;
             } else if (WebSocketDotNames.WEB_SOCKET_CLIENT_CONNECTION.equals(paramTypeName)) {
                 throw new WebSocketException("@WebSocket callback method may not accept WebSocketClientConnection");
             }
         } else if (context.callbackTarget() == Target.CLIENT) {
-            if (WebSocketDotNames.WEB_SOCKET_CLIENT_CONNECTION.equals(paramTypeName)) {
+            if (WebSocketDotNames.WEB_SOCKET_CLIENT_CONNECTION.equals(paramTypeName)
+                    || WebSocketDotNames.CONNECTION.equals(paramTypeName)) {
                 return true;
             } else if (WebSocketDotNames.WEB_SOCKET_CONNECTION.equals(paramTypeName)) {
                 throw new WebSocketException("@WebSocketClient callback method may not accept WebSocketConnection");

extensions/websockets-next/deployment/src/main/java/io/quarkus/websockets/next/deployment/WebSocketDotNames.java

@@ -5,6 +5,7 @@
 import org.jboss.jandex.DotName;
 
 import io.quarkus.websockets.next.CloseReason;
+import io.quarkus.websockets.next.Connection;
 import io.quarkus.websockets.next.HandshakeRequest;
 import io.quarkus.websockets.next.OnBinaryMessage;
 import io.quarkus.websockets.next.OnClose;
@@ -32,6 +33,7 @@ final class WebSocketDotNames {
 
     static final DotName WEB_SOCKET = DotName.createSimple(WebSocket.class);
     static final DotName WEB_SOCKET_CLIENT = DotName.createSimple(WebSocketClient.class);
+    static final DotName CONNECTION = DotName.createSimple(Connection.class);
     static final DotName WEB_SOCKET_CONNECTION = DotName.createSimple(WebSocketConnection.class);
     static final DotName WEB_SOCKET_CLIENT_CONNECTION = DotName.createSimple(WebSocketClientConnection.class);
     static final DotName WEB_SOCKET_CONNECTOR = DotName.createSimple(WebSocketConnector.class);

extensions/websockets-next/deployment/src/test/java/io/quarkus/websockets/next/test/args/ConnectionArgumentTest.java

@@ -1,6 +1,7 @@
 package io.quarkus.websockets.next.test.args;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 import java.net.URI;
 
@@ -11,6 +12,8 @@
 
 import io.quarkus.test.QuarkusUnitTest;
 import io.quarkus.test.common.http.TestHTTPResource;
+import io.quarkus.websockets.next.Connection;
+import io.quarkus.websockets.next.OnOpen;
 import io.quarkus.websockets.next.OnTextMessage;
 import io.quarkus.websockets.next.WebSocket;
 import io.quarkus.websockets.next.WebSocketConnection;
@@ -52,6 +55,11 @@ public static class Echo {
         @Inject
         WebSocketConnection c;
 
+        @OnOpen
+        void connect(Connection connection) {
+            assertInstanceOf(WebSocketConnection.class, connection);
+        }
+
         @OnTextMessage
         Uni<Void> process(WebSocketConnection connection, String message) throws InterruptedException {
             assertEquals(c.id(), connection.id());

extensions/websockets-next/deployment/src/test/java/io/quarkus/websockets/next/test/client/ClientEndpointTest.java

@@ -1,13 +1,15 @@
 package io.quarkus.websockets.next.test.client;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.net.URI;
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 
 import jakarta.inject.Inject;
 
@@ -16,6 +18,7 @@
 
 import io.quarkus.test.QuarkusUnitTest;
 import io.quarkus.test.common.http.TestHTTPResource;
+import io.quarkus.websockets.next.Connection;
 import io.quarkus.websockets.next.OnClose;
 import io.quarkus.websockets.next.OnOpen;
 import io.quarkus.websockets.next.OnPingMessage;
@@ -50,6 +53,8 @@ void testClient() throws InterruptedException {
                 // The value will be encoded automatically
                 .pathParam("name", "Lu=")
                 .connectAndAwait();
+        assertTrue(ClientEndpoint.OPEN_LATCH.await(5, TimeUnit.SECONDS));
+        assertInstanceOf(WebSocketClientConnection.class, ClientEndpoint.CONNECTION.get());
         assertEquals("Lu=", connection.pathParam("name"));
         connection.sendPingAndAwait(ping);
         connection.sendTextAndAwait("Hi!");
@@ -103,6 +108,10 @@ void close() {
     @WebSocketClient(path = "/endpoint/{name}")
     public static class ClientEndpoint {
 
+        static final CountDownLatch OPEN_LATCH = new CountDownLatch(1);
+
+        static final AtomicReference<Connection> CONNECTION = new AtomicReference<>();
+
         static final CountDownLatch PING_LATCH = new CountDownLatch(1);
 
         static final CountDownLatch MESSAGE_LATCH = new CountDownLatch(2);
@@ -111,6 +120,12 @@ public static class ClientEndpoint {
 
         static final CountDownLatch CLOSED_LATCH = new CountDownLatch(1);
 
+        @OnOpen
+        void onOpen(Connection connection) {
+            CONNECTION.set(connection);
+            OPEN_LATCH.countDown();
+        }
+
         @OnTextMessage
         void onMessage(@PathParam String name, String message, WebSocketClientConnection connection) {
             if (!name.equals(connection.pathParam("name"))) {