diff --git a/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/WebsocketRoutingFilter.java b/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/WebsocketRoutingFilter.java index 8db39154..f15ef98f 100644 --- a/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/WebsocketRoutingFilter.java +++ b/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/WebsocketRoutingFilter.java @@ -203,6 +203,10 @@ public class WebsocketRoutingFilter implements GlobalFilter, Ordered { return client.execute(url, this.headers, new WebSocketHandler() { @Override public Mono handle(WebSocketSession proxySession) { + Mono serverClose = proxySession.closeStatus().filter(__ -> session.isOpen()) + .flatMap(session::close); + Mono proxyClose = session.closeStatus().filter(__ -> proxySession.isOpen()) + .flatMap(proxySession::close); // Use retain() for Reactor Netty Mono proxySessionSend = proxySession .send(session.receive().doOnNext(WebSocketMessage::retain)); @@ -210,6 +214,9 @@ public class WebsocketRoutingFilter implements GlobalFilter, Ordered { Mono serverSessionSend = session .send(proxySession.receive().doOnNext(WebSocketMessage::retain)); // .log("sessionSend", Level.FINE); + // Ensure closeStatus from one propagates to the other + Mono.when(serverClose, proxyClose).subscribe(); + // Complete when both sessions are done return Mono.zip(proxySessionSend, serverSessionSend).then(); } diff --git a/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java b/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java index 43426a52..88b7bc17 100644 --- a/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java +++ b/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java @@ -36,6 +36,7 @@ import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import reactor.core.publisher.MonoProcessor; import reactor.core.publisher.ReplayProcessor; +import reactor.core.publisher.Sinks; import org.springframework.beans.factory.annotation.Value; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; @@ -101,6 +102,8 @@ public class WebSocketIntegrationTests { private int gatewayPort; + private static final Sinks.One serverCloseStatusSink = Sinks.one(); + private static Mono doSend(WebSocketSession session, Publisher output) { return session.send(output); // workaround for suspected RxNetty WebSocket client issue @@ -241,13 +244,25 @@ public class WebSocketIntegrationTests { } @Test - public void sessionClosing() throws Exception { - this.client.execute(getUrl("/close"), session -> { + public void serverClosing() throws Exception { + AtomicReference> closeStatus = new AtomicReference<>(); + this.client.execute(getUrl("/server-close"), session -> { logger.debug("Starting.."); + closeStatus.set(session.closeStatus()); return session.receive().doOnNext(s -> logger.debug("inbound " + s)).then().doFinally(signalType -> { logger.debug("Completed with: " + signalType); }); }).block(Duration.ofMillis(5000)); + assertThat(closeStatus.get().block(Duration.ofMillis(5000))) + .isEqualTo(CloseStatus.create(4999, "server-close")); + } + + @Test + public void clientClosing() throws Exception { + this.client.execute(getUrl("/client-close"), session -> session.close(CloseStatus.create(4999, "client-close"))) + .block(Duration.ofMillis(5000)); + assertThat(serverCloseStatusSink.asMono().block(Duration.ofMillis(5000))) + .isEqualTo(CloseStatus.create(4999, "client-close")); } @Configuration(proxyBeanMethods = false) @@ -279,7 +294,8 @@ public class WebSocketIntegrationTests { map.put("/echoForHttp", new EchoWebSocketHandler()); map.put("/sub-protocol", new SubProtocolWebSocketHandler()); map.put("/custom-header", new CustomHeaderHandler()); - map.put("/close", new SessionClosingHandler()); + map.put("/server-close", new ServerClosingHandler()); + map.put("/client-close", new ClientClosingHandler()); SimpleUrlHandlerMapping mapping = new SimpleUrlHandlerMapping(); mapping.setUrlMap(map); @@ -334,11 +350,20 @@ public class WebSocketIntegrationTests { } - private static class SessionClosingHandler implements WebSocketHandler { + private static class ServerClosingHandler implements WebSocketHandler { @Override public Mono handle(WebSocketSession session) { - return Flux.never().mergeWith(session.close(CloseStatus.GOING_AWAY)).then(); + return Flux.never().mergeWith(session.close(CloseStatus.create(4999, "server-close"))).then(); + } + + } + + private static class ClientClosingHandler implements WebSocketHandler { + + @Override + public Mono handle(WebSocketSession session) { + return session.closeStatus().doOnNext(serverCloseStatusSink::tryEmitValue).then(); } }