Sub-protocol negotiation for reactive WebSocket support

Issue: SPR-14527

spring-web-reactive/src/test/java/org/springframework/web/reactive/socket/server/WebSocketIntegrationTests.java

@@ -18,16 +18,22 @@ package org.springframework.web.reactive.socket.server;
 import java.net.URISyntaxException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 
+import org.hamcrest.Matchers;
+import org.junit.Ignore;
 import org.junit.Test;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.core.publisher.MonoProcessor;
 import reactor.core.publisher.ReplayProcessor;
 
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.util.StringUtils;
 import org.springframework.web.reactive.HandlerMapping;
 import org.springframework.web.reactive.handler.SimpleUrlHandlerMapping;
+import org.springframework.web.reactive.socket.HandshakeInfo;
 import org.springframework.web.reactive.socket.WebSocketHandler;
 import org.springframework.web.reactive.socket.WebSocketMessage;
 import org.springframework.web.reactive.socket.WebSocketSession;
@@ -36,6 +42,7 @@ import org.springframework.web.reactive.socket.client.RxNettyWebSocketClient;
 import org.springframework.web.reactive.socket.client.WebSocketClient;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThat;
 
 /**
  * Integration tests with server-side {@link WebSocketHandler}s.
@@ -52,7 +59,7 @@ public class WebSocketIntegrationTests extends AbstractWebSocketIntegrationTests
 
 
 	@Test
-	public void echoReactorNettyClient() throws Exception {
+	public void echoReactorClient() throws Exception {
 		testEcho(new ReactorNettyWebSocketClient());
 	}
 
@@ -77,6 +84,49 @@ public class WebSocketIntegrationTests extends AbstractWebSocketIntegrationTests
 		assertEquals(input.collectList().blockMillis(5000), output.collectList().blockMillis(5000));
 	}
 
+	@Test
+	@Ignore("https://github.com/reactor/reactor-netty/issues/20")
+	public void subProtocolReactorNettyClient() throws Exception {
+		testSubProtocol(new ReactorNettyWebSocketClient());
+	}
+
+	@Test
+	public void subProtocolRxNettyClient() throws Exception {
+		testSubProtocol(new RxNettyWebSocketClient());
+	}
+
+	private void testSubProtocol(WebSocketClient client) throws URISyntaxException {
+		String protocol = "echo-v1";
+		AtomicReference<HandshakeInfo> infoRef = new AtomicReference<>();
+		MonoProcessor<Object> output = MonoProcessor.create();
+
+		client.execute(getUrl("/sub-protocol"),
+				new SubProtocolWebSocketHandler(protocol) {
+
+					@Override
+					public Mono<Void> handle(WebSocketSession session) {
+						infoRef.set(session.getHandshakeInfo());
+						return session.receive()
+								.map(WebSocketMessage::getPayloadAsText)
+								.subscribeWith(output)
+								.then();
+					}
+				})
+				.blockMillis(5000);
+
+		HandshakeInfo info = infoRef.get();
+		assertThat(info.getHeaders().getFirst("Upgrade"), Matchers.equalToIgnoringCase("websocket"));
+		assertEquals(protocol, info.getHeaders().getFirst("Sec-WebSocket-Protocol"));
+		assertEquals("Wrong protocol accepted", protocol, info.getSubProtocol().orElse("none"));
+		assertEquals("Wrong protocol detected on the server side", protocol, output.blockMillis(5000));
+	}
+
+
+	@Test
+	public void customHeaders() throws Exception {
+		// TODO
+	}
+
 
 	@Configuration
 	static class WebConfig {
@@ -86,6 +136,7 @@ public class WebSocketIntegrationTests extends AbstractWebSocketIntegrationTests
 
 			Map<String, WebSocketHandler> map = new HashMap<>();
 			map.put("/echo", new EchoWebSocketHandler());
+			map.put("/sub-protocol", new SubProtocolWebSocketHandler("echo-v1"));
 
 			SimpleUrlHandlerMapping mapping = new SimpleUrlHandlerMapping();
 			mapping.setUrlMap(map);
@@ -102,4 +153,25 @@ public class WebSocketIntegrationTests extends AbstractWebSocketIntegrationTests
 		}
 	}
 
+	private static class SubProtocolWebSocketHandler implements WebSocketHandler {
+
+		private final String subProtocols;
+
+		public SubProtocolWebSocketHandler(String subProtocols) {
+			this.subProtocols = subProtocols;
+		}
+
+		@Override
+		public String[] getSubProtocols() {
+			return StringUtils.commaDelimitedListToStringArray(this.subProtocols);
+		}
+
+		@Override
+		public Mono<Void> handle(WebSocketSession session) {
+			String protocol = session.getHandshakeInfo().getSubProtocol().orElse("none");
+			WebSocketMessage message = session.textMessage(protocol);
+			return session.send(Mono.just(message));
+		}
+	}
+
 }