Allow configuring WebSocket message size in a session

Issue: SPR-11575

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

@@ -83,6 +83,26 @@ public interface WebSocketSession {
 	 */
 	String getAcceptedProtocol();
 
+	/**
+	 * Configure the maximum size for an incoming text message.
+	 */
+	void setTextMessageSizeLimit(int messageSizeLimit);
+
+	/**
+	 * Get the configured maximum size for an incoming text message.
+	 */
+	int getTextMessageSizeLimit();
+
+	/**
+	 * Configure the maximum size for an incoming binary message.
+	 */
+	void setBinaryMessageSizeLimit(int messageSizeLimit);
+
+	/**
+	 * Get the configured maximum size for an incoming binary message.
+	 */
+	int getBinaryMessageSizeLimit();
+
 	/**
 	 * Return the negotiated extensions or {@code null} if none was specified or
 	 * negotiated successfully.

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

@@ -126,6 +126,30 @@ public class JettyWebSocketSession extends AbstractWebSocketSession<Session> {
 		return getNativeSession().getUpgradeResponse().getAcceptedSubProtocol();
 	}
 
+	@Override
+	public void setTextMessageSizeLimit(int messageSizeLimit) {
+		checkNativeSessionInitialized();
+		getNativeSession().getPolicy().setMaxTextMessageSize(messageSizeLimit);
+	}
+
+	@Override
+	public int getTextMessageSizeLimit() {
+		checkNativeSessionInitialized();
+		return getNativeSession().getPolicy().getMaxTextMessageSize();
+	}
+
+	@Override
+	public void setBinaryMessageSizeLimit(int messageSizeLimit) {
+		checkNativeSessionInitialized();
+		getNativeSession().getPolicy().setMaxBinaryMessageSize(messageSizeLimit);
+	}
+
+	@Override
+	public int getBinaryMessageSizeLimit() {
+		checkNativeSessionInitialized();
+		return getNativeSession().getPolicy().getMaxBinaryMessageSize();
+	}
+
 	@Override
 	public List<WebSocketExtension> getExtensions() {
 		checkNativeSessionInitialized();

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2013 the original author or authors.
+ * Copyright 2002-2014 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.
@@ -140,6 +140,30 @@ public class StandardWebSocketSession extends AbstractWebSocketSession<Session>
 		return StringUtils.isEmpty(protocol)? null : protocol;
 	}
 
+	@Override
+	public void setTextMessageSizeLimit(int messageSizeLimit) {
+		checkNativeSessionInitialized();
+		getNativeSession().setMaxTextMessageBufferSize(messageSizeLimit);
+	}
+
+	@Override
+	public int getTextMessageSizeLimit() {
+		checkNativeSessionInitialized();
+		return getNativeSession().getMaxTextMessageBufferSize();
+	}
+
+	@Override
+	public void setBinaryMessageSizeLimit(int messageSizeLimit) {
+		checkNativeSessionInitialized();
+		getNativeSession().setMaxBinaryMessageBufferSize(messageSizeLimit);
+	}
+
+	@Override
+	public int getBinaryMessageSizeLimit() {
+		checkNativeSessionInitialized();
+		return getNativeSession().getMaxBinaryMessageBufferSize();
+	}
+
 	@Override
 	public List<WebSocketExtension> getExtensions() {
 		checkNativeSessionInitialized();

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

@@ -118,6 +118,26 @@ public class WebSocketSessionDecorator implements WebSocketSession {
 		return this.delegate.getExtensions();
 	}
 
+	@Override
+	public void setTextMessageSizeLimit(int messageSizeLimit) {
+		this.delegate.setTextMessageSizeLimit(messageSizeLimit);
+	}
+
+	@Override
+	public int getTextMessageSizeLimit() {
+		return this.delegate.getTextMessageSizeLimit();
+	}
+
+	@Override
+	public void setBinaryMessageSizeLimit(int messageSizeLimit) {
+		this.delegate.setBinaryMessageSizeLimit(messageSizeLimit);
+	}
+
+	@Override
+	public int getBinaryMessageSizeLimit() {
+		return this.delegate.getBinaryMessageSizeLimit();
+	}
+
 	@Override
 	public boolean isOpen() {
 		return this.delegate.isOpen();

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

@@ -61,6 +61,16 @@ import org.springframework.web.socket.sockjs.transport.SockJsSession;
  */
 public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationEventPublisherAware {
 
+	/**
+	 * This protocol handler supports assembling large STOMP messages split into
+	 * multiple WebSocket messages. STOMP clients (like stomp.js) split large STOMP
+	 * messages at 16K boundaries.
+	 *
+	 * <p>We need to ensure the WebSocket server buffer is configured to support
+	 * that size at a minimum plus a little extra for any potential SockJS framing.
+	 */
+	public static final int MINIMUM_WEBSOCKET_MESSAGE_SIZE = 16 * 1024 + 256;
+
 	/**
 	 * The name of the header set on the CONNECTED frame indicating the name
 	 * of the user authenticated on the WebSocket session.
@@ -332,6 +342,9 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE
 
 	@Override
 	public void afterSessionStarted(WebSocketSession session, MessageChannel outputChannel) {
+		if (session.getTextMessageSizeLimit() < MINIMUM_WEBSOCKET_MESSAGE_SIZE) {
+			session.setTextMessageSizeLimit(MINIMUM_WEBSOCKET_MESSAGE_SIZE);
+		}
 		this.decoders.put(session.getId(), new BufferingStompDecoder(getMessageSizeLimit()));
 	}
 

spring-websocket/src/main/java/org/springframework/web/socket/sockjs/transport/session/AbstractHttpSockJsSession.java

@@ -148,6 +148,26 @@ public abstract class AbstractHttpSockJsSession extends AbstractSockJsSession {
 		return (control != null && !control.isCompleted());
 	}
 
+	@Override
+	public void setTextMessageSizeLimit(int messageSizeLimit) {
+		// ignore
+	}
+
+	@Override
+	public int getTextMessageSizeLimit() {
+		return -1;
+	}
+
+	@Override
+	public void setBinaryMessageSizeLimit(int messageSizeLimit) {
+		// ignore
+	}
+
+	@Override
+	public int getBinaryMessageSizeLimit() {
+		return -1;
+	}
+
 	@Override
 	public List<WebSocketExtension> getExtensions() {
 		return Collections.emptyList();

spring-websocket/src/main/java/org/springframework/web/socket/sockjs/transport/session/WebSocketServerSockJsSession.java

@@ -91,6 +91,30 @@ public class WebSocketServerSockJsSession extends AbstractSockJsSession implemen
 		return this.webSocketSession.getAcceptedProtocol();
 	}
 
+	@Override
+	public void setTextMessageSizeLimit(int messageSizeLimit) {
+		checkDelegateSessionInitialized();
+		this.webSocketSession.setTextMessageSizeLimit(messageSizeLimit);
+	}
+
+	@Override
+	public int getTextMessageSizeLimit() {
+		checkDelegateSessionInitialized();
+		return this.webSocketSession.getTextMessageSizeLimit();
+	}
+
+	@Override
+	public void setBinaryMessageSizeLimit(int messageSizeLimit) {
+		checkDelegateSessionInitialized();
+		this.webSocketSession.setBinaryMessageSizeLimit(messageSizeLimit);
+	}
+
+	@Override
+	public int getBinaryMessageSizeLimit() {
+		checkDelegateSessionInitialized();
+		return this.webSocketSession.getBinaryMessageSizeLimit();
+	}
+
 	@Override
 	public List<WebSocketExtension> getExtensions() {
 		checkDelegateSessionInitialized();

spring-websocket/src/test/java/org/springframework/web/socket/handler/TestWebSocketSession.java

@@ -145,6 +145,24 @@ public class TestWebSocketSession implements WebSocketSession {
 		this.protocol = protocol;
 	}
 
+	@Override
+	public void setTextMessageSizeLimit(int messageSizeLimit) {
+	}
+
+	@Override
+	public int getTextMessageSizeLimit() {
+		return 0;
+	}
+
+	@Override
+	public void setBinaryMessageSizeLimit(int messageSizeLimit) {
+	}
+
+	@Override
+	public int getBinaryMessageSizeLimit() {
+		return 0;
+	}
+
 	@Override
 	public List<WebSocketExtension> getExtensions() {
 		return this.extensions;

spring-websocket/src/test/java/org/springframework/web/socket/sockjs/transport/session/TestSockJsSession.java

@@ -128,6 +128,24 @@ public class TestSockJsSession extends AbstractSockJsSession {
 		this.subProtocol = protocol;
 	}
 
+	@Override
+	public void setTextMessageSizeLimit(int messageSizeLimit) {
+	}
+
+	@Override
+	public int getTextMessageSizeLimit() {
+		return 0;
+	}
+
+	@Override
+	public void setBinaryMessageSizeLimit(int messageSizeLimit) {
+	}
+
+	@Override
+	public int getBinaryMessageSizeLimit() {
+		return 0;
+	}
+
 	@Override
 	public List<WebSocketExtension> getExtensions() {
 		return this.extensions;