记一次Spring boot使用stomp协议栈时从服务端发起关闭
前言
这篇文章是在开发过程中发生的问题，会主要根据本人在本次解决问题的角度进行分析。面向的是一个即时通信项目，与客户端使用websocket做连接接口，使用spring boot的stomp协议栈进行通信。即如下代码形式：
@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {
@Override
public void registerStompEndpoints(StompEndpointRegistry registry) {
registry.addEndpoint("/im/ws").setAllowedOrigins("*").withSockJS();
}
...
}
在分析过程中发现stomp是一个闭环，很多东西都是私有的，可能人家并不想让使用者从服务端发起关闭吧。但是如果客户端进行了订阅，虽然可以通过拦截器的方法拒绝他的消息，但却无法做到不向该通道进行消息推送，所以需要从服务端主动断掉推送思路有二，一种是通过是通过推送关闭消息让客服端发起关闭，第二种webSocketSession直接直接close进行关闭。
解决思路
我们配置项目支持stomp是通过@EnableWebSocketMessageBroker标签进行配置的，追踪过去发现它引入了DelegatingWebSocketMessageBrokerConfiguration.class。
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.TYPE)
@Documented
@Import(DelegatingWebSocketMessageBrokerConfiguration.class)
public @interface EnableWebSocketMessageBroker {
}
继续追踪我们发现配置类WebSocketMessageBrokerConfigurationSupport.class中有个subProtocolWebSocketHandler被声明。
public abstract class WebSocketMessageBrokerConfigurationSupport extends AbstractMessageBrokerConfiguration {
...
@Bean
public WebSocketHandler subProtocolWebSocketHandler() {
return new SubProtocolWebSocketHandler(clientInboundChannel(), clientOutboundChannel());
}
...
}
继续分析里面有个函数 handleMessage 是向客户端推送消息的，并且是通过sessionId进行精确推送的。
public class SubProtocolWebSocketHandler
implements WebSocketHandler, SubProtocolCapable, MessageHandler, SmartLifecycle {
....
/**
* Handle an outbound Spring Message to a WebSocket client.
*/
@Override
public void handleMessage(Message<?> message) throws MessagingException {
...
}
...
}
所以我们可以通过@Resource(name = “subProtocolWebSocketHandler”)进行注入获取SubProtocolWebSocketHandler进行操作。因StompCommand中供服务端使用只有CONNECTED（客户端发起连接时回复连接成功）、RECEIPT（接收消息回复接收成功）、MESSAGE（消息推送）、ERROR（错误）可以使用，并且可以从客户端主动发起的只有MESSAGE和ERROR可以使用，通过ERROR可以知道，推送至客户端后，stomp协议就会被动发起关闭流程。以下是一种演示，当然也可以通过MESSAGE的方式自己写逻辑进行主动关闭。
@Controller
public class WSController {
...
@Resource(name = "subProtocolWebSocketHandler")
private SubProtocolWebSocketHandler subProtocolWebSocketHandler;
// 这里只是一个demo程序，我是通过http请求的方式控制关闭具体流程和安全性就要自己掌握了
@RequestMapping("/closeAllSession")
@ResponseBody
public String closeAllSession() {
// 此处只做演示，具体的用户信息和sessionId的对应可以通过configureClientInboundChannel进行映射
String[] array = WebSocketConfig.sessionSet.toArray(new String[0]);
WebSocketConfig.sessionSet.clear();
StompHeaderAccessor headerAccessor = StompHeaderAccessor.create(StompCommand.ERROR);
headerAccessor.setSessionId(array[0]);
Message<byte[]> createMessage = MessageBuilder.createMessage(new byte[0], headerAccessor.getMessageHeaders());
subProtocolWebSocketHandler.handleMessage(createMessage);
return "success";
}
...
}
至此一个简单的从服务端发起关闭stomp接口也就完成了。
但是。。。
流程上可以看出，发起关闭的是stomp协议分析后发起DISCONNECT消息进行关闭，也就是说如果这个流程被恶意劫持，那么这个连接还是不能被关闭的，安全性上还是欠缺一些，所以我们的第二种方法通过webSocketSession进行关闭。查看了其他相关的方法，比较合适做session处理的还是在 SubProtocolWebSocketHandler 中；我们继续分析发现sessions是private修饰的，并且没有对外的获取map或根据sessionId获取session的方法，所以这里只能pass掉；继续分析在afterConnectionEstablished函数中调用了decorateSession方法中对session进行自定义封装，并且decorateSession是protected修饰的。
public class SubProtocolWebSocketHandler
implements WebSocketHandler, SubProtocolCapable, MessageHandler, SmartLifecycle {
...
@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);
}
....
protected WebSocketSession decorateSession(WebSocketSession session) {
return new ConcurrentWebSocketSessionDecorator(session, getSendTimeLimit(), getSendBufferSizeLimit());
}
...
}
所以我们可以通过继承 SubProtocolWebSocketHandler 重写decorateSession的方式拿到webSocketSession并自己维护。
我们要用自定义的 CustomSubProtocolWebSocketHandler 则需要重写 DelegatingWebSocketMessageBrokerConfiguration 配置类。又因为我们自己实现了 DelegatingWebSocketMessageBrokerConfiguration 所以前面的@EnableWebSocketMessageBroker也就要去掉了。
@Configuration
public class CustomWebSocketMessageBrokerConfiguration extends DelegatingWebSocketMessageBrokerConfiguration {
@Bean
public CustomSubProtocolWebSocketHandler customSubProtocolWebSocketHandler() {
return new CustomSubProtocolWebSocketHandler(clientInboundChannel(), clientOutboundChannel());
}
@Override
public WebSocketHandler subProtocolWebSocketHandler() {
return customSubProtocolWebSocketHandler();
}
}
websocket通道处理器 继承 SubProtocolWebSocketHandler 的实现如下：
/**
* 自定义的websocket通道处理器
*/
public class CustomSubProtocolWebSocketHandler extends SubProtocolWebSocketHandler {
protected Map<String, Set<String>> userSessionIdMap = new HashMap<>();
protected Map<String, CustomWebSocketSessionDecorator> sessionMap = new HashMap<>();
private final Lock mapLock = new ReentrantLock();
@Autowired
private WsuserServer wsuserServer;
public CustomSubProtocolWebSocketHandler(MessageChannel clientInboundChannel,
SubscribableChannel clientOutboundChannel) {
super(clientInboundChannel, clientOutboundChannel);
}
@Override
protected WebSocketSession decorateSession(WebSocketSession session) {
CustomWebSocketSessionDecorator decorator = new CustomWebSocketSessionDecorator(session, getSendTimeLimit(),
getSendBufferSizeLimit(), this);
sessionMap.put(session.getId(), decorator);
return decorator;
}
/**
* 设置用户信息
*
* @param sessionId
* @param userInfo
*/
public void setUserInfo(String sessionId, WebsocketUserInfo userInfo) {
String userUnoinAddress = WebsocketFormat.userUnoinAddress(userInfo);
mapLock.lock();
try {
CustomWebSocketSessionDecorator decorator = sessionMap.get(sessionId);
Assert.notNull(decorator, String.format("sessionId [%s] 的decorator为空", sessionId));
decorator.setPrincipal(userInfo);
Set<String> set = userSessionIdMap.get(userUnoinAddress);
if (set == null) {
set = new HashSet<>();
userSessionIdMap.put(userUnoinAddress, set);
}
set.add(sessionId);
} finally {
mapLock.unlock();
}
}
@Override
public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
super.afterConnectionClosed(session, closeStatus);
this.sessionClear(session, closeStatus);
}
/**
* 清理session
*
* @param session
* @param closeStatus
*/
public void sessionClear(WebSocketSession session, CloseStatus closeStatus) {
// 关闭websocket session
mapLock.lock();
try {
CustomWebSocketSessionDecorator remove = sessionMap.remove(session.getId());
if (remove != null) {
WebsocketUserInfo userInfo = (WebsocketUserInfo) remove.getPrincipal();
if (userInfo != null) {
String userUnoinAddress = WebsocketFormat.userUnoinAddress(userInfo);
Set<String> set = userSessionIdMap.get(userUnoinAddress);
if (set != null) {
set.remove(session.getId());
if (set.isEmpty()) {
userSessionIdMap.remove(userUnoinAddress);
}
}
}
}
} finally {
mapLock.unlock();
}
}
/**
* 关闭指定用户
*
* @param userInfo
*/
public void closeByUserInfo(WebsocketUserInfo userInfo) {
List<WebSocketSession> sessionList = new LinkedList<>();
String userUnoinAddress = WebsocketFormat.userUnoinAddress(userInfo);
Assert.isTrue(LockActuator.trylock(userUnoinAddress), String.format("锁定[%s]失败", userUnoinAddress));
try {
WebsocketUserInfoWrapper wrapper = WebsocketUserInfoWrapper.of(userInfo);
WebsocketUserInfo cacheUserInfo = wsuserServer.getUserInfo(wrapper);
mapLock.lock();
try {
Set<String> set = userSessionIdMap.get(userUnoinAddress);
if (set != null) {
Iterator<String> iterator = set.iterator();
while (iterator.hasNext()) {
String next = iterator.next();
CustomWebSocketSessionDecorator customWebSocketSessionDecorator = sessionMap.get(next);
if (customWebSocketSessionDecorator == null) {
// 此处是将那些无关联的进行清理，防止内存泄漏，从目前逻辑上来看应该不会发生，做个保险
iterator.remove();
continue;
}
WebsocketUserInfo tempUserInfo = (WebsocketUserInfo) customWebSocketSessionDecorator
.getPrincipal();
if (tempUserInfo == null || !tempUserInfo.sameUser(cacheUserInfo)) {
// 此处只判断哪些需要关闭，不清理sessionMap和userSessionIdMap，交给close逻辑处理
sessionList.add(customWebSocketSessionDecorator);
continue;
}
}
if (set.isEmpty()) {
userSessionIdMap.remove(userUnoinAddress);
}
}
} finally {
mapLock.unlock();
}
} finally {
LockActuator.releaseLock();
}
for (WebSocketSession session : sessionList) {
try {
session.close();
} catch (Exception e) {
// do nothing
}
}
}
}
CustomWebSocketSessionDecorator的实现如下：
/**
* 自定义的websocket控制器
*/
public class CustomWebSocketSessionDecorator extends ConcurrentWebSocketSessionDecorator {
protected CustomSubProtocolWebSocketHandler webSocketHandler;
protected Principal user;
public CustomWebSocketSessionDecorator(WebSocketSession delegate, int sendTimeLimit, int bufferSizeLimit,
CustomSubProtocolWebSocketHandler webSocketHandler) {
super(delegate, sendTimeLimit, bufferSizeLimit);
this.webSocketHandler = webSocketHandler;
}
public CustomWebSocketSessionDecorator(WebSocketSession delegate, int sendTimeLimit, int bufferSizeLimit,
OverflowStrategy overflowStrategy, CustomSubProtocolWebSocketHandler webSocketHandler) {
super(delegate, sendTimeLimit, bufferSizeLimit, overflowStrategy);
this.webSocketHandler = webSocketHandler;
}
public void setPrincipal(Principal user) {
this.user = user;
}
@Override
public Principal getPrincipal() {
return this.user;
}
@Override
public void close() throws IOException {
this.close(CloseStatus.GOING_AWAY);
}
@Override
public void close(CloseStatus status) throws IOException {
super.close(status);
webSocketHandler.sessionClear(this, status);
}
}
上面CustomWebSocketSessionDecorator的实现中有个setPrincipal，需要注意的是stomp中没有直接将 用户信息 设置到webSocketSession中，而是维护在了 StompSubProtocolHandler.class 中，所以如果我们想让用户信息注入到session中还要主动执行一下保存。
结语
以上是使用SpringBoot封装的stomp消息的情况下从服务端发起关闭连接的两种思路。通过分析来看，stomp更想自己维护session，形成一个闭环，毕竟stomp是一个协议栈，并不只是应用在webSocket上，所以这里的方案也只是适用webSocket的通信方案上。
写了一个相关的demo，可以更好的展示出第二种解决方案：
https://github.com/Pluto-Whong/stomp-demo.git