Add support for Ping and Pong WebSocket messages

Issue: SPR-10876

Author: rstoyanchev

spring-websocket/src/main/java/org/springframework/web/socket/BinaryMessage.java

@@ -31,15 +31,15 @@ public final class BinaryMessage extends WebSocketMessage<ByteBuffer> {
 
 	/**
 	 * Create a new {@link BinaryMessage} instance.
-	 * @param payload a non-null payload
+	 * @param payload the non-null payload
 	 */
 	public BinaryMessage(ByteBuffer payload) {
 		this(payload, true);
 	}
 
 	/**
 	 * Create a new {@link BinaryMessage} instance.
-	 * @param payload a non-null payload
+	 * @param payload the non-null payload
 	 * @param isLast if the message is the last of a series of partial messages
 	 */
 	public BinaryMessage(ByteBuffer payload, boolean isLast) {
@@ -49,15 +49,15 @@ public BinaryMessage(ByteBuffer payload, boolean isLast) {
 
 	/**
 	 * Create a new {@link BinaryMessage} instance.
-	 * @param payload a non-null payload
+	 * @param payload the non-null payload
 	 */
 	public BinaryMessage(byte[] payload) {
 		this(payload, 0, (payload == null ? 0 : payload.length), true);
 	}
 
 	/**
 	 * Create a new {@link BinaryMessage} instance.
-	 * @param payload a non-null payload
+	 * @param payload the non-null payload
 	 * @param isLast if the message is the last of a series of partial messages
 	 */
 	public BinaryMessage(byte[] payload, boolean isLast) {

spring-websocket/src/main/java/org/springframework/web/socket/PingMessage.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2002-2013 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.socket;
+
+import java.nio.ByteBuffer;
+
+/**
+ * A WebSocket ping message.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public final class PingMessage extends WebSocketMessage<ByteBuffer> {
+
+
+	public PingMessage(ByteBuffer payload) {
+		super(payload);
+	}
+
+	@Override
+	protected int getPayloadSize() {
+		return (getPayload() != null) ? getPayload().remaining() : 0;
+	}
+
+	@Override
+	protected String toStringPayload() {
+		return (getPayload() != null) ? getPayload().toString() : null;
+	}
+
+}

spring-websocket/src/main/java/org/springframework/web/socket/PongMessage.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2002-2013 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.socket;
+
+import java.nio.ByteBuffer;
+
+/**
+ * A WebSocket pong message.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.0
+ */
+public final class PongMessage extends WebSocketMessage<ByteBuffer> {
+
+
+	public PongMessage(ByteBuffer payload) {
+		super(payload);
+	}
+
+	@Override
+	protected int getPayloadSize() {
+		return (getPayload() != null) ? getPayload().remaining() : 0;
+	}
+
+	@Override
+	protected String toStringPayload() {
+		return (getPayload() != null) ? getPayload().toString() : null;
+	}
+
+}

spring-websocket/src/main/java/org/springframework/web/socket/TextMessage.java

@@ -26,15 +26,15 @@ public final class TextMessage extends WebSocketMessage<String> {
 
 	/**
 	 * Create a new {@link TextMessage} instance.
-	 * @param payload the payload
+	 * @param payload the non-null payload
 	 */
 	public TextMessage(CharSequence payload) {
 		super(payload.toString(), true);
 	}
 
 	/**
 	 * Create a new {@link TextMessage} instance.
-	 * @param payload the payload
+	 * @param payload the non-null payload
 	 * @param isLast whether this the last part of a message received or transmitted in parts
 	 */
 	public TextMessage(CharSequence payload, boolean isLast) {

spring-websocket/src/main/java/org/springframework/web/socket/WebSocketMessage.java

@@ -24,8 +24,6 @@
  *
  * @author Rossen Stoyanchev
  * @since 4.0
- * @see BinaryMessage
- * @see TextMessage
  */
 public abstract class WebSocketMessage<T> {
 
@@ -34,12 +32,19 @@ public abstract class WebSocketMessage<T> {
 	private final boolean last;
 
 
+	/**
+	 * Create a new {@link WebSocketMessage} instance with the given payload.
+	 */
+	WebSocketMessage(T payload) {
+		this(payload, true);
+	}
+
 	/**
 	 * Create a new {@link WebSocketMessage} instance with the given payload.
 	 * @param payload a non-null payload
 	 */
 	WebSocketMessage(T payload, boolean isLast) {
-		Assert.notNull(payload, "Payload must not be null");
+		Assert.notNull(payload, "payload is required");
 		this.payload = payload;
 		this.last = isLast;
 	}

spring-websocket/src/main/java/org/springframework/web/socket/adapter/AbstractWebSocketSesssion.java

@@ -23,6 +23,8 @@
 import org.springframework.util.Assert;
 import org.springframework.web.socket.BinaryMessage;
 import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.PingMessage;
+import org.springframework.web.socket.PongMessage;
 import org.springframework.web.socket.TextMessage;
 import org.springframework.web.socket.WebSocketMessage;
 import org.springframework.web.socket.WebSocketSession;
@@ -101,15 +103,24 @@ public final void sendMessage(WebSocketMessage message) throws IOException {
 		else if (message instanceof BinaryMessage) {
 			sendBinaryMessage((BinaryMessage) message);
 		}
+		else if (message instanceof PingMessage) {
+			sendPingMessage((PingMessage) message);
+		}
+		else if (message instanceof PongMessage) {
+			sendPongMessage((PongMessage) message);
+		}
 		else {
 			throw new IllegalStateException("Unexpected WebSocketMessage type: " + message);
 		}
 	}
 
-	protected abstract void sendTextMessage(TextMessage message) throws IOException ;
+	protected abstract void sendTextMessage(TextMessage message) throws IOException;
+
+	protected abstract void sendBinaryMessage(BinaryMessage message) throws IOException;
 
-	protected abstract void sendBinaryMessage(BinaryMessage message) throws IOException ;
+	protected abstract void sendPingMessage(PingMessage message) throws IOException;
 
+	protected abstract void sendPongMessage(PongMessage message) throws IOException;
 
 	@Override
 	public final void close() throws IOException {

spring-websocket/src/main/java/org/springframework/web/socket/adapter/JettyWebSocketSession.java

@@ -26,6 +26,8 @@
 import org.springframework.util.ObjectUtils;
 import org.springframework.web.socket.BinaryMessage;
 import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.PingMessage;
+import org.springframework.web.socket.PongMessage;
 import org.springframework.web.socket.TextMessage;
 import org.springframework.web.socket.WebSocketSession;
 
@@ -117,6 +119,16 @@ protected void sendBinaryMessage(BinaryMessage message) throws IOException {
 		getNativeSession().getRemote().sendBytes(message.getPayload());
 	}
 
+	@Override
+	protected void sendPingMessage(PingMessage message) throws IOException {
+		getNativeSession().getRemote().sendPing(message.getPayload());
+	}
+
+	@Override
+	protected void sendPongMessage(PongMessage message) throws IOException {
+		getNativeSession().getRemote().sendPong(message.getPayload());
+	}
+
 	@Override
 	protected void closeInternal(CloseStatus status) throws IOException {
 		getNativeSession().close(status.getCode(), status.getReason());

spring-websocket/src/main/java/org/springframework/web/socket/adapter/StandardWebSocketHandlerAdapter.java

@@ -28,6 +28,7 @@
 import org.springframework.util.Assert;
 import org.springframework.web.socket.BinaryMessage;
 import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.PongMessage;
 import org.springframework.web.socket.TextMessage;
 import org.springframework.web.socket.WebSocketHandler;
 import org.springframework.web.socket.support.ExceptionWebSocketHandlerDecorator;
@@ -89,6 +90,13 @@ public void onMessage(ByteBuffer message) {
 			});
 		}
 
+		session.addMessageHandler(new MessageHandler.Whole<javax.websocket.PongMessage>() {
+			@Override
+			public void onMessage(javax.websocket.PongMessage message) {
+				handlePongMessage(session, message.getApplicationData());
+			}
+		});
+
 		try {
 			this.handler.afterConnectionEstablished(this.wsSession);
 		}
@@ -118,6 +126,16 @@ private void handleBinaryMessage(javax.websocket.Session session, ByteBuffer pay
 		}
 	}
 
+	private void handlePongMessage(javax.websocket.Session session, ByteBuffer payload) {
+		PongMessage pongMessage = new PongMessage(payload);
+		try {
+			this.handler.handleMessage(this.wsSession, pongMessage);
+		}
+		catch (Throwable t) {
+			ExceptionWebSocketHandlerDecorator.tryCloseWithError(this.wsSession, t, logger);
+		}
+	}
+
 	@Override
 	public void onClose(javax.websocket.Session session, CloseReason reason) {
 		CloseStatus closeStatus = new CloseStatus(reason.getCloseCode().getCode(), reason.getReasonPhrase());

spring-websocket/src/main/java/org/springframework/web/socket/adapter/StandardWebSocketSession.java

@@ -29,6 +29,8 @@
 import org.springframework.util.StringUtils;
 import org.springframework.web.socket.BinaryMessage;
 import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.PingMessage;
+import org.springframework.web.socket.PongMessage;
 import org.springframework.web.socket.TextMessage;
 import org.springframework.web.socket.WebSocketSession;
 
@@ -121,6 +123,16 @@ protected void sendBinaryMessage(BinaryMessage message) throws IOException {
 		getNativeSession().getBasicRemote().sendBinary(message.getPayload(), message.isLast());
 	}
 
+	@Override
+	protected void sendPingMessage(PingMessage message) throws IOException {
+		getNativeSession().getBasicRemote().sendPing(message.getPayload());
+	}
+
+	@Override
+	protected void sendPongMessage(PongMessage message) throws IOException {
+		getNativeSession().getBasicRemote().sendPong(message.getPayload());
+	}
+
 	@Override
 	protected void closeInternal(CloseStatus status) throws IOException {
 		getNativeSession().close(new CloseReason(CloseCodes.getCloseCode(status.getCode()), status.getReason()));

spring-websocket/src/main/java/org/springframework/web/socket/adapter/WebSocketHandlerAdapter.java

@@ -18,6 +18,7 @@
 
 import org.springframework.web.socket.BinaryMessage;
 import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.PongMessage;
 import org.springframework.web.socket.TextMessage;
 import org.springframework.web.socket.WebSocketHandler;
 import org.springframework.web.socket.WebSocketMessage;
@@ -32,6 +33,7 @@
  */
 public class WebSocketHandlerAdapter implements WebSocketHandler {
 
+
 	@Override
 	public void afterConnectionEstablished(WebSocketSession session) throws Exception {
 	}
@@ -44,6 +46,9 @@ public void handleMessage(WebSocketSession session, WebSocketMessage<?> message)
 		else if (message instanceof BinaryMessage) {
 			handleBinaryMessage(session, (BinaryMessage) message);
 		}
+		else if (message instanceof PongMessage) {
+			handlePongMessage(session, (PongMessage) message);
+		}
 		else {
 			throw new IllegalStateException("Unexpected WebSocket message type: " + message);
 		}
@@ -55,6 +60,9 @@ protected void handleTextMessage(WebSocketSession session, TextMessage message)
 	protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
 	}
 
+	protected void handlePongMessage(WebSocketSession session, PongMessage message) throws Exception {
+	}
+
 	@Override
 	public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
 	}