org.springframework.web.socket.messaging.SubProtocolWebSocketHandler Maven / Gradle / Ivy
/*
* Copyright 2002-2020 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
*
* https://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.messaging;
import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantLock;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.context.SmartLifecycle;
import org.springframework.lang.Nullable;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.MessageHandler;
import org.springframework.messaging.MessagingException;
import org.springframework.messaging.SubscribableChannel;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.SubProtocolCapable;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.ConcurrentWebSocketSessionDecorator;
import org.springframework.web.socket.handler.SessionLimitExceededException;
import org.springframework.web.socket.sockjs.transport.session.PollingSockJsSession;
import org.springframework.web.socket.sockjs.transport.session.StreamingSockJsSession;
/**
* An implementation of {@link WebSocketHandler} that delegates incoming WebSocket
* messages to a {@link SubProtocolHandler} along with a {@link MessageChannel} to which
* the sub-protocol handler can send messages from WebSocket clients to the application.
*
* Also an implementation of {@link MessageHandler} that finds the WebSocket session
* associated with the {@link Message} and passes it, along with the message, to the
* sub-protocol handler to send messages from the application back to the client.
*
* @author Rossen Stoyanchev
* @author Juergen Hoeller
* @author Andy Wilkinson
* @author Artem Bilan
* @since 4.0
*/
public class SubProtocolWebSocketHandler
implements WebSocketHandler, SubProtocolCapable, MessageHandler, SmartLifecycle {
/** The default value for {@link #setTimeToFirstMessage(int) timeToFirstMessage}. */
private static final int DEFAULT_TIME_TO_FIRST_MESSAGE = 60 * 1000;
private final Log logger = LogFactory.getLog(SubProtocolWebSocketHandler.class);
private final MessageChannel clientInboundChannel;
private final SubscribableChannel clientOutboundChannel;
private final Map protocolHandlerLookup =
new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
private final Set protocolHandlers = new LinkedHashSet<>();
@Nullable
private SubProtocolHandler defaultProtocolHandler;
private final Map sessions = new ConcurrentHashMap<>();
private int sendTimeLimit = 10 * 1000;
private int sendBufferSizeLimit = 512 * 1024;
private int timeToFirstMessage = DEFAULT_TIME_TO_FIRST_MESSAGE;
private volatile long lastSessionCheckTime = System.currentTimeMillis();
private final ReentrantLock sessionCheckLock = new ReentrantLock();
private final DefaultStats stats = new DefaultStats();
private volatile boolean running;
private final Object lifecycleMonitor = new Object();
/**
* Create a new {@code SubProtocolWebSocketHandler} for the given inbound and outbound channels.
* @param clientInboundChannel the inbound {@code MessageChannel}
* @param clientOutboundChannel the outbound {@code MessageChannel}
*/
public SubProtocolWebSocketHandler(MessageChannel clientInboundChannel, SubscribableChannel clientOutboundChannel) {
Assert.notNull(clientInboundChannel, "Inbound MessageChannel must not be null");
Assert.notNull(clientOutboundChannel, "Outbound MessageChannel must not be null");
this.clientInboundChannel = clientInboundChannel;
this.clientOutboundChannel = clientOutboundChannel;
}
/**
* Configure one or more handlers to use depending on the sub-protocol requested by
* the client in the WebSocket handshake request.
* @param protocolHandlers the sub-protocol handlers to use
*/
public void setProtocolHandlers(List protocolHandlers) {
this.protocolHandlerLookup.clear();
this.protocolHandlers.clear();
for (SubProtocolHandler handler : protocolHandlers) {
addProtocolHandler(handler);
}
}
public List getProtocolHandlers() {
return new ArrayList<>(this.protocolHandlers);
}
/**
* Register a sub-protocol handler.
*/
public void addProtocolHandler(SubProtocolHandler handler) {
List protocols = handler.getSupportedProtocols();
if (CollectionUtils.isEmpty(protocols)) {
if (logger.isErrorEnabled()) {
logger.error("No sub-protocols for " + handler);
}
return;
}
for (String protocol : protocols) {
SubProtocolHandler replaced = this.protocolHandlerLookup.put(protocol, handler);
if (replaced != null && replaced != handler) {
throw new IllegalStateException("Cannot map " + handler +
" to protocol '" + protocol + "': already mapped to " + replaced + ".");
}
}
this.protocolHandlers.add(handler);
}
/**
* Return the sub-protocols keyed by protocol name.
*/
public Map getProtocolHandlerMap() {
return this.protocolHandlerLookup;
}
/**
* Set the {@link SubProtocolHandler} to use when the client did not request a
* sub-protocol.
* @param defaultProtocolHandler the default handler
*/
public void setDefaultProtocolHandler(@Nullable SubProtocolHandler defaultProtocolHandler) {
this.defaultProtocolHandler = defaultProtocolHandler;
if (this.protocolHandlerLookup.isEmpty()) {
setProtocolHandlers(Collections.singletonList(defaultProtocolHandler));
}
}
/**
* Return the default sub-protocol handler to use.
*/
@Nullable
public SubProtocolHandler getDefaultProtocolHandler() {
return this.defaultProtocolHandler;
}
/**
* Return all supported protocols.
*/
@Override
public List getSubProtocols() {
return new ArrayList<>(this.protocolHandlerLookup.keySet());
}
/**
* Specify the send-time limit (milliseconds).
* @see ConcurrentWebSocketSessionDecorator
*/
public void setSendTimeLimit(int sendTimeLimit) {
this.sendTimeLimit = sendTimeLimit;
}
/**
* Return the send-time limit (milliseconds).
*/
public int getSendTimeLimit() {
return this.sendTimeLimit;
}
/**
* Specify the buffer-size limit (number of bytes).
* @see ConcurrentWebSocketSessionDecorator
*/
public void setSendBufferSizeLimit(int sendBufferSizeLimit) {
this.sendBufferSizeLimit = sendBufferSizeLimit;
}
/**
* Return the buffer-size limit (number of bytes).
*/
public int getSendBufferSizeLimit() {
return this.sendBufferSizeLimit;
}
/**
* Set the maximum time allowed in milliseconds after the WebSocket connection
* is established and before the first sub-protocol message is received.
* This handler is for WebSocket connections that use a sub-protocol.
* Therefore, we expect the client to send at least one sub-protocol message
* in the beginning, or else we assume the connection isn't doing well, e.g.
* proxy issue, slow network, and can be closed.
*
By default this is set to {@code 60,000} (1 minute).
* @param timeToFirstMessage the maximum time allowed in milliseconds
* @since 5.1
* @see #checkSessions()
*/
public void setTimeToFirstMessage(int timeToFirstMessage) {
this.timeToFirstMessage = timeToFirstMessage;
}
/**
* Return the maximum time allowed after the WebSocket connection is
* established and before the first sub-protocol message.
* @since 5.1
*/
public int getTimeToFirstMessage() {
return this.timeToFirstMessage;
}
/**
* Return a String describing internal state and counters.
* Effectively {@code toString()} on {@link #getStats() getStats()}.
*/
public String getStatsInfo() {
return this.stats.toString();
}
/**
* Return a structured object with various session counters.
* @since 5.2
*/
public Stats getStats() {
return this.stats;
}
@Override
public final void start() {
Assert.isTrue(this.defaultProtocolHandler != null || !this.protocolHandlers.isEmpty(), "No handlers");
synchronized (this.lifecycleMonitor) {
this.clientOutboundChannel.subscribe(this);
this.running = true;
}
}
@Override
public final void stop() {
synchronized (this.lifecycleMonitor) {
this.running = false;
this.clientOutboundChannel.unsubscribe(this);
}
// Proactively notify all active WebSocket sessions
for (WebSocketSessionHolder holder : this.sessions.values()) {
try {
holder.getSession().close(CloseStatus.GOING_AWAY);
}
catch (Throwable ex) {
if (logger.isWarnEnabled()) {
logger.warn("Failed to close '" + holder.getSession() + "': " + ex);
}
}
}
}
@Override
public final void stop(Runnable callback) {
synchronized (this.lifecycleMonitor) {
stop();
callback.run();
}
}
@Override
public final boolean isRunning() {
return this.running;
}
@Override
public void afterConnectionEstablished(WebSocketSession session) throws Exception {
// WebSocketHandlerDecorator could close the session
if (!session.isOpen()) {
return;
}
this.stats.incrementSessionCount(session);
session = decorateSession(session);
this.sessions.put(session.getId(), new WebSocketSessionHolder(session));
findProtocolHandler(session).afterSessionStarted(session, this.clientInboundChannel);
}
/**
* Handle an inbound message from a WebSocket client.
*/
@Override
public void handleMessage(WebSocketSession session, WebSocketMessage message) throws Exception {
WebSocketSessionHolder holder = this.sessions.get(session.getId());
if (holder != null) {
session = holder.getSession();
}
SubProtocolHandler protocolHandler = findProtocolHandler(session);
protocolHandler.handleMessageFromClient(session, message, this.clientInboundChannel);
if (holder != null) {
holder.setHasHandledMessages();
}
checkSessions();
}
/**
* Handle an outbound Spring Message to a WebSocket client.
*/
@Override
public void handleMessage(Message message) throws MessagingException {
String sessionId = resolveSessionId(message);
if (sessionId == null) {
if (logger.isErrorEnabled()) {
logger.error("Could not find session id in " + message);
}
return;
}
WebSocketSessionHolder holder = this.sessions.get(sessionId);
if (holder == null) {
if (logger.isDebugEnabled()) {
// The broker may not have removed the session yet
logger.debug("No session for " + message);
}
return;
}
WebSocketSession session = holder.getSession();
try {
findProtocolHandler(session).handleMessageToClient(session, message);
}
catch (SessionLimitExceededException ex) {
try {
if (logger.isDebugEnabled()) {
logger.debug("Terminating '" + session + "'", ex);
}
else if (logger.isWarnEnabled()) {
logger.warn("Terminating '" + session + "': " + ex.getMessage());
}
this.stats.incrementLimitExceededCount();
clearSession(session, ex.getStatus()); // clear first, session may be unresponsive
session.close(ex.getStatus());
}
catch (Exception secondException) {
logger.debug("Failure while closing session " + sessionId + ".", secondException);
}
}
catch (Exception ex) {
// Could be part of normal workflow (e.g. browser tab closed)
if (logger.isDebugEnabled()) {
logger.debug("Failed to send message to client in " + session + ": " + message, ex);
}
}
}
@Override
public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
this.stats.incrementTransportError();
}
@Override
public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
clearSession(session, closeStatus);
}
@Override
public boolean supportsPartialMessages() {
return false;
}
/**
* Decorate the given {@link WebSocketSession}, if desired.
*
The default implementation builds a {@link ConcurrentWebSocketSessionDecorator}
* with the configured {@link #getSendTimeLimit() send-time limit} and
* {@link #getSendBufferSizeLimit() buffer-size limit}.
* @param session the original {@code WebSocketSession}
* @return the decorated {@code WebSocketSession}, or potentially the given session as-is
* @since 4.3.13
*/
protected WebSocketSession decorateSession(WebSocketSession session) {
return new ConcurrentWebSocketSessionDecorator(session, getSendTimeLimit(), getSendBufferSizeLimit());
}
/**
* Find a {@link SubProtocolHandler} for the given session.
* @param session the {@code WebSocketSession} to find a handler for
*/
protected final SubProtocolHandler findProtocolHandler(WebSocketSession session) {
String protocol = null;
try {
protocol = session.getAcceptedProtocol();
}
catch (Exception ex) {
// Shouldn't happen
logger.error("Failed to obtain session.getAcceptedProtocol(): " +
"will use the default protocol handler (if configured).", ex);
}
SubProtocolHandler handler;
if (StringUtils.hasLength(protocol)) {
handler = this.protocolHandlerLookup.get(protocol);
if (handler == null) {
throw new IllegalStateException(
"No handler for '" + protocol + "' among " + this.protocolHandlerLookup);
}
}
else {
if (this.defaultProtocolHandler != null) {
handler = this.defaultProtocolHandler;
}
else if (this.protocolHandlers.size() == 1) {
handler = this.protocolHandlers.iterator().next();
}
else {
throw new IllegalStateException("Multiple protocol handlers configured and " +
"no protocol was negotiated. Consider configuring a default SubProtocolHandler.");
}
}
return handler;
}
@Nullable
private String resolveSessionId(Message message) {
for (SubProtocolHandler handler : this.protocolHandlerLookup.values()) {
String sessionId = handler.resolveSessionId(message);
if (sessionId != null) {
return sessionId;
}
}
if (this.defaultProtocolHandler != null) {
String sessionId = this.defaultProtocolHandler.resolveSessionId(message);
if (sessionId != null) {
return sessionId;
}
}
return null;
}
/**
* When a session is connected through a higher-level protocol it has a chance
* to use heartbeat management to shut down sessions that are too slow to send
* or receive messages. However, after a WebSocketSession is established and
* before the higher level protocol is fully connected there is a possibility for
* sessions to hang. This method checks and closes any sessions that have been
* connected for more than 60 seconds without having received a single message.
*/
private void checkSessions() {
long currentTime = System.currentTimeMillis();
if (!isRunning() || (currentTime - this.lastSessionCheckTime < getTimeToFirstMessage())) {
return;
}
if (this.sessionCheckLock.tryLock()) {
try {
for (WebSocketSessionHolder holder : this.sessions.values()) {
if (holder.hasHandledMessages()) {
continue;
}
long timeSinceCreated = currentTime - holder.getCreateTime();
if (timeSinceCreated < getTimeToFirstMessage()) {
continue;
}
WebSocketSession session = holder.getSession();
if (logger.isInfoEnabled()) {
logger.info("No messages received after " + timeSinceCreated + " ms. " +
"Closing " + holder.getSession() + ".");
}
try {
this.stats.incrementNoMessagesReceivedCount();
session.close(CloseStatus.SESSION_NOT_RELIABLE);
}
catch (Throwable ex) {
if (logger.isWarnEnabled()) {
logger.warn("Failed to close unreliable " + session, ex);
}
}
}
}
finally {
this.lastSessionCheckTime = currentTime;
this.sessionCheckLock.unlock();
}
}
}
private void clearSession(WebSocketSession session, CloseStatus closeStatus) throws Exception {
if (logger.isDebugEnabled()) {
logger.debug("Clearing session " + session.getId());
}
if (this.sessions.remove(session.getId()) != null) {
this.stats.decrementSessionCount(session);
}
findProtocolHandler(session).afterSessionEnded(session, closeStatus, this.clientInboundChannel);
}
@Override
public String toString() {
return "SubProtocolWebSocketHandler" + this.protocolHandlers;
}
private static class WebSocketSessionHolder {
private final WebSocketSession session;
private final long createTime;
private volatile boolean hasHandledMessages;
public WebSocketSessionHolder(WebSocketSession session) {
this.session = session;
this.createTime = System.currentTimeMillis();
}
public WebSocketSession getSession() {
return this.session;
}
public long getCreateTime() {
return this.createTime;
}
public void setHasHandledMessages() {
this.hasHandledMessages = true;
}
public boolean hasHandledMessages() {
return this.hasHandledMessages;
}
@Override
public String toString() {
return "WebSocketSessionHolder[session=" + this.session + ", createTime=" +
this.createTime + ", hasHandledMessages=" + this.hasHandledMessages + "]";
}
}
/**
* Contract for access to session counters.
* @since 5.2
*/
public interface Stats {
int getTotalSessions();
int getWebSocketSessions();
int getHttpStreamingSessions();
int getHttpPollingSessions();
int getLimitExceededSessions();
int getNoMessagesReceivedSessions();
int getTransportErrorSessions();
}
private class DefaultStats implements Stats {
private final AtomicInteger total = new AtomicInteger();
private final AtomicInteger webSocket = new AtomicInteger();
private final AtomicInteger httpStreaming = new AtomicInteger();
private final AtomicInteger httpPolling = new AtomicInteger();
private final AtomicInteger limitExceeded = new AtomicInteger();
private final AtomicInteger noMessagesReceived = new AtomicInteger();
private final AtomicInteger transportError = new AtomicInteger();
@Override
public int getTotalSessions() {
return this.total.get();
}
@Override
public int getWebSocketSessions() {
return this.webSocket.get();
}
@Override
public int getHttpStreamingSessions() {
return this.httpStreaming.get();
}
@Override
public int getHttpPollingSessions() {
return this.httpPolling.get();
}
@Override
public int getLimitExceededSessions() {
return this.limitExceeded.get();
}
@Override
public int getNoMessagesReceivedSessions() {
return this.noMessagesReceived.get();
}
@Override
public int getTransportErrorSessions() {
return this.transportError.get();
}
void incrementSessionCount(WebSocketSession session) {
getCountFor(session).incrementAndGet();
this.total.incrementAndGet();
}
void decrementSessionCount(WebSocketSession session) {
getCountFor(session).decrementAndGet();
}
void incrementLimitExceededCount() {
this.limitExceeded.incrementAndGet();
}
void incrementNoMessagesReceivedCount() {
this.noMessagesReceived.incrementAndGet();
}
void incrementTransportError() {
this.transportError.incrementAndGet();
}
AtomicInteger getCountFor(WebSocketSession session) {
if (session instanceof PollingSockJsSession) {
return this.httpPolling;
}
else if (session instanceof StreamingSockJsSession) {
return this.httpStreaming;
}
else {
return this.webSocket;
}
}
@Override
public String toString() {
return SubProtocolWebSocketHandler.this.sessions.size() +
" current WS(" + this.webSocket.get() +
")-HttpStream(" + this.httpStreaming.get() +
")-HttpPoll(" + this.httpPolling.get() + "), " +
this.total.get() + " total, " +
(this.limitExceeded.get() + this.noMessagesReceived.get()) + " closed abnormally (" +
this.noMessagesReceived.get() + " connect failure, " +
this.limitExceeded.get() + " send limit, " +
this.transportError.get() + " transport error)";
}
}
}