java中使用sockjs、stomp完成websocket通信

主要配置

import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;

import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketTransportRegistration;



/**
 * @author
 * WebSocket配置类
 */
@Slf4j
@Configuration
@AllArgsConstructor
@EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {

    private final WebSocketInterceptor webSocketInterceptor;
    private final WebSocketHandshakeInterceptor webSocketHandshakeInterceptor;
    public static final String USER_DESTINATION_PREFIX = "/salaryother/";
    public static final String CALL_DEVICE_NOTIFY_PATH = USER_DESTINATION_PREFIX + "CALL_DEVICE_NOTIFY/";

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        log.info("WebSocket服务器注册");
        registry.addEndpoint("/ws")
                .addInterceptors(webSocketHandshakeInterceptor)
                .setHandshakeHandler(new WebSocketHandshakeHandler())
                .setAllowedOrigins("*")
                .withSockJS();

        registry.addEndpoint("/wsapp")
                .addInterceptors(webSocketHandshakeInterceptor)
                .setHandshakeHandler(new WebSocketHandshakeHandler())
                .setAllowedOrigins("*");
    }

    @Override
    public void configureMessageBroker(MessageBrokerRegistry registry) {
        log.info("WebSocket服务器启动");
        //心跳检测
        ThreadPoolTaskScheduler threadPoolTaskScheduler = new ThreadPoolTaskScheduler();
        threadPoolTaskScheduler.setPoolSize(10);
        threadPoolTaskScheduler.setThreadNamePrefix("wss-heartbeat-thread-");
        threadPoolTaskScheduler.initialize();
        ///信息接收头
        registry.enableSimpleBroker("/topic", USER_DESTINATION_PREFIX)
                .setHeartbeatValue(new long[]{10000, 10000}).setTaskScheduler(threadPoolTaskScheduler);
        //接收前缀
        registry.setApplicationDestinationPrefixes("/topic");
        //请求前缀
        registry.setUserDestinationPrefix("/user");
    }

    /**
     * 配置发送与接收的消息参数，可以指定消息字节大小，缓存大小，发送超时时间
     *
     * @param registration
     */
    @Override
    public void configureWebSocketTransport(WebSocketTransportRegistration registration) {
        /*
         * 1. setMessageSizeLimit 设置消息缓存的字节数大小 字节
         * 2. setSendBufferSizeLimit 设置websocket会话时，缓存的大小 字节
         * 3. setSendTimeLimit 设置消息发送会话超时时间，毫秒
         */
        registration.setMessageSizeLimit(10240)
                .setSendBufferSizeLimit(10240)
                .setSendTimeLimit(10000);
    }

    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {

        /*
         * 配置消息线程池
         * 1. corePoolSize 配置核心线程池，当线程数小于此配置时，不管线程中有无空闲的线程，都会产生新线程处理任务
         * 2. maxPoolSize 配置线程池最大数，当线程池数等于此配置时，不会产生新线程
         * 3. keepAliveSeconds 线程池维护线程所允许的空闲时间，单位秒
         */
        registration.taskExecutor().corePoolSize(10)
                .maxPoolSize(60)
                .keepAliveSeconds(60);
        registration.interceptors(webSocketInterceptor);
    }

    //  这个是为了解决和调度任务的冲突重写的bean
    @Primary
    @Bean
    public TaskScheduler taskScheduler() {
        ThreadPoolTaskScheduler taskScheduler = new ThreadPoolTaskScheduler();
        taskScheduler.setPoolSize(10);
        taskScheduler.initialize();
        return taskScheduler;
    }


}

握手拦截(这套方案好像前端无法补充Header,就不在这里做权限校验)这里采用的方法是直接问号拼接token,前端 new SockJS(这里带问号),sockjs使用的是http所以没毛病,本文使用的是OAuth2权限校验

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.joolun.cloud.common.security.entity.BaseUser;
import com.joolun.cloud.common.security.util.SecurityUtils;
import io.micrometer.core.lang.NonNullApi;
import io.micrometer.core.lang.Nullable;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.RemoteTokenServices;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import javax.servlet.http.HttpServletRequest;
import java.util.Map;

@Slf4j
@Component
@NonNullApi
@AllArgsConstructor
public class WebSocketHandshakeInterceptor implements HandshakeInterceptor {
    private WebSocketManager webSocketManager;
    private RemoteTokenServices tokenService;
//SecurityUtils.getUser()代码
//    SecurityContextHolder.getContext().getAuthentication().getPrincipal()
//     principal instanceof BaseUser ? (BaseUser)principal : null;

    @Override
    public boolean beforeHandshake(ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse, WebSocketHandler webSocketHandler, Map<String, Object> map) {
        String token = getToken(serverHttpRequest);
        if (StrUtil.isBlank(token)) return false;
        // 验证令牌信息
        try {
            OAuth2Authentication auth2Authentication = tokenService.loadAuthentication(token);
            if (ObjectUtil.isNull(auth2Authentication)) return false;
            SecurityContextHolder.getContext().setAuthentication(auth2Authentication);
        } catch (Exception e) {
            log.error("token验证失败");
            return false;
        }
        BaseUser user = SecurityUtils.getUser();
        String userId = user.getId();
        map.put("WebSocket-user", new WebSocketUserAuthentication(userId, user.getUsername(), token));
        webSocketManager.addUser(userId, token);
        log.info("userId:" + userId + "用户名:" + user.getUsername() + ":开始建立连接");
        return true;

    }

    @Override
    public void afterHandshake(ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse, WebSocketHandler webSocketHandler, @Nullable Exception e) {
        BaseUser user = SecurityUtils.getUser();
        log.info("userId:" + user.getId() + "用户名:" + user.getUsername() + ":建立连接完成");
    }


    private String getToken(ServerHttpRequest serverHttpRequest) {
        ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) serverHttpRequest;
        HttpServletRequest httpServletRequest = servletRequest.getServletRequest();
        String token = httpServletRequest.getParameter("Authorization");
        return StrUtil.isBlank(token) ? "" : token;
    }

之后可以设置握手之后的身份注入(配置了这个可以在单对单订阅时直接使用)

import io.micrometer.core.lang.NonNullApi;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.support.DefaultHandshakeHandler;

import java.security.Principal;
import java.util.Map;

@NonNullApi
public class WebSocketHandshakeHandler extends DefaultHandshakeHandler {
    @Override
    protected Principal determineUser(ServerHttpRequest request, WebSocketHandler wsHandler, Map<String, Object> attributes) {
        return (Principal) attributes.get("WebSocket-user");
    }


}

import lombok.Data;

import java.security.Principal;

@Data
public class WebSocketUserAuthentication implements Principal {

    private String token;
    private String userId;
    private String userName;

    public WebSocketUserAuthentication(String userId, String userName,String token ) {
        this.token = token;
        this.userId = userId;
        this.userName = userName;
    }

    public WebSocketUserAuthentication() {
    }

    @Override
    public String getName() {
        return token;
    }

}

储存用户数据

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.util.concurrent.TimeUnit;


@Slf4j
@Component
public class WebSocketManager {

    Cache<String, String> webSocketUser;

    @PostConstruct
    public void init() {
        webSocketUser = Caffeine.newBuilder().initialCapacity(16).expireAfterWrite(60, TimeUnit.MINUTES).build();
    }

    public boolean isOnline(String userId) {
        return StringUtils.isNotBlank(webSocketUser.getIfPresent(userId));
    }

    public void addUser(String userId, String token) {
        webSocketUser.put(userId, token);
    }

    public String getTokenById(String userId) {
        return webSocketUser.getIfPresent(userId);
    }

    public void deleteUser(String userId) {
        webSocketUser.invalidate(userId);
    }
}

之后就是通道拦截,如果不使用握手拦截可以在这里鉴权,这里就可以拿到握手后发送的Header,前端就在headers里面添加

this.stompClient.connect(
        headers,
        () => {
        .....

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.joolun.cloud.common.security.entity.BaseUser;
import io.micrometer.core.lang.NonNullApi;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.MessagingException;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.RemoteTokenServices;
import org.springframework.stereotype.Component;

import java.util.Optional;


@Slf4j
@NonNullApi
@Component
@Order(Ordered.HIGHEST_PRECEDENCE + 99)
@AllArgsConstructor
public class WebSocketInterceptor implements ChannelInterceptor {

    private WebSocketManager webSocketManager;
    private RemoteTokenServices tokenService;


    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
        if (ObjectUtil.isNull(accessor)) throw new MessagingException("获取数据失败");
        StompCommand stompCommand = accessor.getCommand();
        // 判断是否是连接或者断开请求
        if (StompCommand.CONNECT != stompCommand && StompCommand.DISCONNECT != stompCommand) return message;
        //拿取用户信息
        Optional<WebSocketUserAuthentication> user = getUser(accessor);
        if (!user.isPresent()) throw new MessagingException("获取用户失败");
        WebSocketUserAuthentication baseUser = user.get();
        String userId = baseUser.getUserId();
        // 下线请求
        if (StompCommand.DISCONNECT == stompCommand) {
            String userName = baseUser.getUserName();
            webSocketManager.deleteUser(userId);
            log.info("userId:" + userId + "用户名:" + userName + ":下线了");
        }
        return message;
    }


    /**
     * 在消息发送后立刻调用，boolean值参数表示该调用的返回值
     */
    @Override
    public void postSend(Message<?> message, MessageChannel messageChannel, boolean sent) {
        StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
        // 忽略心跳消息等非STOMP消息
        StompCommand command = accessor.getCommand();
        if (command == null) return;
        switch (command) {
            case CONNECT:
                log.info("上线了");
                break;
            case CONNECTED:
                log.info("已连接");
                break;
            case SUBSCRIBE:
                log.info("订阅:" + accessor.getDestination());
                break;
            case UNSUBSCRIBE:
                log.info("取消订阅:" + accessor.getDestination());
                break;
            case DISCONNECT:
                log.info("下线了");
                break;
            default:
                log.info("不匹配以上情况");
                break;
        }
    }

    private boolean isUserOAuth2Authentication(StompHeaderAccessor accessor) {
        String token = getToken(accessor);
        if (StrUtil.isBlank(token)) return false;
        try {
            // 验证令牌信息
            OAuth2Authentication auth2Authentication = tokenService.loadAuthentication(token);
            if (ObjectUtil.isNull(auth2Authentication)) return false;
            SecurityContextHolder.getContext().setAuthentication(auth2Authentication);
        } catch (Exception e) {
            log.error("token验证失败");
            return false;
        }
        return true;
    }

    private Optional<WebSocketUserAuthentication> getUser(StompHeaderAccessor accessor) {
        return accessor.getUser() instanceof WebSocketUserAuthentication ?
                Optional.of((WebSocketUserAuthentication) accessor.getUser()) :
                getSystemUserToWebSocketUserAuthentication(accessor)
                ;


    }

    private Optional<WebSocketUserAuthentication> getSystemUserToWebSocketUserAuthentication(StompHeaderAccessor accessor) {
        Authentication authentication = getAuthentication();
        if (ObjectUtil.isNull(authentication)) {
            if (isUserOAuth2Authentication(accessor)) {
                authentication = getAuthentication();
            } else {
                return Optional.empty();
            }
        }
        Object principal = authentication.getPrincipal();
        if (ObjectUtil.isNull(principal)) return Optional.empty();
        BaseUser user = principal instanceof BaseUser ? (BaseUser) principal : null;
        if (ObjectUtil.isNotNull(user)) {
            WebSocketUserAuthentication webSocketUserAuthentication = new WebSocketUserAuthentication(user.getId(), user.getUsername(), getToken(accessor));
            accessor.setUser(webSocketUserAuthentication);
            return Optional.of(webSocketUserAuthentication);
        }
        return Optional.empty();

    }

    private String getToken(StompHeaderAccessor accessor) {
        String tokens = accessor.getFirstNativeHeader("Authorization");
        if (StrUtil.isBlank(tokens)) return "";
        return tokens.split(" ")[1];
    }

    private Authentication getAuthentication() {
        return SecurityContextHolder.getContext().getAuthentication();
    }

}

配置完成之后就是发送消息这里就不举详细的例子了就拿SimpMessageSendingOperations来说

//引入
import org.springframework.messaging.simp.SimpMessageSendingOperations;

private final SimpMessageSendingOperations simpMessageSendingOperations;
private final WebSocketManager webSocketManager;

如果配置了握手拦截器返回了Principal 个人消息

订阅地址:/user/salaryother/activistIncoming
发送地址:/salaryother/activistIncoming
发送消息:
 simpMessageSendingOperations.convertAndSendToUser
                    (webSocketManager.getTokenById(用户id), 发送地址, 消息);

如果没配置那就得多加几个前缀具体参考请点击:
Spring Springboot实现websocket通讯-2 这个详细

微信小程序连接websocket