springboot+websocket+token验证+jedis支持集群部署发消息

目录

websocket主要代码

pom.xml

MyWebSocketHandler

WSInterceptor

WebSocketConfig

Redis订阅广播实现Session共享

pom.xml

WebSocketSub

WebSocketSubscriber

服务端发送消息

JedisUtil

---

websocket主要代码

pom.xml

            org.springframework.boot
            spring-boot-starter-websocket

MyWebSocketHandler

通过继承 TextWebSocketHandler 类并覆盖相应方法，可以对 websocket 的事件进行处理，这里可以同原生注解的那几个注解连起来看

1. afterConnectionEstablished 方法是在 socket 连接成功后被触发，同原生注解里的 @OnOpen 功能
2. afterConnectionClosed 方法是在 socket 连接关闭后被触发，同原生注解里的 @OnClose 功能
3. handleTextMessage 方法是在客户端发送信息时触发，同原生注解里的 @OnMessage 功能

ConcurrentHashMap实现本地session池，用来保存已经登录的websocket的session。服务端发送消息给客户端必须要通过这个session。Map的key根据你的需求来指定，我们这个系统可能一个账号多个地方登陆使用所以没有用userId来做key，而是用的token。

@Slf4j
@Component
public class MyWebSocketHandler extends TextWebSocketHandler {

    private static final ConcurrentHashMap SESSION_POOL = new ConcurrentHashMap<>();

    /**
     * socket 建立成功事件
     * @param session
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        Object token = session.getAttributes().get("token");
        if (token != null) {
            // 用户连接成功，放入在线用户缓存
            SESSION_POOL.put(token.toString(), session);
        } else {
            throw new RuntimeException("用户登录已经失效!");
        }
    }

    /**
     * 接收消息事件
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // 获得客户端传来的消息
        String payload = message.getPayload();
        Object token = session.getAttributes().get("token");
        Object userId = session.getAttributes().get("userId");
        log.info("server 接收到 {}:{} 发送的 {}", userId, token, payload);
        session.sendMessage(new TextMessage("server 发送给 " + token + " 消息 " + payload + " " + LocalDateTime.now().toString()));
    }

    /**
     * socket 断开连接时
     * @param session
     * @param status
     * @throws Exception
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        Object token = session.getAttributes().get("token");
        if (token != null) {
            // 用户退出，移除缓存
            SESSION_POOL.remove(token.toString());
        }
    }

    /**
     * 给某个用户(token)发送消息
     * @param token
     * @param message
     */
    public void send(String token, String message) {
        WebSocketSession session = SESSION_POOL.get(token);
        if (session != null) {
            try {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage(message));
                }
            } catch (IOException e) {
                log.error("发送消息给{}:{}失败", session.getAttributes().get("userId"), token, e);
            }
        }
    }

    public void sendAll(String message) {
        for (String token : SESSION_POOL.keySet()) {
            send(token, message);
        }
    }
}

WSInterceptor

通过实现 HandshakeInterceptor 接口来定义握手拦截器，注意这里与上面 Handler 的事件是不同的，这里是建立握手时的事件，分为握手前与握手后，而  Handler 的事件是在握手成功后的基础上建立 socket 的连接。所以在如果把认证放在这个步骤相对来说最节省服务器资源。它主要有两个方法 beforeHandshake(握手前触发) 与 afterHandshake(握手后触发)。

测试的时候是用的postman做的前端websocket请求，token放到header里了。但是前端同事说找不到增加header的方法，后面就改成了把token放到参数里。

@Slf4j
@Component
public class WSInterceptor implements HandshakeInterceptor {
    @Autowired
    JedisUtil jedisUtil;

    /**
     * 握手前
     * @param request
     * @param response
     * @param wsHandler
     * @param attributes
     * @return
     * @throws Exception
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map attributes) throws Exception {
        // 获得请求参数
        String authToken = ((ServletServerHttpRequest) request).getServletRequest().getParameter("token");
        if (!StringUtils.isEmpty(authToken) && jedisUtil.exists(TokenConstants.USER_KEY + authToken)) {
            String json = jedisUtil.get(TokenConstants.USER_KEY + authToken);
            if (!StringUtils.isEmpty(json)) {
                SecurityUserCommon user = JsonUtil.jsonToPojo(json, SecurityUserCommon.class);
                // 放入属性域
                attributes.put("token", authToken);
                attributes.put("userId", user.getId());
                log.info("用户{}握手成功", user.getId());
                return true;
            }
        }
        log.info("握手失败，用户登录已失效");
        return false;
    }

    /**
     * 握手后
     * @param request
     * @param response
     * @param wsHandler
     * @param exception
     */
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
        log.info("握手完成");
    }
}

WebSocketConfig

通过实现 WebSocketConfigurer 类并覆盖相应的方法进行 websocket 的配置。我们主要覆盖 registerWebSocketHandlers 这个方法。通过向 WebSocketHandlerRegistry 设置不同参数来进行配置。其中 addHandler 方法添加我们上面的写的 ws 的  handler 处理类，第二个参数是你暴露出的 ws 路径。addInterceptors 添加我们写的握手过滤器。setAllowedOrigins("*") 这个是关闭跨域校验，方便本地调试，线上推荐打开。 

@Configuration
@EnableWebSocket
public class WebSocketConfig  implements WebSocketConfigurer {
    @Autowired
    private MyWebSocketHandler myWebSocketHandler;
    @Autowired
    private WSInterceptor wsInterceptor;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry
                .addHandler(myWebSocketHandler, "/wsserver")
                .addInterceptors(wsInterceptor)
                .setAllowedOrigins("*");
    }
}

Redis订阅广播实现Session共享

服务端如果要主动发送消息给客户端一定要用到 session。而大家都知道的是 session 这个东西是不跨 jvm 的。如果有多台服务器，在 http 请求的情况下，我们可以通过把 session 放入缓存中间件中来共享解决这个问题，通过 spring session 几条配置就解决了。但是 websocket 不可以。他的 session 是不能序列化的，当然这样设计的目的不是为了为难你，而是出于对 http与websocket 请求的差异导致的。
目前网上找到的最简单方案就是通过 redis 订阅广播的形式，在本地放个 map 保存请求的 session。也就是说每台服务器都会保存与他连接的 session 于本地。服务器要发消息的时候，你通过 redis 广播这条消息，所有订阅的服务端都会收到这个消息，然后本地尝试发送。最后肯定只有有这个对应用户 session 的那台才能发送出去。

用redisTemplate做订阅广播的网上有很多，本文用的是Jedis

pom.xml

        
            redis.clients
            jedis
            2.8.2
        

WebSocketSub

继承自JedisPubSub，用于处理订阅相关事件。在收到订阅消息后，调用handler发送消息，token就是上面提到的存储Session的Map里的key，用于找到实际的Session发送消息。

@Slf4j
@Component
public class WebSocketSub extends JedisPubSub {
    @Autowired
    private MyWebSocketHandler myWebSocketHandler;

    // 取得订阅的消息后的处理
    @Override
    public void onMessage(String channel, String message) {
        log.info("订阅成功,接收到的消息为:频道-{},消息-{}", channel, message);
        JSONObject json = JSONObject.parseObject(message);
        myWebSocketHandler.send(json.get("token").toString(), json.get("msg").toString());
    }

    // 初始化订阅时候的处理
    @Override
    public void onSubscribe(String channel, int subscribedChannels) {
        log.info("初始化订阅信息:频道-{},订阅频道-{}", channel, subscribedChannels);
    }

    // 取消订阅时候的处理
    @Override
    public void onUnsubscribe(String channel, int subscribedChannels) {
        log.info("已取消订阅频道{}", channel);
    }
}

WebSocketSubscriber

继承自ApplicationRunner，系统启动完成后启动一个线程订阅redis websocket的频道

@Slf4j
@Component
public class WebSocketSubscriber implements ApplicationRunner {
    @Autowired
    private JedisUtil jedisUtil;
    @Autowired
    private WebSocketSub webSocketSub;

    @Override
    public void run(ApplicationArguments args) throws Exception {
        //起线程订阅redis ws频道
        ExecutorService cachedThreadPool = Executors.newCachedThreadPool();
        cachedThreadPool.execute(() -> {
            jedisUtil.subscribe(webSocketSub, RedisConstant.CHANNEL_WS);
        });
    }
}

服务端发送消息

就是通过redis发布消息到websocket的频道

JSONObject json = new JSONObject();
json.put("token", token);
json.put("msg", msg);
jedisUtil.publish(RedisConstant.CHANNEL_WS, json.toJSONString());

JedisUtil

就是一个jedis的工具类，只展示出这里用到的publish和subscribe方法

@Component
@Slf4j
public class JedisUtil {

    @Autowired
    private JedisPool jedisPool;

    @Value("${spring.redis.database}")
    private int indexdb;

    /**
     * 发布消息
     * @param channel
     * @param message
     * @return
     */
    public Long publish(String channel, String message) {
        Jedis jedis = null;
        Long count = null;
        try {
            jedis = jedisPool.getResource();
            jedis.select(indexdb);
            count = jedis.publish(channel, message);//返回订阅者数量
        } catch (Exception e) {
            log.error(e.getMessage());
        } finally {
            if (jedis != null) {
                jedis.close();
            }
        }
        return count;
    }

    /**
     * 订阅频道
     * @param jedisPubSub
     * @param channel
     */
    public void subscribe(JedisPubSub jedisPubSub, String channel) {
        Jedis jedis = null;
        try {
            jedis = jedisPool.getResource();
            jedis.select(indexdb);
            jedis.subscribe(jedisPubSub, channel);
        } catch (Exception e) {
            log.error(e.getMessage());
        } finally {
            if (jedis != null) {
                jedis.close();
            }
        }
    }
}

参考：

【websocket】spring boot 集成 websocket 的四种方式 - KIWI的碎碎念 - 博客园

Redis 订阅发布 - Jedis实现 - 叶云轩 - 博客园