Automatic websocket codecs reflection registration

extensions/websockets/server/deployment/src/main/java/io/quarkus/websockets/deployment/ServerWebSocketProcessor.java

@@ -2,8 +2,12 @@
 
 import java.lang.reflect.Modifier;
 import java.util.Collection;
+import java.util.List;
 import java.util.Optional;
+import java.util.stream.Collectors;
 
+import jakarta.websocket.Decoder;
+import jakarta.websocket.Encoder;
 import jakarta.websocket.server.ServerEndpoint;
 
 import org.jboss.jandex.AnnotationInstance;
@@ -19,7 +23,10 @@
 import io.quarkus.deployment.annotations.Record;
 import io.quarkus.deployment.builditem.CombinedIndexBuildItem;
 import io.quarkus.deployment.builditem.FeatureBuildItem;
+import io.quarkus.deployment.builditem.IndexDependencyBuildItem;
 import io.quarkus.deployment.builditem.nativeimage.ReflectiveClassBuildItem;
+import io.quarkus.deployment.builditem.nativeimage.ReflectiveHierarchyBuildItem;
+import io.quarkus.deployment.util.JandexUtil;
 import io.quarkus.vertx.http.deployment.FilterBuildItem;
 import io.quarkus.websockets.client.deployment.AnnotatedWebsocketEndpointBuildItem;
 import io.quarkus.websockets.client.deployment.ServerWebSocketContainerBuildItem;
@@ -32,6 +39,17 @@ public class ServerWebSocketProcessor {
 
     private static final DotName SERVER_ENDPOINT = DotName.createSimple(ServerEndpoint.class.getName());
 
+    public static final Collection<DotName> CODECS = List.of(
+            Decoder.TextStream.class,
+            Decoder.Text.class,
+            Decoder.BinaryStream.class,
+            Decoder.Binary.class,
+            Encoder.TextStream.class,
+            Encoder.Text.class,
+            Encoder.BinaryStream.class,
+            Encoder.Binary.class).stream().map(Class::getName).map(DotName::createSimple)
+            .collect(Collectors.toList());
+
     @BuildStep
     void holdConfig(BuildProducer<FeatureBuildItem> feature) {
         feature.produce(new FeatureBuildItem(Feature.WEBSOCKETS));
@@ -54,6 +72,24 @@ void scanForAnnotatedEndpoints(CombinedIndexBuildItem indexBuildItem,
         }
     }
 
+    @BuildStep
+    void buildIndexDependencies(BuildProducer<IndexDependencyBuildItem> indexDependencyProduer) {
+        indexDependencyProduer.produce(new IndexDependencyBuildItem("jakarta.websocket", "jakarta.websocket-client-api"));
+    }
+
+    @BuildStep
+    void scanForCodecs(CombinedIndexBuildItem index,
+            BuildProducer<ReflectiveHierarchyBuildItem> reflectiveHierarchyBuildItemProducer) {
+        CODECS.stream().forEach(
+                codec -> index.getIndex().getAllKnownImplementors(codec).stream()
+                        .filter(implementor -> !Modifier.isAbstract(implementor.flags()))
+                        .forEach(implementor -> JandexUtil.resolveTypeParameters(
+                                implementor.name(),
+                                codec, index.getIndex()).forEach(
+                                        typeParameter -> reflectiveHierarchyBuildItemProducer.produce(
+                                                new ReflectiveHierarchyBuildItem.Builder().type(typeParameter).build()))));
+    }
+
     @BuildStep
     @Record(ExecutionTime.STATIC_INIT)
     public ServerWebSocketContainerFactoryBuildItem factory(WebsocketServerRecorder recorder) {

integration-tests/websockets/pom.xml

@@ -18,6 +18,10 @@
       <groupId>io.quarkus</groupId>
       <artifactId>quarkus-websockets</artifactId>
     </dependency>
+    <dependency>
+      <groupId>io.quarkus</groupId>
+      <artifactId>quarkus-jsonb</artifactId>
+    </dependency>
     <dependency>
       <groupId>io.quarkus</groupId>
       <artifactId>quarkus-arc</artifactId>
@@ -53,6 +57,19 @@
         </exclusion>
       </exclusions>
     </dependency>
+    <dependency>
+      <groupId>io.quarkus</groupId>
+      <artifactId>quarkus-jsonb-deployment</artifactId>
+      <version>${project.version}</version>
+      <type>pom</type>
+      <scope>test</scope>
+      <exclusions>
+        <exclusion>
+          <groupId>*</groupId>
+          <artifactId>*</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
     <dependency>
       <groupId>io.quarkus</groupId>
       <artifactId>quarkus-websockets-deployment</artifactId>

integration-tests/websockets/src/main/java/io/quarkus/websockets/ChatCodecServer.java

@@ -0,0 +1,57 @@
+package io.quarkus.websockets;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.UUID;
+
+import jakarta.enterprise.context.ApplicationScoped;
+import jakarta.websocket.OnClose;
+import jakarta.websocket.OnError;
+import jakarta.websocket.OnMessage;
+import jakarta.websocket.OnOpen;
+import jakarta.websocket.Session;
+import jakarta.websocket.server.ServerEndpoint;
+
+@ServerEndpoint(value = "/codec", encoders = ChatMessageEncoder.class, decoders = ChatMessageDecoder.class)
+@ApplicationScoped
+public class ChatCodecServer {
+
+    Collection<Session> sessions = new ArrayList<>();
+
+    @OnOpen
+    public void onOpen(Session session) {
+        sessions.add(session);
+    }
+
+    @OnClose
+    public void onClose(Session session) {
+        sessions.remove(session);
+        broadcast("Session " + session.getId() + " closed");
+    }
+
+    @OnError
+    public void onError(Session session, Throwable throwable) {
+        sessions.remove(session);
+        broadcast("Session " + session.getId() + " closed on error: " + throwable);
+    }
+
+    @OnMessage
+    public void onMessage(ChatMessageDTO message) {
+        broadcast(String.format("%s in message [%s] said: %s", message.getFrom(), message.getId(), message.getContent()));
+    }
+
+    private void broadcast(String message) {
+        sessions.forEach(s -> {
+            ChatMessageDTO chatMessageDTO = new ChatMessageDTO();
+            chatMessageDTO.setId(UUID.randomUUID().toString());
+            chatMessageDTO.setFrom("SuperCoolWebsocket");
+            chatMessageDTO.setContent(message);
+            s.getAsyncRemote().sendObject(chatMessageDTO, result -> {
+                if (result.getException() != null) {
+                    System.out.println("Unable to send message: " + result.getException());
+                }
+            });
+        });
+    }
+
+}

integration-tests/websockets/src/main/java/io/quarkus/websockets/ChatMessageDTO.java

@@ -0,0 +1,35 @@
+package io.quarkus.websockets;
+
+public class ChatMessageDTO {
+
+    private String id;
+
+    private String from;
+
+    private String content;
+
+    public String getId() {
+        return id;
+    }
+
+    public void setId(String id) {
+        this.id = id;
+    }
+
+    public String getFrom() {
+        return from;
+    }
+
+    public void setFrom(String from) {
+        this.from = from;
+    }
+
+    public String getContent() {
+        return content;
+    }
+
+    public void setContent(String content) {
+        this.content = content;
+    }
+
+}

integration-tests/websockets/src/main/java/io/quarkus/websockets/ChatMessageDecoder.java

@@ -0,0 +1,24 @@
+package io.quarkus.websockets;
+
+import jakarta.json.bind.Jsonb;
+import jakarta.json.bind.JsonbBuilder;
+import jakarta.websocket.DecodeException;
+import jakarta.websocket.Decoder;
+
+public class ChatMessageDecoder implements Decoder.Text<ChatMessageDTO> {
+
+    private final Jsonb jsonb = JsonbBuilder.create();
+
+    @Override
+    public ChatMessageDTO decode(String string) throws DecodeException {
+        return jsonb.fromJson(
+                string,
+                ChatMessageDTO.class);
+    }
+
+    @Override
+    public boolean willDecode(String string) {
+        return Boolean.TRUE;
+    }
+
+}

integration-tests/websockets/src/main/java/io/quarkus/websockets/ChatMessageEncoder.java

@@ -0,0 +1,17 @@
+package io.quarkus.websockets;
+
+import jakarta.json.bind.Jsonb;
+import jakarta.json.bind.JsonbBuilder;
+import jakarta.websocket.EncodeException;
+import jakarta.websocket.Encoder;
+
+public class ChatMessageEncoder implements Encoder.Text<ChatMessageDTO> {
+
+    private final Jsonb jsonb = JsonbBuilder.create();
+
+    @Override
+    public String encode(ChatMessageDTO t) throws EncodeException {
+        return jsonb.toJson(t);
+    }
+
+}

integration-tests/websockets/src/test/java/io/quarkus/websockets/ChatCodecIT.java

@@ -0,0 +1,7 @@
+package io.quarkus.websockets;
+
+import io.quarkus.test.junit.QuarkusIntegrationTest;
+
+@QuarkusIntegrationTest
+public class ChatCodecIT extends ChatCodecTest {
+}

integration-tests/websockets/src/test/java/io/quarkus/websockets/ChatCodecTest.java

@@ -0,0 +1,51 @@
+package io.quarkus.websockets;
+
+import java.net.URI;
+import java.util.UUID;
+import java.util.concurrent.LinkedBlockingDeque;
+import java.util.concurrent.TimeUnit;
+
+import jakarta.json.bind.JsonbBuilder;
+import jakarta.websocket.ClientEndpoint;
+import jakarta.websocket.ContainerProvider;
+import jakarta.websocket.OnMessage;
+import jakarta.websocket.Session;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import io.quarkus.test.common.http.TestHTTPResource;
+import io.quarkus.test.junit.QuarkusTest;
+
+@QuarkusTest
+public class ChatCodecTest {
+
+    private static final LinkedBlockingDeque<ChatMessageDTO> MESSAGES = new LinkedBlockingDeque<>();
+
+    @TestHTTPResource("/codec")
+    URI uri;
+
+    @Test
+    public void testWebsocketChat() throws Exception {
+        try (Session session = ContainerProvider.getWebSocketContainer().connectToServer(Client.class, uri)) {
+            ChatMessageDTO chatMessageDTO = new ChatMessageDTO();
+            chatMessageDTO.setId(UUID.randomUUID().toString());
+            chatMessageDTO.setFrom("SuperCoolProgrammer");
+            chatMessageDTO.setContent("Hello my young padawan!");
+            session.getAsyncRemote().sendText(JsonbBuilder.create().toJson(chatMessageDTO));
+            Assertions.assertEquals(String.format("%s in message [%s] said: %s", chatMessageDTO.getFrom(),
+                    chatMessageDTO.getId(), chatMessageDTO.getContent()), MESSAGES.poll(10, TimeUnit.SECONDS).getContent());
+        }
+    }
+
+    @ClientEndpoint(encoders = ChatMessageEncoder.class, decoders = ChatMessageDecoder.class)
+    public static class Client {
+
+        @OnMessage
+        void message(ChatMessageDTO msg) {
+            MESSAGES.add(msg);
+        }
+
+    }
+
+}