diff --git a/framework-docs/modules/ROOT/pages/web/webflux-websocket.adoc b/framework-docs/modules/ROOT/pages/web/webflux-websocket.adoc index e08eed918259..10ccef66b39b 100644 --- a/framework-docs/modules/ROOT/pages/web/webflux-websocket.adoc +++ b/framework-docs/modules/ROOT/pages/web/webflux-websocket.adoc @@ -446,9 +446,21 @@ specify CORS settings by URL pattern. If both are specified, they are combined b === Client Spring WebFlux provides a `WebSocketClient` abstraction with implementations for -Reactor Netty, Tomcat, Jetty, and standard Java (that is, JSR-356). +Reactor Netty, Tomcat, Jetty, the Jakarta WebSocket API, and the JDK `HttpClient`. -NOTE: The Tomcat client is effectively an extension of the standard Java one with some extra +`JdkWebSocketClient` uses the WebSocket support in `java.net.http` without an additional +WebSocket client library. You can supply an `HttpClient` to configure connection settings. +Incoming text and binary fragments are aggregated into complete messages, with a default +limit of 64 KiB per message that you can change through `setMaxMessageSize`. + +NOTE: The JDK does not expose successful handshake response headers, so +`HandshakeInfo.getHeaders()` is empty for this client. The negotiated sub-protocol is +available through `HandshakeInfo.getSubProtocol()`. Incoming callbacks, including remote +close notifications, follow demand for `WebSocketSession.receive()`. Local closure and +cancellation of `execute` do not require receive demand. Close codes and reasons must +satisfy the constraints of `java.net.http.WebSocket.sendClose`. + +NOTE: The Tomcat client is effectively an extension of the Jakarta WebSocket client with some extra functionality in the `WebSocketSession` handling to take advantage of the Tomcat-specific API to suspend receiving messages for back pressure. diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/JdkWebSocketSession.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/JdkWebSocketSession.java new file mode 100644 index 000000000000..0424f2b0ae62 --- /dev/null +++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/JdkWebSocketSession.java @@ -0,0 +1,301 @@ +/* + * Copyright 2002-present the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.web.reactive.socket.adapter; + +import java.io.ByteArrayOutputStream; +import java.net.http.WebSocket; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.jspecify.annotations.Nullable; +import org.reactivestreams.Publisher; +import reactor.core.publisher.Flux; +import reactor.core.publisher.FluxSink; +import reactor.core.publisher.Mono; +import reactor.core.publisher.Sinks; + +import org.springframework.core.io.buffer.DataBuffer; +import org.springframework.core.io.buffer.DataBufferFactory; +import org.springframework.core.io.buffer.DataBufferLimitException; +import org.springframework.util.Assert; +import org.springframework.util.ObjectUtils; +import org.springframework.web.reactive.socket.CloseStatus; +import org.springframework.web.reactive.socket.HandshakeInfo; +import org.springframework.web.reactive.socket.WebSocketMessage; + +/** + * Adapt a JDK {@link WebSocket} to a reactive WebSocket session. + * + *

Incoming fragments are aggregated into complete messages. JDK receive + * callbacks, including the close callback, are requested according to demand + * for {@link #receive()}. + * + * @author Goutam Adwant + * @since 7.1 + */ +public class JdkWebSocketSession extends AbstractWebSocketSession implements WebSocket.Listener { + + private final int maxMessageSize; + + private final Flux receiveFlux; + + private final AtomicBoolean receiveSubscribed = new AtomicBoolean(); + + private final AtomicBoolean sendSubscribed = new AtomicBoolean(); + + private final AtomicBoolean awaitingCallback = new AtomicBoolean(); + + private final AtomicBoolean terminated = new AtomicBoolean(); + + private final Sinks.Empty termination = Sinks.empty(); + + private final Sinks.One closeStatus = Sinks.one(); + + private volatile @Nullable FluxSink receiveSink; + + private @Nullable ByteArrayOutputStream fragments; + + + /** + * Create a session for the given JDK WebSocket. + * @param webSocket the native WebSocket + * @param info handshake information + * @param bufferFactory the buffer factory to use + * @param maxMessageSize maximum incoming text or binary message size in bytes + */ + public JdkWebSocketSession(WebSocket webSocket, HandshakeInfo info, + DataBufferFactory bufferFactory, int maxMessageSize) { + + super(webSocket, ObjectUtils.getIdentityHexString(webSocket), info, bufferFactory); + Assert.isTrue(maxMessageSize > 0, "Max message size must be positive"); + this.maxMessageSize = maxMessageSize; + this.receiveFlux = Flux.create(sink -> { + if (!this.receiveSubscribed.compareAndSet(false, true)) { + sink.error(new IllegalStateException("receive() supports only one subscriber")); + return; + } + this.receiveSink = sink; + sink.onRequest(n -> requestCallback()); + sink.onCancel(this::clearFragments); + }, FluxSink.OverflowStrategy.ERROR).takeUntilOther(this.termination.asMono()) + .doOnDiscard(WebSocketMessage.class, WebSocketMessage::release); + } + + + @Override + public Flux receive() { + return this.receiveFlux; + } + + private void requestCallback() { + FluxSink sink = this.receiveSink; + if (!this.terminated.get() && sink != null && !sink.isCancelled() && sink.requestedFromDownstream() > 0 && + !getDelegate().isInputClosed() && this.awaitingCallback.compareAndSet(false, true)) { + getDelegate().request(1); + } + } + + @Override + public Mono send(Publisher messages) { + return Mono.defer(() -> { + if (!this.sendSubscribed.compareAndSet(false, true)) { + return Mono.error(new IllegalStateException("send() has already been called")); + } + return Flux.from(messages).concatMap(this::sendMessage) + .takeUntilOther(this.termination.asMono()) + .doOnDiscard(WebSocketMessage.class, WebSocketMessage::release).then(); + }); + } + + private Mono sendMessage(WebSocketMessage message) { + try { + DataBuffer payload = message.getPayload(); + if (message.getType() == WebSocketMessage.Type.TEXT) { + String text = payload.toString(StandardCharsets.UTF_8); + return Mono.fromFuture(getDelegate().sendText(text, true)).then(); + } + // The JDK may access this copy until its send future completes, + // independently of cancellation and pooled buffer ownership. + ByteBuffer bytes = ByteBuffer.allocate(payload.readableByteCount()); + payload.toByteBuffer(bytes); + return switch (message.getType()) { + case BINARY -> Mono.fromFuture(getDelegate().sendBinary(bytes, true)).then(); + case PING -> Mono.fromFuture(getDelegate().sendPing(bytes)).then(); + case PONG -> Mono.fromFuture(getDelegate().sendPong(bytes)).then(); + default -> Mono.error(new IllegalArgumentException("Unexpected message type: " + message.getType())); + }; + } + finally { + message.release(); + } + } + + @Override + public boolean isOpen() { + return !getDelegate().isInputClosed() && !getDelegate().isOutputClosed(); + } + + /** + * Send a close frame. The status code and reason must satisfy the constraints + * of {@link WebSocket#sendClose(int, String)}. + */ + @Override + public Mono close(CloseStatus status) { + return Mono.defer(() -> { + if (getDelegate().isOutputClosed()) { + return Mono.empty(); + } + String reason = status.getReason(); + return Mono.fromFuture(getDelegate().sendClose(status.getCode(), (reason != null ? reason : ""))) + .doOnSuccess(webSocket -> this.closeStatus.tryEmitValue(status)).then(); + }); + } + + @Override + public Mono closeStatus() { + return this.closeStatus.asMono(); + } + + /** + * Abort the native connection, for example on cancellation. + */ + public void abort() { + this.terminated.set(true); + clearFragments(); + getDelegate().abort(); + this.closeStatus.tryEmitEmpty(); + this.termination.tryEmitEmpty(); + } + + @Override + public void onOpen(WebSocket webSocket) { + // Demand is driven by the receive subscriber, not the default listener. + } + + @Override + public synchronized @Nullable CompletionStage onText(WebSocket webSocket, CharSequence data, boolean last) { + if (!isReceiving()) { + return null; + } + if (data.length() > remainingCapacity()) { + throw new DataBufferLimitException("Exceeded maximum WebSocket message size: " + this.maxMessageSize); + } + byte[] bytes = data.toString().getBytes(StandardCharsets.UTF_8); + handleFragment(WebSocketMessage.Type.TEXT, ByteBuffer.wrap(bytes), last); + return null; + } + + @Override + public @Nullable CompletionStage onBinary(WebSocket webSocket, ByteBuffer data, boolean last) { + handleFragment(WebSocketMessage.Type.BINARY, data, last); + return null; + } + + private int remainingCapacity() { + return this.maxMessageSize - (this.fragments != null ? this.fragments.size() : 0); + } + + private synchronized void handleFragment(WebSocketMessage.Type type, ByteBuffer data, boolean last) { + if (!isReceiving()) { + return; + } + if (data.remaining() > remainingCapacity()) { + throw new DataBufferLimitException("Exceeded maximum WebSocket message size: " + this.maxMessageSize); + } + byte[] bytes = new byte[data.remaining()]; + data.get(bytes); + if (!last || this.fragments != null) { + if (this.fragments == null) { + this.fragments = new ByteArrayOutputStream(); + } + this.fragments.writeBytes(bytes); + if (!last) { + this.awaitingCallback.set(false); + requestCallback(); + return; + } + bytes = this.fragments.toByteArray(); + this.fragments = null; + } + handleMessage(new WebSocketMessage(type, bufferFactory().wrap(bytes))); + } + + private boolean isReceiving() { + FluxSink sink = this.receiveSink; + return !this.terminated.get() && sink != null && !sink.isCancelled(); + } + + private void handleMessage(WebSocketMessage message) { + FluxSink sink = this.receiveSink; + if (sink != null && !sink.isCancelled()) { + sink.next(message); + } + else { + message.release(); + } + this.awaitingCallback.set(false); + requestCallback(); + } + + @Override + public @Nullable CompletionStage onPing(WebSocket webSocket, ByteBuffer message) { + handleControlMessage(WebSocketMessage.Type.PING, message); + return null; + } + + @Override + public @Nullable CompletionStage onPong(WebSocket webSocket, ByteBuffer message) { + handleControlMessage(WebSocketMessage.Type.PONG, message); + return null; + } + + private void handleControlMessage(WebSocketMessage.Type type, ByteBuffer data) { + byte[] bytes = new byte[data.remaining()]; + data.get(bytes); + handleMessage(new WebSocketMessage(type, bufferFactory().wrap(bytes))); + } + + @Override + public @Nullable CompletionStage onClose(WebSocket webSocket, int statusCode, String reason) { + try { + CloseStatus status = CloseStatus.create(statusCode, reason); + this.terminated.set(true); + clearFragments(); + this.closeStatus.tryEmitValue(status); + this.termination.tryEmitEmpty(); + } + catch (IllegalArgumentException ex) { + onError(webSocket, ex); + } + return null; + } + + @Override + public void onError(WebSocket webSocket, Throwable error) { + this.terminated.set(true); + clearFragments(); + this.closeStatus.tryEmitEmpty(); + this.termination.tryEmitError(error); + } + + private synchronized void clearFragments() { + this.fragments = null; + } + +} diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/client/JdkWebSocketClient.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/client/JdkWebSocketClient.java new file mode 100644 index 000000000000..667395b30788 --- /dev/null +++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/client/JdkWebSocketClient.java @@ -0,0 +1,225 @@ +/* + * Copyright 2002-present the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.web.reactive.socket.client; + +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.WebSocket; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; + +import org.jspecify.annotations.Nullable; +import reactor.core.publisher.Mono; +import reactor.core.publisher.Sinks; + +import org.springframework.core.io.buffer.DefaultDataBufferFactory; +import org.springframework.http.HttpHeaders; +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.WebSocketHandler; +import org.springframework.web.reactive.socket.adapter.JdkWebSocketSession; + +/** + * {@link WebSocketClient} implementation for the JDK {@link HttpClient}. + * + *

The JDK API does not expose handshake response headers. Therefore, + * {@link HandshakeInfo#getHeaders()} is empty, while the negotiated sub-protocol + * is available through {@link HandshakeInfo#getSubProtocol()}. + * + * @author Goutam Adwant + * @since 7.1 + */ +public class JdkWebSocketClient implements WebSocketClient { + + private final HttpClient httpClient; + + private int maxMessageSize = 64 * 1024; + + + /** + * Create a client with a default {@link HttpClient}. + */ + public JdkWebSocketClient() { + this(HttpClient.newHttpClient()); + } + + /** + * Create a client with the given {@link HttpClient}. + * @param httpClient the client to use + */ + public JdkWebSocketClient(HttpClient httpClient) { + Assert.notNull(httpClient, "HttpClient must not be null"); + this.httpClient = httpClient; + } + + + /** + * Return the configured {@link HttpClient}. + */ + public HttpClient getHttpClient() { + return this.httpClient; + } + + /** + * Set the maximum size in bytes of an incoming text or binary message, + * including all of its fragments. The default is 64 KiB. + */ + public void setMaxMessageSize(int maxMessageSize) { + Assert.isTrue(maxMessageSize > 0, "Max message size must be positive"); + this.maxMessageSize = maxMessageSize; + } + + /** + * Return the maximum incoming message size in bytes. + */ + public int getMaxMessageSize() { + return this.maxMessageSize; + } + + + @Override + public Mono execute(URI url, WebSocketHandler handler) { + return execute(url, new HttpHeaders(), handler); + } + + @Override + public Mono execute(URI url, HttpHeaders headers, WebSocketHandler handler) { + return Mono.defer(() -> { + WebSocket.Builder builder = this.httpClient.newWebSocketBuilder(); + headers.forEach((name, values) -> values.forEach(value -> builder.header(name, value))); + List protocols = handler.getSubProtocols(); + if (!protocols.isEmpty()) { + builder.subprotocols(protocols.get(0), protocols.subList(1, protocols.size()).toArray(String[]::new)); + } + JdkWebSocketListener listener = new JdkWebSocketListener(url, this.maxMessageSize); + return Mono.fromFuture(() -> { + CompletableFuture future = builder.buildAsync(url, listener); + future.thenAccept(listener::setWebSocket); + return future; + }) + .then(listener.sessionReady.asMono()) + .flatMap(session -> Mono.defer(() -> handler.handle(session)).then(Mono.defer(session::close))) + .takeUntilOther(listener.errors.asMono()) + .doFinally(signal -> listener.abort()); + }); + } + + + private static final class JdkWebSocketListener implements WebSocket.Listener { + + private final URI url; + + private final int maxMessageSize; + + private final Sinks.Empty errors = Sinks.empty(); + + private final Sinks.One sessionReady = Sinks.one(); + + private volatile @Nullable WebSocket webSocket; + + private volatile @Nullable JdkWebSocketSession session; + + private volatile boolean aborted; + + JdkWebSocketListener(URI url, int maxMessageSize) { + this.url = url; + this.maxMessageSize = maxMessageSize; + } + + JdkWebSocketSession getSession() { + JdkWebSocketSession session = this.session; + Assert.state(session != null, "WebSocket session must be initialized"); + return session; + } + + void setWebSocket(WebSocket webSocket) { + this.webSocket = webSocket; + if (this.aborted) { + webSocket.abort(); + } + } + + void abort() { + this.aborted = true; + JdkWebSocketSession session = this.session; + if (session != null) { + session.abort(); + } + else { + WebSocket webSocket = this.webSocket; + if (webSocket != null) { + webSocket.abort(); + } + } + } + + @Override + public void onOpen(WebSocket webSocket) { + String protocol = webSocket.getSubprotocol(); + HandshakeInfo info = new HandshakeInfo(this.url, HttpHeaders.EMPTY, Mono.empty(), + (!protocol.isEmpty() ? protocol : null)); + this.session = new JdkWebSocketSession(webSocket, info, + DefaultDataBufferFactory.sharedInstance, this.maxMessageSize); + if (this.aborted) { + this.session.abort(); + } + this.sessionReady.tryEmitValue(this.session); + } + + @Override + public @Nullable CompletionStage onText(WebSocket webSocket, CharSequence data, boolean last) { + return getSession().onText(webSocket, data, last); + } + + @Override + public @Nullable CompletionStage onBinary(WebSocket webSocket, ByteBuffer data, boolean last) { + return getSession().onBinary(webSocket, data, last); + } + + @Override + public @Nullable CompletionStage onPing(WebSocket webSocket, ByteBuffer message) { + return getSession().onPing(webSocket, message); + } + + @Override + public @Nullable CompletionStage onPong(WebSocket webSocket, ByteBuffer message) { + return getSession().onPong(webSocket, message); + } + + @Override + public @Nullable CompletionStage onClose(WebSocket webSocket, int statusCode, String reason) { + try { + CloseStatus.create(statusCode, reason); + return getSession().onClose(webSocket, statusCode, reason); + } + catch (IllegalArgumentException ex) { + onError(webSocket, ex); + return null; + } + } + + @Override + public void onError(WebSocket webSocket, Throwable error) { + getSession().onError(webSocket, error); + this.errors.tryEmitError(error); + } + } + +} diff --git a/spring-webflux/src/test/java/org/springframework/web/reactive/socket/AbstractReactiveWebSocketIntegrationTests.java b/spring-webflux/src/test/java/org/springframework/web/reactive/socket/AbstractReactiveWebSocketIntegrationTests.java index 44f8c8403efa..020550077003 100644 --- a/spring-webflux/src/test/java/org/springframework/web/reactive/socket/AbstractReactiveWebSocketIntegrationTests.java +++ b/spring-webflux/src/test/java/org/springframework/web/reactive/socket/AbstractReactiveWebSocketIntegrationTests.java @@ -45,6 +45,7 @@ import org.springframework.http.server.reactive.HttpHandler; import org.springframework.web.filter.reactive.ServerWebExchangeContextFilter; import org.springframework.web.reactive.DispatcherHandler; +import org.springframework.web.reactive.socket.client.JdkWebSocketClient; import org.springframework.web.reactive.socket.client.JettyWebSocketClient; import org.springframework.web.reactive.socket.client.ReactorNettyWebSocketClient; import org.springframework.web.reactive.socket.client.TomcatWebSocketClient; @@ -90,6 +91,7 @@ abstract class AbstractReactiveWebSocketIntegrationTests { static Stream arguments() throws IOException { List> clients = List.of( + named(JdkWebSocketClient.class.getSimpleName(), new JdkWebSocketClient()), named(TomcatWebSocketClient.class.getSimpleName(), new TomcatWebSocketClient()), named(JettyWebSocketClient.class.getSimpleName(), new JettyWebSocketClient()), named(ReactorNettyWebSocketClient.class.getSimpleName(), new ReactorNettyWebSocketClient()) diff --git a/spring-webflux/src/test/java/org/springframework/web/reactive/socket/WebSocketIntegrationTests.java b/spring-webflux/src/test/java/org/springframework/web/reactive/socket/WebSocketIntegrationTests.java index 6754116f0d6f..d9f8abfc5434 100644 --- a/spring-webflux/src/test/java/org/springframework/web/reactive/socket/WebSocketIntegrationTests.java +++ b/spring-webflux/src/test/java/org/springframework/web/reactive/socket/WebSocketIntegrationTests.java @@ -41,6 +41,7 @@ import org.springframework.web.reactive.HandlerMapping; import org.springframework.web.reactive.handler.SimpleUrlHandlerMapping; import org.springframework.web.reactive.socket.adapter.NettyWebSocketSessionSupport; +import org.springframework.web.reactive.socket.client.JdkWebSocketClient; import org.springframework.web.reactive.socket.client.JettyWebSocketClient; import org.springframework.web.reactive.socket.client.ReactorNettyWebSocketClient; import org.springframework.web.reactive.socket.client.TomcatWebSocketClient; @@ -101,6 +102,23 @@ private void testEcho() { assertThat(actualRef.get()).isEqualTo(input.collectList().block()); } + @ParameterizedWebSocketTest + void binaryEcho(WebSocketClient client, HttpServer server, Class serverConfigClass) throws Exception { + startServer(client, server, serverConfigClass); + byte[] input = new byte[] {0, 1, 2, 3, (byte) 255}; + AtomicReference actual = new AtomicReference<>(); + this.client.execute(getUrl("/echo"), session -> + session.send(Mono.just(session.binaryMessage(factory -> factory.wrap(input)))) + .thenMany(session.receive().take(1)) + .doOnNext(message -> { + byte[] bytes = new byte[message.getPayload().readableByteCount()]; + message.getPayload().read(bytes); + actual.set(bytes); + }) + .then()).block(TIMEOUT); + assertThat(actual.get()).containsExactly(input); + } + @ParameterizedWebSocketTest void subProtocol(WebSocketClient client, HttpServer server, Class serverConfigClass) throws Exception { startServer(client, server, serverConfigClass); @@ -129,8 +147,13 @@ public Mono handle(WebSocketSession session) { .block(TIMEOUT); HandshakeInfo info = infoRef.get(); - assertThat(info.getHeaders().getFirst("Upgrade")).isEqualToIgnoringCase("websocket"); - assertThat(info.getHeaders().getFirst("Sec-WebSocket-Protocol")).isEqualTo(protocol); + if (client instanceof JdkWebSocketClient) { + assertThat(info.getHeaders().isEmpty()).isTrue(); + } + else { + assertThat(info.getHeaders().getFirst("Upgrade")).isEqualToIgnoringCase("websocket"); + assertThat(info.getHeaders().getFirst("Sec-WebSocket-Protocol")).isEqualTo(protocol); + } assertThat(info.getSubProtocol()).as("Wrong protocol accepted").isEqualTo(protocol); assertThat(protocolRef.get()).as("Wrong protocol detected on the server side").isEqualTo(protocol); } @@ -192,7 +215,12 @@ void cookie(WebSocketClient client, HttpServer server, Class serverConfigClas }) .block(TIMEOUT); assertThat(receivedCookieRef.get()).isEqualTo("cookie"); - assertThat(cookie.get()).isEqualTo("project=spring"); + if (client instanceof JdkWebSocketClient) { + assertThat(cookie.get()).isNull(); + } + else { + assertThat(cookie.get()).isEqualTo("project=spring"); + } } @ParameterizedWebSocketTest @@ -221,6 +249,10 @@ void largePayload(WebSocketClient client, HttpServer server, Class serverConf } private WebSocketClient extendLimits(WebSocketClient client, int limit) { + if (client instanceof JdkWebSocketClient jdk) { + jdk.setMaxMessageSize(limit); + } + if (client instanceof ReactorNettyWebSocketClient netty) { client = new ReactorNettyWebSocketClient( netty.getHttpClient(), diff --git a/spring-webflux/src/test/java/org/springframework/web/reactive/socket/adapter/JdkWebSocketSessionTests.java b/spring-webflux/src/test/java/org/springframework/web/reactive/socket/adapter/JdkWebSocketSessionTests.java new file mode 100644 index 000000000000..77e80d0bd9f9 --- /dev/null +++ b/spring-webflux/src/test/java/org/springframework/web/reactive/socket/adapter/JdkWebSocketSessionTests.java @@ -0,0 +1,348 @@ +/* + * Copyright 2002-present the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.web.reactive.socket.adapter; + +import java.net.URI; +import java.net.http.WebSocket; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; + +import io.netty.buffer.Unpooled; +import io.netty.buffer.UnpooledByteBufAllocator; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.mockito.ArgumentCaptor; +import reactor.core.Disposable; +import reactor.core.publisher.Flux; +import reactor.core.publisher.FluxSink; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +import org.springframework.core.io.buffer.DataBufferLimitException; +import org.springframework.core.io.buffer.DefaultDataBufferFactory; +import org.springframework.core.io.buffer.NettyDataBuffer; +import org.springframework.core.io.buffer.NettyDataBufferFactory; +import org.springframework.http.HttpHeaders; +import org.springframework.web.reactive.socket.CloseStatus; +import org.springframework.web.reactive.socket.HandshakeInfo; +import org.springframework.web.reactive.socket.WebSocketMessage; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatCode; +import static org.assertj.core.api.Assertions.assertThatExceptionOfType; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@Timeout(10) +class JdkWebSocketSessionTests { + + private final WebSocket webSocket = mock(); + + private final JdkWebSocketSession session = new JdkWebSocketSession(this.webSocket, + new HandshakeInfo(URI.create("ws://example.org"), HttpHeaders.EMPTY, Mono.empty(), null), + DefaultDataBufferFactory.sharedInstance, 16); + + + @Test + void fragmentsConsumeOneMessageOfDemand() { + StepVerifier.create(this.session.receive(), 0) + .then(() -> verify(this.webSocket, never()).request(1)) + .thenRequest(1) + .then(() -> { + verify(this.webSocket).request(1); + this.session.onText(this.webSocket, "hel", false); + verify(this.webSocket, times(2)).request(1); + this.session.onText(this.webSocket, "lo", true); + }) + .assertNext(message -> assertThat(message.getPayloadAsText()).isEqualTo("hello")) + .then(() -> verify(this.webSocket, times(2)).request(1)) + .thenRequest(1) + .then(() -> verify(this.webSocket, times(3)).request(1)) + .thenCancel() + .verify(); + } + + @Test + void pingInterleavedWithFragments() { + StepVerifier.create(this.session.receive(), 1) + .then(() -> { + this.session.onText(this.webSocket, "hel", false); + this.session.onPing(this.webSocket, ByteBuffer.wrap(new byte[] {1})); + }) + .assertNext(message -> assertThat(message.getType()).isEqualTo(WebSocketMessage.Type.PING)) + .then(() -> verify(this.webSocket, times(2)).request(1)) + .thenRequest(1) + .then(() -> this.session.onText(this.webSocket, "lo", true)) + .assertNext(message -> assertThat(message.getPayloadAsText()).isEqualTo("hello")) + .thenCancel() + .verify(); + } + + @Test + void binaryInputIsCopiedBeforeCallbackReturns() { + ByteBuffer buffer = ByteBuffer.wrap(new byte[] {1, 2}); + StepVerifier.create(this.session.receive(), 1) + .then(() -> { + this.session.onBinary(this.webSocket, buffer, false); + buffer.put(0, (byte) 9); + this.session.onBinary(this.webSocket, ByteBuffer.wrap(new byte[] {3}), true); + }) + .assertNext(message -> { + byte[] bytes = new byte[3]; + message.getPayload().read(bytes); + assertThat(bytes).containsExactly(1, 2, 3); + }) + .thenCancel() + .verify(); + } + + @Test + void cancellationStopsDemandBetweenFragments() { + Disposable subscription = this.session.receive().subscribe(); + this.session.onText(this.webSocket, "first", false); + subscription.dispose(); + this.session.onText(this.webSocket, "last", true); + verify(this.webSocket, times(2)).request(1); + verify(this.webSocket, never()).abort(); + } + + @Test + void cancelledReceiveIgnoresOversizedOutstandingCallback() { + Disposable subscription = this.session.receive().subscribe(); + subscription.dispose(); + assertThatCode(() -> this.session.onText(this.webSocket, "x".repeat(17), true)) + .doesNotThrowAnyException(); + verify(this.webSocket).request(1); + verify(this.webSocket, never()).abort(); + } + + @Test + void messageLimitIncludesFragments() { + this.session.receive().subscribe(); + this.session.onText(this.webSocket, "1234567890", false); + assertThatExceptionOfType(DataBufferLimitException.class) + .isThrownBy(() -> this.session.onText(this.webSocket, "1234567", true)); + } + + @Test + void messageLimitCountsUtf8Bytes() { + this.session.receive().subscribe(); + assertThatExceptionOfType(DataBufferLimitException.class) + .isThrownBy(() -> this.session.onText(this.webSocket, "€€€€€€", true)); + } + + @Test + void closeCompletesReceiveAndPreservesStatus() { + StepVerifier.create(this.session.receive()) + .then(() -> this.session.onClose(this.webSocket, 1001, "going away")) + .verifyComplete(); + StepVerifier.create(this.session.closeStatus()) + .expectNext(CloseStatus.GOING_AWAY.withReason("going away")) + .verifyComplete(); + } + + @Test + void errorBeforeReceiveSubscriptionIsReplayed() { + IllegalStateException error = new IllegalStateException("connection failed"); + this.session.onError(this.webSocket, error); + StepVerifier.create(this.session.receive()).expectErrorSatisfies(actual -> assertThat(actual).isSameAs(error)) + .verify(); + } + + @Test + void sendIsSerialAndIncludesEmptyMessages() { + CompletableFuture first = new CompletableFuture<>(); + when(this.webSocket.sendText("", true)).thenReturn(first); + when(this.webSocket.sendText("next", true)).thenReturn(CompletableFuture.completedFuture(this.webSocket)); + StepVerifier.create(this.session.send(Flux.just(this.session.textMessage(""), this.session.textMessage("next")))) + .then(() -> { + verify(this.webSocket).sendText("", true); + verify(this.webSocket, never()).sendText("next", true); + first.complete(this.webSocket); + }) + .verifyComplete(); + verify(this.webSocket).sendText("next", true); + } + + @Test + void pooledOutputIsCopiedAndReleasedEvenWhenSendIsCancelled() { + NettyDataBuffer payload = new NettyDataBufferFactory(UnpooledByteBufAllocator.DEFAULT) + .wrap(Unpooled.wrappedBuffer("hello".getBytes(StandardCharsets.UTF_8))); + CompletableFuture future = new CompletableFuture<>(); + when(this.webSocket.sendBinary(any(), eq(true))).thenReturn(future); + Disposable subscription = this.session.send(Mono.just(new WebSocketMessage(WebSocketMessage.Type.BINARY, payload))) + .subscribe(); + ArgumentCaptor bytes = ArgumentCaptor.forClass(ByteBuffer.class); + verify(this.webSocket).sendBinary(bytes.capture(), eq(true)); + assertThat(payload.getNativeBuffer().refCnt()).isZero(); + subscription.dispose(); + assertThat(StandardCharsets.UTF_8.decode(bytes.getValue()).toString()).isEqualTo("hello"); + } + + @Test + void sendCompositeBinaryPayload() { + NettyDataBuffer payload = new NettyDataBufferFactory(UnpooledByteBufAllocator.DEFAULT).wrap( + Unpooled.wrappedBuffer(Unpooled.wrappedBuffer(new byte[] {1, 2}), + Unpooled.wrappedBuffer(new byte[] {3, 4}))); + when(this.webSocket.sendBinary(any(), eq(true))) + .thenReturn(CompletableFuture.completedFuture(this.webSocket)); + this.session.send(Mono.just(new WebSocketMessage(WebSocketMessage.Type.BINARY, payload))).block(); + ArgumentCaptor bytes = ArgumentCaptor.forClass(ByteBuffer.class); + verify(this.webSocket).sendBinary(bytes.capture(), eq(true)); + assertThat(bytes.getValue().array()).containsExactly(1, 2, 3, 4); + assertThat(payload.getNativeBuffer().refCnt()).isZero(); + } + + @Test + void sendEmptyBinaryPayload() { + when(this.webSocket.sendBinary(any(), eq(true))) + .thenReturn(CompletableFuture.completedFuture(this.webSocket)); + this.session.send(Mono.just(this.session.binaryMessage(factory -> factory.wrap(new byte[0])))).block(); + ArgumentCaptor bytes = ArgumentCaptor.forClass(ByteBuffer.class); + verify(this.webSocket).sendBinary(bytes.capture(), eq(true)); + assertThat(bytes.getValue().remaining()).isZero(); + } + + @Test + void sendPingAndPongPayloads() { + when(this.webSocket.sendPing(any())).thenReturn(CompletableFuture.completedFuture(this.webSocket)); + when(this.webSocket.sendPong(any())).thenReturn(CompletableFuture.completedFuture(this.webSocket)); + this.session.send(Flux.just( + this.session.pingMessage(factory -> factory.wrap(new byte[] {1})), + this.session.pongMessage(factory -> factory.wrap(new byte[] {2})))).block(); + ArgumentCaptor ping = ArgumentCaptor.forClass(ByteBuffer.class); + ArgumentCaptor pong = ArgumentCaptor.forClass(ByteBuffer.class); + verify(this.webSocket).sendPing(ping.capture()); + verify(this.webSocket).sendPong(pong.capture()); + assertThat(ping.getValue().array()).containsExactly(1); + assertThat(pong.getValue().array()).containsExactly(2); + } + + @Test + void failedSendIsPropagated() { + IllegalStateException failure = new IllegalStateException("write failed"); + when(this.webSocket.sendText("hello", true)).thenReturn(CompletableFuture.failedFuture(failure)); + StepVerifier.create(this.session.send(Mono.just(this.session.textMessage("hello")))) + .expectErrorSatisfies(actual -> assertThat(actual).isSameAs(failure)).verify(); + } + + @Test + void sendDiscardsLateMessagesAfterCancellation() { + AtomicReference> source = new AtomicReference<>(); + Disposable subscription = this.session.send(Flux.create(source::set)).subscribe(); + subscription.dispose(); + NettyDataBuffer payload = new NettyDataBufferFactory(UnpooledByteBufAllocator.DEFAULT) + .wrap(Unpooled.wrappedBuffer(new byte[] {1})); + source.get().next(new WebSocketMessage(WebSocketMessage.Type.BINARY, payload)); + assertThat(payload.getNativeBuffer().refCnt()).isZero(); + } + + @Test + void sendCancellationReleasesQueuedPayload() { + NettyDataBufferFactory factory = new NettyDataBufferFactory(UnpooledByteBufAllocator.DEFAULT); + NettyDataBuffer first = factory.wrap(Unpooled.wrappedBuffer(new byte[] {1})); + NettyDataBuffer second = factory.wrap(Unpooled.wrappedBuffer(new byte[] {2})); + when(this.webSocket.sendBinary(any(), eq(true))).thenReturn(new CompletableFuture<>()); + Disposable subscription = this.session.send(Flux.create(sink -> { + sink.next(new WebSocketMessage(WebSocketMessage.Type.BINARY, first)); + sink.next(new WebSocketMessage(WebSocketMessage.Type.BINARY, second)); + })).subscribe(); + subscription.dispose(); + assertThat(first.getNativeBuffer().refCnt()).isZero(); + assertThat(second.getNativeBuffer().refCnt()).isZero(); + } + + @Test + @Timeout(30) + void concurrentSendCancellationReleasesPayload() throws Exception { + NettyDataBufferFactory factory = new NettyDataBufferFactory(UnpooledByteBufAllocator.DEFAULT); + when(this.webSocket.sendBinary(any(), eq(true))).thenAnswer(invocation -> new CompletableFuture<>()); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + for (int i = 0; i < 1000; i++) { + JdkWebSocketSession session = new JdkWebSocketSession(this.webSocket, + this.session.getHandshakeInfo(), factory, 16); + AtomicReference> source = new AtomicReference<>(); + Disposable subscription = session.send(Flux.create(source::set)).subscribe(); + NettyDataBuffer payload = factory.wrap(Unpooled.wrappedBuffer(new byte[] {1})); + CyclicBarrier barrier = new CyclicBarrier(2); + Future sending = executor.submit(() -> { + barrier.await(); + source.get().next(new WebSocketMessage(WebSocketMessage.Type.BINARY, payload)); + return null; + }); + Future cancelling = executor.submit(() -> { + barrier.await(); + subscription.dispose(); + return null; + }); + sending.get(); + cancelling.get(); + assertThat(payload.getNativeBuffer().refCnt()).as("iteration %s", i).isZero(); + } + } + finally { + executor.shutdownNow(); + } + } + + @Test + void closeDoesNotRequireReceiveDemandAndNormalizesReason() { + when(this.webSocket.sendClose(1000, "")).thenReturn(CompletableFuture.completedFuture(this.webSocket)); + this.session.close().block(); + verify(this.webSocket).sendClose(1000, ""); + verify(this.webSocket, never()).request(1); + StepVerifier.create(this.session.closeStatus()).expectNext(CloseStatus.NORMAL).verifyComplete(); + } + + @Test + void remoteCloseCancelsPendingSendSource() { + AtomicBoolean cancelled = new AtomicBoolean(); + StepVerifier.create(this.session.send(Flux.never().doOnCancel(() -> cancelled.set(true)))) + .then(() -> this.session.onClose(this.webSocket, 1000, "")) + .verifyComplete(); + assertThat(cancelled).isTrue(); + } + + @Test + void nativeErrorTerminatesPendingSend() { + IllegalStateException error = new IllegalStateException("connection failed"); + StepVerifier.create(this.session.send(Flux.never())) + .then(() -> this.session.onError(this.webSocket, error)) + .expectErrorSatisfies(actual -> assertThat(actual).isSameAs(error)).verify(); + } + + @Test + void unsupportedIncomingCloseStatusTerminatesReceive() { + StepVerifier.create(this.session.receive()) + .then(() -> this.session.onClose(this.webSocket, 65535, "")) + .verifyError(IllegalArgumentException.class); + } + +} diff --git a/spring-webflux/src/test/java/org/springframework/web/reactive/socket/client/JdkWebSocketClientTests.java b/spring-webflux/src/test/java/org/springframework/web/reactive/socket/client/JdkWebSocketClientTests.java new file mode 100644 index 000000000000..1b13deaf38e4 --- /dev/null +++ b/spring-webflux/src/test/java/org/springframework/web/reactive/socket/client/JdkWebSocketClientTests.java @@ -0,0 +1,218 @@ +/* + * Copyright 2002-present the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.web.reactive.socket.client; + +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.WebSocket; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicReference; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import reactor.core.Disposable; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +import org.springframework.http.HttpHeaders; +import org.springframework.web.reactive.socket.WebSocketHandler; +import org.springframework.web.reactive.socket.WebSocketSession; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@Timeout(10) +class JdkWebSocketClientTests { + + private static final URI URL = URI.create("ws://example.org"); + + private final HttpClient httpClient = mock(); + + private final WebSocket.Builder builder = mock(); + + private final WebSocket webSocket = mock(); + + private final JdkWebSocketClient client = new JdkWebSocketClient(this.httpClient); + + private final AtomicReference listener = new AtomicReference<>(); + + + @BeforeEach + void setUp() { + when(this.httpClient.newWebSocketBuilder()).thenReturn(this.builder); + when(this.webSocket.getSubprotocol()).thenReturn(""); + when(this.webSocket.sendClose(1000, "")).thenReturn(CompletableFuture.completedFuture(this.webSocket)); + when(this.builder.buildAsync(eq(URL), any())).thenAnswer(invocation -> { + WebSocket.Listener listener = invocation.getArgument(1); + this.listener.set(listener); + listener.onOpen(this.webSocket); + return CompletableFuture.completedFuture(this.webSocket); + }); + } + + + @Test + void connectIsDeferredAndRepeatedPerSubscription() { + Mono execute = this.client.execute(URL, session -> Mono.empty()); + verify(this.httpClient, never()).newWebSocketBuilder(); + execute.block(); + execute.block(); + verify(this.httpClient, times(2)).newWebSocketBuilder(); + verify(this.webSocket, times(2)).sendClose(1000, ""); + verify(this.webSocket, times(2)).abort(); + } + + @Test + void requestHeadersAndSubprotocols() { + HttpHeaders headers = new HttpHeaders(); + headers.add("X-Test", "one"); + headers.add("X-Test", "two"); + when(this.webSocket.getSubprotocol()).thenReturn("echo"); + this.client.execute(URL, headers, new WebSocketHandler() { + @Override + public List getSubProtocols() { + return List.of("echo", "other"); + } + @Override + public Mono handle(WebSocketSession session) { + assertThat(session.getHandshakeInfo().getHeaders().isEmpty()).isTrue(); + assertThat(session.getHandshakeInfo().getSubProtocol()).isEqualTo("echo"); + return Mono.empty(); + } + }).block(); + verify(this.builder).header("X-Test", "one"); + verify(this.builder).header("X-Test", "two"); + verify(this.builder).subprotocols("echo", "other"); + } + + @Test + void subscriberContextReachesHandler() { + this.client.execute(URL, session -> Mono.deferContextual(context -> { + assertThat(context.get("key").toString()).isEqualTo("value"); + return Mono.empty(); + })).contextWrite(context -> context.put("key", "value")).block(); + } + + @Test + void cancellationCancelsHandshake() { + CompletableFuture handshake = new CompletableFuture<>(); + when(this.builder.buildAsync(eq(URL), any())).thenReturn(handshake); + Disposable subscription = this.client.execute(URL, session -> Mono.never()).subscribe(); + subscription.dispose(); + assertThat(handshake).isCancelled(); + } + + @Test + void handshakeFutureCanCompleteBeforeOnOpen() { + when(this.builder.buildAsync(eq(URL), any())).thenAnswer(invocation -> { + this.listener.set(invocation.getArgument(1)); + return CompletableFuture.completedFuture(this.webSocket); + }); + StepVerifier.create(this.client.execute(URL, session -> Mono.empty())) + .then(() -> this.listener.get().onOpen(this.webSocket)) + .verifyComplete(); + } + + @Test + void cancellationWhileWaitingForOnOpenAbortsConnection() { + when(this.builder.buildAsync(eq(URL), any())).thenAnswer(invocation -> { + this.listener.set(invocation.getArgument(1)); + return CompletableFuture.completedFuture(this.webSocket); + }); + WebSocketHandler handler = mock(); + Disposable subscription = this.client.execute(URL, handler).subscribe(); + subscription.dispose(); + verify(this.webSocket).abort(); + this.listener.get().onOpen(this.webSocket); + verify(handler, never()).handle(any()); + } + + @Test + void connectionOpenedAfterCancellationIsAborted() { + CompletableFuture handshake = new CompletableFuture<>(); + when(this.builder.buildAsync(eq(URL), any())).thenAnswer(invocation -> { + this.listener.set(invocation.getArgument(1)); + return handshake; + }); + Disposable subscription = this.client.execute(URL, session -> Mono.never()).subscribe(); + subscription.dispose(); + this.listener.get().onOpen(this.webSocket); + verify(this.webSocket).abort(); + } + + @Test + void cancellationAbortsConnectedSession() { + Disposable subscription = this.client.execute(URL, session -> Mono.never()).subscribe(); + subscription.dispose(); + verify(this.webSocket).abort(); + } + + @Test + void nativeErrorTerminatesHandlerWithoutReceiveSubscription() { + IllegalStateException error = new IllegalStateException("connection failed"); + StepVerifier.create(this.client.execute(URL, session -> Mono.never())) + .then(() -> this.listener.get().onError(this.webSocket, error)) + .expectErrorSatisfies(actual -> assertThat(actual).isSameAs(error)).verify(); + verify(this.webSocket).abort(); + } + + @Test + void unsupportedCloseStatusTerminatesHandlerWithoutReceiveSubscription() { + StepVerifier.create(this.client.execute(URL, session -> Mono.never())) + .then(() -> this.listener.get().onClose(this.webSocket, 65535, "")) + .verifyError(IllegalArgumentException.class); + verify(this.webSocket).abort(); + } + + @Test + void synchronousHandlerFailureAbortsSession() { + IllegalStateException error = new IllegalStateException("handler failed"); + StepVerifier.create(this.client.execute(URL, session -> { throw error; })) + .expectErrorSatisfies(actual -> assertThat(actual).isSameAs(error)).verify(); + verify(this.webSocket).abort(); + } + + @Test + void closeFailureIsPropagatedWithoutSuccessfulCloseStatus() { + IllegalArgumentException error = new IllegalArgumentException("close rejected"); + when(this.webSocket.sendClose(1000, "")).thenReturn(CompletableFuture.failedFuture(error)); + AtomicReference sessionRef = new AtomicReference<>(); + StepVerifier.create(this.client.execute(URL, session -> { + sessionRef.set(session); + return Mono.empty(); + })).expectErrorSatisfies(actual -> assertThat(actual).isSameAs(error)).verify(); + verify(this.webSocket).abort(); + StepVerifier.create(sessionRef.get().closeStatus()).verifyComplete(); + } + + @Test + void handshakeFailureIsPropagated() { + IllegalStateException error = new IllegalStateException("handshake failed"); + when(this.builder.buildAsync(eq(URL), any())).thenReturn(CompletableFuture.failedFuture(error)); + StepVerifier.create(this.client.execute(URL, session -> Mono.empty())) + .expectErrorSatisfies(actual -> assertThat(actual).isSameAs(error)).verify(); + } + +}