Add reactive WebSocketClient and RxNetty implementation

Issue: SPR-14527

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

@@ -0,0 +1,121 @@
+/*
+ * Copyright 2002-2016 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
+ *
+ *      http://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.security.NoSuchAlgorithmException;
+import java.util.function.Function;
+
+import javax.net.ssl.SSLContext;
+import javax.net.ssl.SSLEngine;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
+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;
+import reactor.core.publisher.Mono;
+import reactor.util.function.Tuples;
+import rx.Observable;
+import rx.RxReactiveStreams;
+
+import org.springframework.core.io.buffer.NettyDataBufferFactory;
+import org.springframework.http.HttpHeaders;
+import org.springframework.web.reactive.socket.WebSocketSession;
+import org.springframework.web.reactive.socket.adapter.HandshakeInfo;
+import org.springframework.web.reactive.socket.adapter.RxNettyWebSocketSession;
+
+/**
+ * A {@link WebSocketClient} based on RxNetty.
+ *
+ * @author Rossen Stoyanchev
+ * @since 5.0
+ */
+public class RxNettyWebSocketClient implements WebSocketClient {
+
+	private final Function<URI, HttpClient<ByteBuf, ByteBuf>> httpClientFactory;
+
+
+	/**
+	 * Default constructor that uses {@link HttpClient#newClient(String, int)}
+	 * to create HTTP client instances when connecting.
+	 */
+	public RxNettyWebSocketClient() {
+		this(RxNettyWebSocketClient::createDefaultHttpClient);
+	}
+
+	/**
+	 * Constructor with a function to create {@link HttpClient} instances.
+	 * @param httpClientFactory factory to create clients
+	 */
+	public RxNettyWebSocketClient(Function<URI, HttpClient<ByteBuf, ByteBuf>> httpClientFactory) {
+		this.httpClientFactory = httpClientFactory;
+	}
+
+	private static HttpClient<ByteBuf, ByteBuf> createDefaultHttpClient(URI url) {
+		boolean secure = "wss".equals(url.getScheme());
+		int port = url.getPort() > 0 ? url.getPort() : secure ? 443 : 80;
+		HttpClient<ByteBuf, ByteBuf> httpClient = HttpClient.newClient(url.getHost(), port);
+		if (secure) {
+			try {
+				SSLContext context = SSLContext.getDefault();
+				SSLEngine engine = context.createSSLEngine(url.getHost(), port);
+				engine.setUseClientMode(true);
+				httpClient.secure(engine);
+			}
+			catch (NoSuchAlgorithmException ex) {
+				throw new IllegalStateException("Failed to create HttpClient for " + url, ex);
+			}
+		}
+		return httpClient;
+	}
+
+
+	@Override
+	public Mono<WebSocketSession> connect(URI url) {
+		return connect(url, new HttpHeaders());
+	}
+
+	@Override
+	public Mono<WebSocketSession> connect(URI url, HttpHeaders headers) {
+		HandshakeInfo info = new HandshakeInfo(url, headers, Mono.empty());
+		Observable<WebSocketSession> observable = connectInternal(info);
+		return Mono.from(RxReactiveStreams.toPublisher(observable));
+	}
+
+	private Observable<WebSocketSession> connectInternal(HandshakeInfo info) {
+		return createWebSocketRequest(info.getUri())
+				.flatMap(response -> {
+					ByteBufAllocator allocator = response.unsafeNettyChannel().alloc();
+					NettyDataBufferFactory bufferFactory = new NettyDataBufferFactory(allocator);
+					Observable<WebSocketConnection> conn = response.getWebSocketConnection();
+					return Observable.zip(conn, Observable.just(bufferFactory), Tuples::of);
+				})
+				.map(tuple -> {
+					WebSocketConnection conn = tuple.getT1();
+					NettyDataBufferFactory bufferFactory = tuple.getT2();
+					return new RxNettyWebSocketSession(conn, info, bufferFactory);
+				});
+	}
+
+	private WebSocketRequest<ByteBuf> createWebSocketRequest(URI url) {
+		String query = url.getRawQuery();
+		return this.httpClientFactory.apply(url)
+				.createGet(url.getRawPath() + (query != null ? "?" + query : ""))
+				.requestWebSocketUpgrade();
+	}
+
+}

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

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2002-2016 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
+ *
+ *      http://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 reactor.core.publisher.Mono;
+
+import org.springframework.http.HttpHeaders;
+import org.springframework.web.reactive.socket.WebSocketSession;
+
+/**
+ * Contract for starting a WebSocket interaction.
+ *
+ * @author Rossen Stoyanchev
+ * @since 5.0
+ */
+public interface WebSocketClient {
+
+	/**
+	 * Start a WebSocket interaction to the given url.
+	 * @param url the handshake url
+	 * @return the session for the WebSocket interaction
+	 */
+	Mono<WebSocketSession> connect(URI url);
+
+	/**
+	 * Start a WebSocket interaction to the given url.
+	 * @param url the handshake url
+	 * @param headers headers for the handshake request
+	 * @return the session for the WebSocket interaction
+	 */
+	Mono<WebSocketSession> connect(URI url, HttpHeaders headers);
+
+}

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

@@ -0,0 +1,4 @@
+/**
+ * Client support for WebSocket interactions.
+ */
+package org.springframework.web.reactive.socket.client;

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

@@ -15,24 +15,22 @@
  */
 package org.springframework.web.reactive.socket.server;
 
-import java.nio.charset.StandardCharsets;
+import java.net.URI;
 import java.util.HashMap;
 import java.util.Map;
 
-import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
-import io.netty.handler.codec.http.websocketx.WebSocketFrame;
-import io.reactivex.netty.protocol.http.client.HttpClient;
-import io.reactivex.netty.protocol.http.ws.client.WebSocketResponse;
 import org.junit.Test;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
-import rx.Observable;
 
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.web.reactive.HandlerMapping;
 import org.springframework.web.reactive.handler.SimpleUrlHandlerMapping;
 import org.springframework.web.reactive.socket.WebSocketHandler;
 import org.springframework.web.reactive.socket.WebSocketSession;
+import org.springframework.web.reactive.socket.client.RxNettyWebSocketClient;
 
 import static org.junit.Assert.assertEquals;
 
@@ -42,7 +40,7 @@ import static org.junit.Assert.assertEquals;
  * @author Rossen Stoyanchev
  */
 @SuppressWarnings({"unused", "WeakerAccess"})
-public class ServerWebSocketIntegrationTests extends AbstractWebSocketIntegrationTests {
+public class WebSocketIntegrationTests extends AbstractWebSocketIntegrationTests {
 
 
 	@Override
@@ -54,21 +52,20 @@ public class ServerWebSocketIntegrationTests extends AbstractWebSocketIntegratio
 	@Test
 	public void echo() throws Exception {
 		int count = 100;
-		Observable<String> input = Observable.range(1, count).map(index -> "msg-" + index);
-		Observable<String> output = HttpClient.newClient("localhost", this.port)
-				.createGet("/echo")
-				.requestWebSocketUpgrade()
-				.flatMap(WebSocketResponse::getWebSocketConnection)
-				.flatMap(conn -> conn
-						.write(input.map(TextWebSocketFrame::new)).cast(WebSocketFrame.class)
-						.mergeWith(conn.getInput())
-						.take(count)
-						.map(frame -> {
-							String text = frame.content().toString(StandardCharsets.UTF_8);
-							frame.release();
-							return text;
-						}));
-		assertEquals(input.toList().toBlocking().first(), output.toList().toBlocking().first());
+		Flux<String> input = Flux.range(1, count).map(index -> "msg-" + index);
+		Flux<String> output = new RxNettyWebSocketClient()
+				.connect(new URI("ws://localhost:" + this.port + "/echo"))
+				.flatMap(session -> session
+						.send(input.map(session::textMessage))
+						.thenMany(session.receive()
+								.take(count)
+								.map(message -> {
+									String text = message.getPayloadAsText();
+									DataBufferUtils.release(message.getPayload());
+									return text;
+								})
+						));
+		assertEquals(input.collectList().blockMillis(5000), output.collectList().blockMillis(5000));
 	}