Add support for WebSocketFrame parameters in OnMessage (#11673)

Fixes #11506

Co-authored-by: Sergio del Amo <[email protected]>

Author: yawkat

http-netty/src/main/java/io/micronaut/http/netty/websocket/AbstractNettyWebSocketHandler.java

@@ -319,30 +319,36 @@ protected void handleWebSocketFrame(ChannelHandlerContext ctx, WebSocketFrame ms
                         CloseReason.UNSUPPORTED_DATA
                 );
             } else {
-                ByteBuf msgContent = msg.content().retain();
-                if (!msg.isFinalFragment()) {
-                    frameBuffer.updateAndGet((buffer) -> {
-                        if (buffer == null) {
-                            buffer = ctx.alloc().compositeBuffer();
-                        }
-                        buffer.addComponent(true, msgContent);
-                        return buffer;
-                    });
-                    return;
-                }
+                Argument<?> bodyArgument = this.getBodyArgument();
+                Object data;
 
-                ByteBuf content;
-                CompositeByteBuf buffer = frameBuffer.getAndSet(null);
-                if (buffer == null) {
-                    content = msgContent;
+                if (WebSocketFrame.class.isAssignableFrom(bodyArgument.getType())) {
+                    data = msg.retain();
                 } else {
-                    buffer.addComponent(true, msgContent);
-                    content = buffer;
-                }
+                    ByteBuf msgContent = msg.content().retain();
+                    if (!msg.isFinalFragment()) {
+                        frameBuffer.updateAndGet((buffer) -> {
+                            if (buffer == null) {
+                                buffer = ctx.alloc().compositeBuffer();
+                            }
+                            buffer.addComponent(true, msgContent);
+                            return buffer;
+                        });
+                        return;
+                    }
 
-                Argument<?> bodyArgument = this.getBodyArgument();
-                Object data = conversionService.convert(content, ByteBuf.class, bodyArgument).orElse(null);
-                content.release();
+                    ByteBuf content;
+                    CompositeByteBuf buffer = frameBuffer.getAndSet(null);
+                    if (buffer == null) {
+                        content = msgContent;
+                    } else {
+                        buffer.addComponent(true, msgContent);
+                        content = buffer;
+                    }
+
+                    data = conversionService.convert(content, ByteBuf.class, bodyArgument).orElse(null);
+                    content.release();
+                }
 
                 if (data == null) {
                     MediaType mediaType;

http-server-netty/src/test/groovy/io/micronaut/http/server/netty/websocket/WebSocketSpec.groovy

@@ -2,7 +2,6 @@ package io.micronaut.http.server.netty.websocket
 
 import io.micronaut.context.ApplicationContext
 import io.micronaut.context.annotation.Requires
-import io.micronaut.core.annotation.NonNull
 import io.micronaut.http.HttpRequest
 import io.micronaut.http.HttpResponse
 import io.micronaut.http.annotation.Filter
@@ -11,26 +10,29 @@ import io.micronaut.http.filter.HttpFilter
 import io.micronaut.http.server.netty.EmbeddedTestUtil
 import io.micronaut.http.server.netty.NettyHttpServer
 import io.micronaut.runtime.server.EmbeddedServer
+import io.micronaut.websocket.WebSocketClient
 import io.micronaut.websocket.WebSocketSession
+import io.micronaut.websocket.annotation.ClientWebSocket
 import io.micronaut.websocket.annotation.OnMessage
 import io.micronaut.websocket.annotation.ServerWebSocket
-import io.netty.channel.ChannelHandlerContext
-import io.netty.channel.SimpleChannelInboundHandler
 import io.netty.channel.embedded.EmbeddedChannel
 import io.netty.handler.codec.http.DefaultHttpHeaders
-import io.netty.handler.codec.http.FullHttpResponse
 import io.netty.handler.codec.http.HttpClientCodec
 import io.netty.handler.codec.http.HttpObjectAggregator
 import io.netty.handler.codec.http.websocketx.TextWebSocketFrame
 import io.netty.handler.codec.http.websocketx.WebSocketClientHandshakerFactory
+import io.netty.handler.codec.http.websocketx.WebSocketFrame
 import io.netty.handler.codec.http.websocketx.WebSocketVersion
 import jakarta.inject.Singleton
 import org.reactivestreams.Publisher
 import reactor.core.publisher.Flux
+import reactor.core.publisher.Mono
 import reactor.core.publisher.Sinks
 import spock.lang.Issue
 import spock.lang.Specification
 
+import java.util.concurrent.CompletableFuture
+
 class WebSocketSpec extends Specification {
     @Issue('https://github.com/micronaut-projects/micronaut-core/issues/7920')
     def 'race condition with channel close from http filter'() {
@@ -99,4 +101,50 @@ class WebSocketSpec extends Specification {
             return Flux.concat(delay.asMono(), Flux.defer(() -> chain.proceed(request)))
         }
     }
+
+    @Issue('https://github.com/micronaut-projects/micronaut-core/issues/11506')
+    def 'netty WebSocketFrame'() {
+        given:
+        ApplicationContext ctx = ApplicationContext.run([
+                'spec.name': 'WebSocketSpec2',
+        ])
+        def embeddedServer = ctx.getBean(EmbeddedServer)
+        embeddedServer.start()
+        def client = ctx.createBean(WebSocketClient, embeddedServer.URI)
+
+        when:
+        def sock = Mono.from(client.connect(ClientSocket, "/ServerSocket")).block()
+        sock.send("foo")
+        then:
+        sock.reply.get() == "reply: foo"
+
+        cleanup:
+        sock.close()
+        client.close()
+        embeddedServer.close()
+        ctx.close()
+    }
+
+    @ServerWebSocket('/ServerSocket')
+    @Requires(property = 'spec.name', value = 'WebSocketSpec2')
+    static class ServerSocket {
+        @OnMessage
+        def onMessage(WebSocketFrame message, WebSocketSession session) {
+            def text = ((TextWebSocketFrame) message).text()
+            message.release()
+            return session.send('reply: ' + text)
+        }
+    }
+
+    @ClientWebSocket
+    static abstract class ClientSocket implements AutoCloseable {
+        def reply = new CompletableFuture<String>()
+
+        abstract void send(String message);
+
+        @OnMessage
+        public void onMessage(String message) {
+            reply.complete(message)
+        }
+    }
 }