Spring/spring-framework
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Support splitting STOMP messages in WebSocketStompClient

Browse Source 
See gh-31970
This commit is contained in:
injae-kim
2024-01-08 00:45:54 +09:00
committed by rstoyanchev
parent bf014ef18b
commit 76d00d78db
6 changed files with 580 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
spring-websocket/src/main/java/org/springframework/web/socket/messaging/WebSocketStompClient.java
View File

@@ -35,6 +35,7 @@ import org.springframework.lang.Nullable;
import org.springframework.messaging.Message;
import org.springframework.messaging.simp.stomp.BufferingStompDecoder;
import org.springframework.messaging.simp.stomp.ConnectionHandlingStompSession;
import org.springframework.messaging.simp.stomp.SplittingStompEncoder;
import org.springframework.messaging.simp.stomp.StompClientSupport;
import org.springframework.messaging.simp.stomp.StompDecoder;
import org.springframework.messaging.simp.stomp.StompEncoder;
@@ -67,15 +68,23 @@ import org.springframework.web.util.UriComponentsBuilder;
* SockJsClient}.
*
* @author Rossen Stoyanchev
* @author Injae Kim
* @since 4.2
*/
public class WebSocketStompClient extends StompClientSupport implements SmartLifecycle {
private static final Log logger = LogFactory.getLog(WebSocketStompClient.class);
/**
* The default max size for in&outbound STOMP message.
*/
private static final int DEFAULT_MESSAGE_MAX_SIZE = 64 * 1024;
private final WebSocketClient webSocketClient;
private int inboundMessageSizeLimit = 64 * 1024;
private int inboundMessageSizeLimit = DEFAULT_MESSAGE_MAX_SIZE;
private int outboundMessageSizeLimit = DEFAULT_MESSAGE_MAX_SIZE;
private boolean autoStartup = true;
@@ -122,7 +131,7 @@ public class WebSocketStompClient extends StompClientSupport implements SmartLif
* Since a STOMP message can be received in multiple WebSocket messages,
* buffering may be required and this property determines the maximum buffer
* size per message.
* <p>By default this is set to 64 * 1024 (64K).
* <p>By default this is set to 64 * 1024 (64K), see {@link WebSocketStompClient#DEFAULT_MESSAGE_MAX_SIZE}.
*/
public void setInboundMessageSizeLimit(int inboundMessageSizeLimit) {
this.inboundMessageSizeLimit = inboundMessageSizeLimit;
@@ -135,6 +144,25 @@ public class WebSocketStompClient extends StompClientSupport implements SmartLif
return this.inboundMessageSizeLimit;
}
/**
* Configure the maximum size allowed for outbound STOMP message.
* If STOMP message's size exceeds {@link WebSocketStompClient#outboundMessageSizeLimit},
* STOMP message is split into multiple frames.
* <p>By default this is set to 64 * 1024 (64K), see {@link WebSocketStompClient#DEFAULT_MESSAGE_MAX_SIZE}.
* @since 6.2
*/
public void setOutboundMessageSizeLimit(int outboundMessageSizeLimit) {
this.outboundMessageSizeLimit = outboundMessageSizeLimit;
}
/**
* Get the configured outbound message buffer size in bytes.
* @since 6.2
*/
public int getOutboundMessageSizeLimit() {
return this.outboundMessageSizeLimit;
}
/**
* Set whether to auto-start the contained WebSocketClient when the Spring
* context has been refreshed.
@@ -373,7 +401,8 @@ public class WebSocketStompClient extends StompClientSupport implements SmartLif
private final TcpConnectionHandler<byte[]> stompSession;
private final StompWebSocketMessageCodec codec = new StompWebSocketMessageCodec(getInboundMessageSizeLimit());
private final StompWebSocketMessageCodec codec =
new StompWebSocketMessageCodec(getInboundMessageSizeLimit(),getOutboundMessageSizeLimit());
@Nullable
private volatile WebSocketSession session;
@@ -450,7 +479,9 @@ public class WebSocketStompClient extends StompClientSupport implements SmartLif
try {
WebSocketSession session = this.session;
Assert.state(session != null, "No WebSocketSession available");
session.sendMessage(this.codec.encode(message, session.getClass()));
for (WebSocketMessage<?> webSocketMessage : this.codec.encode(message, session.getClass())) {
session.sendMessage(webSocketMessage);
}
future.complete(null);
}
catch (Throwable ex) {
@@ -561,8 +592,11 @@ public class WebSocketStompClient extends StompClientSupport implements SmartLif
private final BufferingStompDecoder bufferingDecoder;
public StompWebSocketMessageCodec(int messageSizeLimit) {
this.bufferingDecoder = new BufferingStompDecoder(DECODER, messageSizeLimit);
private final SplittingStompEncoder splittingEncoder;
public StompWebSocketMessageCodec(int inboundMessageSizeLimit, int outboundMessageSizeLimit) {
this.bufferingDecoder = new BufferingStompDecoder(DECODER, inboundMessageSizeLimit);
this.splittingEncoder = new SplittingStompEncoder(ENCODER, outboundMessageSizeLimit);
}
public List<Message<byte[]>> decode(WebSocketMessage<?> webSocketMessage) {
@@ -588,17 +622,21 @@ public class WebSocketStompClient extends StompClientSupport implements SmartLif
return result;
}
public WebSocketMessage<?> encode(Message<byte[]> message, Class<? extends WebSocketSession> sessionType) {
public List<WebSocketMessage<?>> encode(Message<byte[]> message, Class<? extends WebSocketSession> sessionType) {
StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
Assert.notNull(accessor, "No StompHeaderAccessor available");
byte[] payload = message.getPayload();
byte[] bytes = ENCODER.encode(accessor.getMessageHeaders(), payload);
List<byte[]> frames = splittingEncoder.encode(accessor.getMessageHeaders(), payload);
boolean useBinary = (payload.length > 0 &&
!(SockJsSession.class.isAssignableFrom(sessionType)) &&
MimeTypeUtils.APPLICATION_OCTET_STREAM.isCompatibleWith(accessor.getContentType()));
return (useBinary ? new BinaryMessage(bytes) : new TextMessage(bytes));
List<WebSocketMessage<?>> messages = new ArrayList<>();
for (byte[] frame : frames) {
messages.add(useBinary ? new BinaryMessage(frame) : new TextMessage(frame));
}
return messages;
}
}

67
spring-websocket/src/test/java/org/springframework/web/socket/messaging/WebSocketStompClientTests.java
View File

@@ -65,6 +65,7 @@ import static org.mockito.Mockito.verifyNoMoreInteractions;
* Tests for {@link WebSocketStompClient}.
*
* @author Rossen Stoyanchev
* @author Injae Kim
*/
@MockitoSettings(strictness = Strictness.LENIENT)
class WebSocketStompClientTests {
@@ -211,6 +212,29 @@ class WebSocketStompClientTests {
assertThat(textMessage.getPayload()).isEqualTo("SEND\ndestination:/topic/foo\ncontent-length:7\n\npayload\0");
}
@Test
void sendWebSocketMessageExceedOutboundMessageSizeLimit() throws Exception {
stompClient.setOutboundMessageSizeLimit(30);
StompHeaderAccessor accessor = StompHeaderAccessor.create(StompCommand.SEND);
accessor.setDestination("/topic/foo");
byte[] payload = "payload".getBytes(StandardCharsets.UTF_8);
getTcpConnection().sendAsync(MessageBuilder.createMessage(payload, accessor.getMessageHeaders()));
ArgumentCaptor<TextMessage> textMessageCaptor = ArgumentCaptor.forClass(TextMessage.class);
verify(this.webSocketSession, times(2)).sendMessage(textMessageCaptor.capture());
TextMessage textMessage = textMessageCaptor.getAllValues().get(0);
assertThat(textMessage).isNotNull();
assertThat(textMessage.getPayload()).isEqualTo("SEND\ndestination:/topic/foo\nco");
assertThat(textMessage.getPayload().getBytes().length).isEqualTo(30);
textMessage = textMessageCaptor.getAllValues().get(1);
assertThat(textMessage).isNotNull();
assertThat(textMessage.getPayload()).isEqualTo("ntent-length:7\n\npayload\0");
assertThat(textMessage.getPayload().getBytes().length).isEqualTo(24);
}
@Test
void sendWebSocketBinary() throws Exception {
StompHeaderAccessor accessor = StompHeaderAccessor.create(StompCommand.SEND);
@@ -228,6 +252,49 @@ class WebSocketStompClientTests {
.isEqualTo("SEND\ndestination:/b\ncontent-type:application/octet-stream\ncontent-length:7\n\npayload\0");
}
@Test
void sendWebSocketBinaryExceedOutboundMessageSizeLimit() throws Exception {
stompClient.setOutboundMessageSizeLimit(50);
StompHeaderAccessor accessor = StompHeaderAccessor.create(StompCommand.SEND);
accessor.setDestination("/b");
accessor.setContentType(MimeTypeUtils.APPLICATION_OCTET_STREAM);
byte[] payload = "payload".getBytes(StandardCharsets.UTF_8);
getTcpConnection().sendAsync(MessageBuilder.createMessage(payload, accessor.getMessageHeaders()));
ArgumentCaptor<BinaryMessage> binaryMessageCaptor = ArgumentCaptor.forClass(BinaryMessage.class);
verify(this.webSocketSession, times(2)).sendMessage(binaryMessageCaptor.capture());
BinaryMessage binaryMessage = binaryMessageCaptor.getAllValues().get(0);
assertThat(binaryMessage).isNotNull();
assertThat(new String(binaryMessage.getPayload().array(), StandardCharsets.UTF_8))
.isEqualTo("SEND\ndestination:/b\ncontent-type:application/octet");
assertThat(binaryMessage.getPayload().array().length).isEqualTo(50);
binaryMessage = binaryMessageCaptor.getAllValues().get(1);
assertThat(binaryMessage).isNotNull();
assertThat(new String(binaryMessage.getPayload().array(), StandardCharsets.UTF_8))
.isEqualTo("-stream\ncontent-length:7\n\npayload\0");
assertThat(binaryMessage.getPayload().array().length).isEqualTo(34);
}
@Test
void reassembleReceivedIFragmentedFrames() throws Exception {
WebSocketHandler handler = connect();
handler.handleMessage(this.webSocketSession, new TextMessage("SEND\ndestination:/topic/foo\nco"));
handler.handleMessage(this.webSocketSession, new TextMessage("ntent-length:7\n\npayload\0"));
ArgumentCaptor<Message> receiveMessageCaptor = ArgumentCaptor.forClass(Message.class);
verify(this.stompSession).handleMessage(receiveMessageCaptor.capture());
Message<byte[]> receiveMessage = receiveMessageCaptor.getValue();
assertThat(receiveMessage).isNotNull();
StompHeaderAccessor headers = StompHeaderAccessor.wrap(receiveMessage);
assertThat(headers.toNativeHeaderMap()).hasSize(2);
assertThat(headers.getContentLength()).isEqualTo(7);
assertThat(headers.getDestination()).isEqualTo("/topic/foo");
assertThat(new String(receiveMessage.getPayload())).isEqualTo("payload");
}
@Test
void heartbeatDefaultValue() {
WebSocketStompClient stompClient = new WebSocketStompClient(mock());