diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/WebsocketRoutingFilter.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/WebsocketRoutingFilter.java index f90d58b1..35a00e15 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/WebsocketRoutingFilter.java +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/WebsocketRoutingFilter.java @@ -1,9 +1,12 @@ package org.springframework.cloud.gateway.filter; import java.net.URI; +import java.util.Collections; +import java.util.List; import java.util.logging.Level; import org.springframework.core.Ordered; +import org.springframework.http.HttpHeaders; import org.springframework.web.reactive.socket.WebSocketHandler; import org.springframework.web.reactive.socket.WebSocketMessage; import org.springframework.web.reactive.socket.WebSocketSession; @@ -21,6 +24,8 @@ import reactor.core.publisher.Mono; * @author Spencer Gibb */ public class WebsocketRoutingFilter implements GlobalFilter, Ordered { + public static final String SEC_WEBSOCKET_PROTOCOL = "Sec-WebSocket-Protocol"; + private final WebSocketClient webSocketClient; private final WebSocketService webSocketService; @@ -48,30 +53,54 @@ public class WebsocketRoutingFilter implements GlobalFilter, Ordered { return chain.filter(exchange); } - return this.webSocketService.handleRequest(exchange, new ProxyWebSocketHandler(requestUrl, this.webSocketClient)); + + return this.webSocketService.handleRequest(exchange, + new ProxyWebSocketHandler(requestUrl, this.webSocketClient, exchange.getRequest().getHeaders())); } private static class ProxyWebSocketHandler implements WebSocketHandler { private final WebSocketClient client; private final URI url; + private final HttpHeaders headers; + private final List subProtocols; - public ProxyWebSocketHandler(URI url, WebSocketClient client) { + public ProxyWebSocketHandler(URI url, WebSocketClient client, HttpHeaders headers) { this.client = client; this.url = url; + this.headers = headers; + List protocols = headers.get(SEC_WEBSOCKET_PROTOCOL); + if (protocols != null) { + this.subProtocols = protocols; + } else { + this.subProtocols = Collections.emptyList(); + } + } + + @Override + public List getSubProtocols() { + return this.subProtocols; } @Override public Mono handle(WebSocketSession session) { - return client.execute(url, proxySession -> { - // Use retain() for Reactor Netty - Mono proxySessionSend = proxySession - .send(session.receive().doOnNext(WebSocketMessage::retain)) - .log("proxySessionSend", Level.FINE); - Mono serverSessionSend = session - .send(proxySession.receive().doOnNext(WebSocketMessage::retain)) - .log("sessionSend", Level.FINE); - return Mono.when(proxySessionSend, serverSessionSend).then(); + return client.execute(url, this.headers, new WebSocketHandler() { + @Override + public Mono handle(WebSocketSession proxySession) { + // Use retain() for Reactor Netty + Mono proxySessionSend = proxySession + .send(session.receive().doOnNext(WebSocketMessage::retain)) + .log("proxySessionSend", Level.FINE); + Mono serverSessionSend = session + .send(proxySession.receive().doOnNext(WebSocketMessage::retain)) + .log("sessionSend", Level.FINE); + return Mono.when(proxySessionSend, serverSessionSend).then(); + } + + @Override + public List getSubProtocols() { + return ProxyWebSocketHandler.this.subProtocols; + } }); } } diff --git a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java index 0e653697..c501b4b8 100644 --- a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java +++ b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java @@ -45,6 +45,7 @@ import org.springframework.context.annotation.Configuration; import org.springframework.core.env.ConfigurableEnvironment; import org.springframework.http.HttpHeaders; import org.springframework.http.server.reactive.HttpHandler; +import org.springframework.util.StringUtils; import org.springframework.web.reactive.DispatcherHandler; import org.springframework.web.reactive.HandlerMapping; import org.springframework.web.reactive.handler.SimpleUrlHandlerMapping; @@ -260,6 +261,9 @@ public class WebSocketIntegrationTests { @Override public Mono handle(WebSocketSession session) { String protocol = session.getHandshakeInfo().getSubProtocol(); + if (!StringUtils.hasText(protocol)) { + return Mono.error(new IllegalStateException("Missing protocol")); + } WebSocketMessage message = session.textMessage(protocol); return doSend(session, Mono.just(message)); } @@ -270,6 +274,9 @@ public class WebSocketIntegrationTests { @Override public Mono handle(WebSocketSession session) { HttpHeaders headers = session.getHandshakeInfo().getHeaders(); + if (!headers.containsKey("my-header")) { + return Mono.error(new IllegalStateException("Missing my-header")); + } String payload = "my-header:" + headers.getFirst("my-header"); WebSocketMessage message = session.textMessage(payload); return doSend(session, Mono.just(message));