import java.util.ArrayList; import java.util.concurrent.atomic.AtomicInteger; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.web.bind.annotation.RestController; import org.springframework.web.socket.CloseStatus; import org.springframework.web.socket.WebSocketHandler; import org.springframework.web.socket.WebSocketMessage; import org.springframework.web.socket.WebSocketSession; @RestController public class TestWebSocketController implements WebSocketHandler { private static AtomicInteger onlineCount = new AtomicInteger(0); private static final ArrayList<WebSocketSession> sessions = new ArrayList<>(); private final Logger LOGGER = LoggerFactory.getLogger(TestWebSocketController.class); @Override public void afterConnectionEstablished(WebSocketSession session) throws Exception { sessions.add(session); int onlineNum = addOnlineCount(); LOGGER.info("Oprn a WebSocket. Current connection number: " + onlineNum); } @Override public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception { sessions.remove(session); int onlineNum = subOnlineCount(); LOGGER.info("Close a webSocket. Current connection number: " + onlineNum); } @Override public void handleMessage(WebSocketSession wsSession, WebSocketMessage<?> message) throws Exception { LOGGER.info("Receive a message from client: " + message.toString()); } @Override public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception { LOGGER.error("Exception occurs on webSocket connection. disconnecting...."); if (session.isOpen()) { session.close(); } sessions.remove(session); subOnlineCount(); } /* * 是否支持消息拆分发送：如果接收的数据量比较大，最好打开(true), 否则可能会导致接收失败。 * 如果出现WebSocket连接接收一次数据后就自动断开，应检查是否是这里的问题。 */ @Override public boolean supportsPartialMessages() { return true; } public static int getOnlineCount() { return onlineCount.get(); } public static int addOnlineCount() { return onlineCount.incrementAndGet(); } public static int subOnlineCount() { return onlineCount.decrementAndGet(); } }
（2）HandShake 拦截器实现
import java.util.Map; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.context.ApplicationContext; import org.springframework.http.server.ServerHttpRequest; import org.springframework.http.server.ServerHttpResponse; import org.springframework.http.server.ServletServerHttpRequest; import org.springframework.web.socket.WebSocketHandler; import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor; public class TestHandShakeInterceptor extends HttpSessionHandshakeInterceptor { private final Logger LOGGER = LoggerFactory.getLogger(TestHandShakeInterceptor.class); /* * 在WebSocket连接建立之前的操作，以鉴权为例 */ @Override public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception { LOGGER.info("Handle before webSocket connected. "); // 获取url传递的参数，通过attributes在Interceptor处理结束后传递给WebSocketHandler // WebSocketHandler可以通过WebSocketSession的getAttributes()方法获取参数 ServletServerHttpRequest serverRequest = (ServletServerHttpRequest) request; String id = serverRequest.getServletRequest().getParameter("id"); String name = serverRequest.getServletRequest().getParameter("name"); if (tokenValidation.validateSign()) { LOGGER.info("Validation passed. WebSocket connecting.... "); attributes.put("id", id); attributes.put("name", name); return super.beforeHandshake(request, response, wsHandler, attributes); } else { LOGGER.error("Validation failed. WebSocket will not connect. "); return false; } } @Override public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception ex) { // 省略 } }
（3）WebSocket 配置类实现（注册WebSocket实现类，绑定接口，同时将实现类和拦截器绑定）