Use WebSocketFrameAggregator for Reactor and RxNetty

WebSocket frames are now aggregated through a Netty decoder so that
we always receive fully assembled messages by default capped at 64K.

Issue: SPR-14527

spring-web-reactive/src/main/java/org/springframework/web/reactive/socket/adapter/NettyWebSocketSessionSupport.java

@@ -16,19 +16,16 @@
 package org.springframework.web.reactive.socket.adapter;
 
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 
 import io.netty.buffer.ByteBuf;
 import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
-import io.netty.handler.codec.http.websocketx.CloseWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.PingWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.PongWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.WebSocketFrame;
-import reactor.core.publisher.Flux;
 
-import org.springframework.core.io.buffer.NettyDataBuffer;
+import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.NettyDataBufferFactory;
 import org.springframework.util.ObjectUtils;
 import org.springframework.web.reactive.socket.HandshakeInfo;
@@ -45,6 +42,12 @@ import org.springframework.web.reactive.socket.WebSocketSession;
  */
 public abstract class NettyWebSocketSessionSupport<T> extends AbstractWebSocketSession<T> {
 
+	/**
+	 * The default max size for aggregating inbound WebSocket frames.
+	 */
+	protected static final int DEFAULT_FRAME_MAX_SIZE = 64 * 1024;
+
+
 	private static final Map<Class<?>, WebSocketMessage.Type> MESSAGE_TYPES;
 
 	static {
@@ -66,27 +69,9 @@ public abstract class NettyWebSocketSessionSupport<T> extends AbstractWebSocketS
 		return (NettyDataBufferFactory) super.bufferFactory();
 	}
 
-
-	protected Flux<WebSocketMessage> toMessageFlux(Flux<WebSocketFrame> frameFlux) {
-		return frameFlux
-				.filter(frame -> !(frame instanceof CloseWebSocketFrame))
-				.window()
-				.concatMap(flux -> flux.takeUntil(WebSocketFrame::isFinalFragment).buffer())
-				.map(this::toMessage);
-	}
-
-	@SuppressWarnings("OptionalGetWithoutIsPresent")
-	private WebSocketMessage toMessage(List<WebSocketFrame> frames) {
-		Class<?> frameType = frames.get(0).getClass();
-		if (frames.size() == 1) {
-			NettyDataBuffer buffer = bufferFactory().wrap(frames.get(0).content());
-			return new WebSocketMessage(MESSAGE_TYPES.get(frameType), buffer);
-		}
-		return frames.stream()
-				.map(socketFrame -> bufferFactory().wrap(socketFrame.content()))
-				.reduce(NettyDataBuffer::write)
-				.map(buffer -> new WebSocketMessage(MESSAGE_TYPES.get(frameType), buffer))
-				.get();
+	protected WebSocketMessage toMessage(WebSocketFrame frame) {
+		DataBuffer payload = bufferFactory().wrap(frame.content());
+		return new WebSocketMessage(MESSAGE_TYPES.get(frame.getClass()), payload);
 	}
 
 	protected WebSocketFrame toFrame(WebSocketMessage message) {

spring-web-reactive/src/main/java/org/springframework/web/reactive/socket/adapter/ReactorNettyWebSocketSession.java

@@ -22,6 +22,8 @@ import reactor.core.publisher.Mono;
 import reactor.ipc.netty.NettyInbound;
 import reactor.ipc.netty.NettyOutbound;
 import reactor.ipc.netty.NettyPipeline;
+import reactor.ipc.netty.http.websocket.WebsocketInbound;
+import reactor.ipc.netty.http.websocket.WebsocketOutbound;
 
 import org.springframework.core.io.buffer.NettyDataBufferFactory;
 import org.springframework.web.reactive.socket.CloseStatus;
@@ -41,7 +43,7 @@ public class ReactorNettyWebSocketSession
 		extends NettyWebSocketSessionSupport<ReactorNettyWebSocketSession.WebSocketConnection> {
 
 
-	public ReactorNettyWebSocketSession(NettyInbound inbound, NettyOutbound outbound,
+	public ReactorNettyWebSocketSession(WebsocketInbound inbound, WebsocketOutbound outbound,
 			HandshakeInfo info, NettyDataBufferFactory bufferFactory) {
 
 		super(new WebSocketConnection(inbound, outbound), info, bufferFactory);
@@ -50,17 +52,19 @@ public class ReactorNettyWebSocketSession
 
 	@Override
 	public Flux<WebSocketMessage> receive() {
-		NettyInbound inbound = getDelegate().getInbound();
-		return toMessageFlux(inbound.receiveObject().cast(WebSocketFrame.class));
+		return getDelegate().getInbound()
+				.aggregateFrames(DEFAULT_FRAME_MAX_SIZE)
+				.receiveFrames()
+				.map(super::toMessage);
 	}
 
 	@Override
 	public Mono<Void> send(Publisher<WebSocketMessage> messages) {
-		Flux<WebSocketFrame> frameFlux = Flux.from(messages).map(this::toFrame);
-		NettyOutbound outbound = getDelegate().getOutbound();
-		return outbound.options(NettyPipeline.SendOptions::flushOnEach)
-		               .sendObject(frameFlux)
-		               .then();
+		Flux<WebSocketFrame> frames = Flux.from(messages).map(this::toFrame);
+		return getDelegate().getOutbound()
+				.options(NettyPipeline.SendOptions::flushOnEach)
+				.sendObject(frames)
+				.then();
 	}
 
 	@Override
@@ -77,21 +81,21 @@ public class ReactorNettyWebSocketSession
 	 */
 	public static class WebSocketConnection {
 
-		private final NettyInbound inbound;
+		private final WebsocketInbound inbound;
 
-		private final NettyOutbound outbound;
+		private final WebsocketOutbound outbound;
 
 
-		public WebSocketConnection(NettyInbound inbound, NettyOutbound outbound) {
+		public WebSocketConnection(WebsocketInbound inbound, WebsocketOutbound outbound) {
 			this.inbound = inbound;
 			this.outbound = outbound;
 		}
 
-		public NettyInbound getInbound() {
+		public WebsocketInbound getInbound() {
 			return this.inbound;
 		}
 
-		public NettyOutbound getOutbound() {
+		public WebsocketOutbound getOutbound() {
 			return this.outbound;
 		}
 	}

spring-web-reactive/src/main/java/org/springframework/web/reactive/socket/adapter/RxNettyWebSocketSession.java

@@ -16,7 +16,11 @@
 
 package org.springframework.web.reactive.socket.adapter;
 
+import io.netty.channel.Channel;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelPipeline;
 import io.netty.handler.codec.http.websocketx.WebSocketFrame;
+import io.netty.handler.codec.http.websocketx.WebSocketFrameAggregator;
 import io.reactivex.netty.protocol.http.ws.WebSocketConnection;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
@@ -25,6 +29,7 @@ import rx.Observable;
 import rx.RxReactiveStreams;
 
 import org.springframework.core.io.buffer.NettyDataBufferFactory;
+import org.springframework.util.Assert;
 import org.springframework.web.reactive.socket.CloseStatus;
 import org.springframework.web.reactive.socket.HandshakeInfo;
 import org.springframework.web.reactive.socket.WebSocketMessage;
@@ -39,6 +44,12 @@ import org.springframework.web.reactive.socket.WebSocketSession;
  */
 public class RxNettyWebSocketSession extends NettyWebSocketSessionSupport<WebSocketConnection> {
 
+	/**
+	 * The name of the {@link WebSocketFrameAggregator} inserted by
+	 * {@link #aggregateFrames(Channel, String)}.
+	 */
+	public static final String FRAME_AGGREGATOR_NAME = "websocket-frame-aggregator";
+
 
 	public RxNettyWebSocketSession(WebSocketConnection conn, HandshakeInfo info,
 			NettyDataBufferFactory factory) {
@@ -47,11 +58,30 @@ public class RxNettyWebSocketSession extends NettyWebSocketSessionSupport<WebSoc
 	}
 
 
+	/**
+	 * Inserts an {@link WebSocketFrameAggregator} after the
+	 * {@code WebSocketFrameDecoder} for receiving full messages.
+	 * @param channel the channel for the session
+	 * @param frameDecoderName the name of the WebSocketFrame decoder
+	 */
+	public RxNettyWebSocketSession aggregateFrames(Channel channel, String frameDecoderName) {
+		ChannelPipeline pipeline = channel.pipeline();
+		if (pipeline.context(FRAME_AGGREGATOR_NAME) != null) {
+			logger.trace("WebSocketFrameAggregator already registered.");
+			return this;
+		}
+		ChannelHandlerContext context = pipeline.context(frameDecoderName);
+		Assert.notNull(context, "WebSocketFrameDecoder not found: " + frameDecoderName);
+		WebSocketFrameAggregator aggregator = new WebSocketFrameAggregator(DEFAULT_FRAME_MAX_SIZE);
+		pipeline.addAfter(context.name(), FRAME_AGGREGATOR_NAME, aggregator);
+		return this;
+	}
+
+
 	@Override
 	public Flux<WebSocketMessage> receive() {
-		Observable<WebSocketFrame> observable = getDelegate().getInput();
-		Flux<WebSocketFrame> flux = Flux.from(RxReactiveStreams.toPublisher(observable));
-		return toMessageFlux(flux);
+		Observable<WebSocketMessage> observable = getDelegate().getInput().map(super::toMessage);
+		return Flux.from(RxReactiveStreams.toPublisher(observable));
 	}
 
 	@Override

spring-web-reactive/src/main/java/org/springframework/web/reactive/socket/client/ReactorNettyWebSocketClient.java

@@ -17,12 +17,10 @@ package org.springframework.web.reactive.socket.client;
 
 import java.net.URI;
 import java.util.Optional;
-import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 
 import io.netty.buffer.ByteBufAllocator;
 import reactor.core.publisher.Mono;
-import reactor.ipc.netty.NettyOutbound;
 import reactor.ipc.netty.http.client.HttpClient;
 import reactor.ipc.netty.http.client.HttpClientOptions;
 import reactor.ipc.netty.http.client.HttpClientRequest;
@@ -63,31 +61,27 @@ public class ReactorNettyWebSocketClient extends WebSocketClientSupport implemen
 	@Override
 	public Mono<Void> execute(URI url, HttpHeaders headers, WebSocketHandler handler) {
 
-		// TODO: https://github.com/reactor/reactor-netty/issues/19
-		AtomicReference<NettyOutbound> outboundRef = new AtomicReference<>();
-
 		String[] protocols = getSubProtocols(headers, handler);
 		// TODO: https://github.com/reactor/reactor-netty/issues/20
 
 		return this.httpClient
 				.get(url.toString(), request -> {
 					addRequestHeaders(request, headers);
-					NettyOutbound outbound = request.sendWebsocket();
-					outboundRef.set(outbound);
-					return outbound;
+					return request.sendWebsocket();
 				})
-				.then(in -> {
-					HttpHeaders responseHeaders = getResponseHeaders(in);
+				.then(response -> {
+					HttpHeaders responseHeaders = getResponseHeaders(response);
 					String protocol = responseHeaders.getFirst(SEC_WEBSOCKET_PROTOCOL);
 					HandshakeInfo info = new HandshakeInfo(url, responseHeaders, Mono.empty(),
 							Optional.ofNullable(protocol));
 
-					ByteBufAllocator allocator = in.channel().alloc();
+					ByteBufAllocator allocator = response.channel().alloc();
 					NettyDataBufferFactory factory = new NettyDataBufferFactory(allocator);
 
-					NettyOutbound out = outboundRef.get();
-					WebSocketSession session = new ReactorNettyWebSocketSession(in, out, info, factory);
-					return handler.handle(session);
+					return response.receiveWebsocket((in, out) -> {
+						WebSocketSession session = new ReactorNettyWebSocketSession(in, out, info, factory);
+						return handler.handle(session);
+					});
 				});
 	}
 

spring-web-reactive/src/main/java/org/springframework/web/reactive/socket/client/RxNettyWebSocketClient.java

@@ -29,6 +29,7 @@ import javax.net.ssl.SSLEngine;
 
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufAllocator;
+import io.reactivex.netty.protocol.http.HttpHandlerNames;
 import io.reactivex.netty.protocol.http.client.HttpClient;
 import io.reactivex.netty.protocol.http.ws.WebSocketConnection;
 import io.reactivex.netty.protocol.http.ws.client.WebSocketRequest;
@@ -119,7 +120,10 @@ public class RxNettyWebSocketClient extends WebSocketClientSupport implements We
 					NettyDataBufferFactory factory = new NettyDataBufferFactory(allocator);
 
 					WebSocketConnection conn = tuple.getT2();
-					WebSocketSession session = new RxNettyWebSocketSession(conn, info, factory);
+					RxNettyWebSocketSession session = new RxNettyWebSocketSession(conn, info, factory);
+					String name = HttpHandlerNames.WsClientDecoder.getName();
+					session.aggregateFrames(response.unsafeNettyChannel(), name);
+
 					return RxReactiveStreams.toObservable(handler.handle(session));
 				});
 	}

spring-web-reactive/src/main/java/org/springframework/web/reactive/socket/server/upgrade/RxNettyRequestUpgradeStrategy.java

@@ -18,6 +18,8 @@ package org.springframework.web.reactive.socket.server.upgrade;
 import java.security.Principal;
 import java.util.Optional;
 
+import io.reactivex.netty.protocol.http.HttpHandlerNames;
+import io.reactivex.netty.protocol.http.server.HttpServerResponse;
 import io.reactivex.netty.protocol.http.ws.server.WebSocketHandshaker;
 import reactor.core.publisher.Mono;
 import rx.RxReactiveStreams;
@@ -27,7 +29,6 @@ import org.springframework.http.server.reactive.RxNettyServerHttpResponse;
 import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.web.reactive.socket.HandshakeInfo;
 import org.springframework.web.reactive.socket.WebSocketHandler;
-import org.springframework.web.reactive.socket.WebSocketSession;
 import org.springframework.web.reactive.socket.adapter.RxNettyWebSocketSession;
 import org.springframework.web.reactive.socket.server.RequestUpgradeStrategy;
 import org.springframework.web.server.ServerWebExchange;
@@ -47,12 +48,16 @@ public class RxNettyRequestUpgradeStrategy implements RequestUpgradeStrategy {
 			Optional<String> subProtocol) {
 
 		RxNettyServerHttpResponse response = (RxNettyServerHttpResponse) exchange.getResponse();
+		HttpServerResponse<?> rxNettyResponse = response.getRxNettyResponse();
+
 		HandshakeInfo info = getHandshakeInfo(exchange, subProtocol);
 		NettyDataBufferFactory factory = (NettyDataBufferFactory) response.bufferFactory();
 
-		WebSocketHandshaker handshaker = response.getRxNettyResponse()
+		WebSocketHandshaker handshaker = rxNettyResponse
 				.acceptWebSocketUpgrade(conn -> {
-					WebSocketSession session = new RxNettyWebSocketSession(conn, info, factory);
+					RxNettyWebSocketSession session = new RxNettyWebSocketSession(conn, info, factory);
+					String name = HttpHandlerNames.WsServerDecoder.getName();
+					session.aggregateFrames(rxNettyResponse.unsafeNettyChannel(), name);
 					return RxReactiveStreams.toObservable(handler.handle(session));
 				});