WebSocket 初体验

import com.hp.water.websocket.HttpAuthHandler;
import com.hp.water.websocket.WebsocketInterceptor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

/**
 * @author 范天予
 * @data 2022/7/7  10:27
 */
@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
    @Autowired
    private HttpAuthHandler httpAuthHandler;
    @Autowired
    private WebsocketInterceptor myInterceptor;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry
                //自定义的websocket服务,这里都可以定义多个
                .addHandler(httpAuthHandler, "websocket")
                //设置拦截器
                .addInterceptors(myInterceptor)
                //关闭跨域校验
                .setAllowedOrigins("*");
    }

}

import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author 范天予
 * @data 2022/7/7  11:12
 */
@Slf4j
public class WsSessionManager {
    /**
     * 保存连接 session 的地方
     */
    private static ConcurrentHashMap<String, WebSocketSession> SESSION_POOL = new ConcurrentHashMap<>();

    /**
     * 添加 session
     * @param key
     */
    public static void add(String key, WebSocketSession session) {
        // 添加 session
        SESSION_POOL.put(key, session);
    }

    /**
     * 删除 session,会返回删除的 session
     * @param key
     * @return
     */
    public static WebSocketSession remove(String key) {
        // 删除 session
        return SESSION_POOL.remove(key);
    }

    /**
     * 删除并同步关闭连接
     * @param key
     */
    public static void removeAndClose(String key) {
        WebSocketSession session = remove(key);
        if (session != null) {
            try {
                // 关闭连接
                session.close();
            } catch (IOException e) {
                // todo: 关闭出现异常处理
                e.printStackTrace();
            }
        }
    }
    /**
     * 获得 session
     * @param key
     * @return
     */
    public static WebSocketSession get(String key) {
        // 获得 session
        return SESSION_POOL.get(key);
    }
}

import cn.hutool.core.date.DateTime;
import cn.hutool.json.JSONUtil;
import com.hp.water.config.WebSocketEntity;
import com.hp.water.config.WsSessionManager;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

/**
 * @author 范天予
 * @data 2022/7/7  11:13
 */
@Component
public class HttpAuthHandler extends TextWebSocketHandler {
    int i = 0;

    /**
     * socket 建立成功事件
     *
     * @param session
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        //这里的值在拦截器中的域属性中复制，后面会自动添加进去
        Object token = session.getAttributes().get("token");
        i++;
        System.out.println("总数量"+i);
        if (token != null) {
            // 用户连接成功，放入在线用户缓存
            WsSessionManager.add(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");
        System.out.println("server 接收到 " + token + " 发送的 " + payload);

        WebSocketEntity webSocketEntity = new WebSocketEntity();
        webSocketEntity.setCode(0);
        webSocketEntity.setSno("123321123456789");
        webSocketEntity.setTime(DateTime.now());
        webSocketEntity.setData(null);
        webSocketEntity.setType(1);
        webSocketEntity.setData(JSONUtil.toJsonStr(webSocketEntity));




        String tojson = JSONUtil.toJsonStr(webSocketEntity);

        session.sendMessage(new TextMessage(tojson));

    }

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

import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HttpUtil;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Map;

/**
 * @author 范天予
 * @data 2022/7/7  11:14
 */
@Component
public class WebsocketInterceptor implements HandshakeInterceptor {
    /**
     * 握手前
     * @param request
     * @param response
     * @param wsHandler
     * @param attributes
     * @return
     * @throws Exception
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        System.out.println("握手开始");
        // 获得请求参数,这里用了hutools工具箱
        HashMap<String, String> paramMap = (HashMap<String, String>) HttpUtil.decodeParamMap(request.getURI().getQuery(), Charset.defaultCharset());
        String uid = paramMap.get("token");
        String sno = paramMap.get("sno");
        System.out.println(uid);
        System.out.println(sno);
        if (StrUtil.isNotBlank(uid)) {
            // 放入属性域，可以在HttpAuthHandler的session的attributes中获取
            attributes.put("token", uid);
            System.out.println("用户 token :" + uid + " 握手成功！");
            if(StrUtil.isNotBlank(sno)){
                attributes.put("sno", sno);
                System.out.println("用户 son :" + sno + " 握手成功！");
            }
            return true;
        }
        System.out.println("用户登录已失效");
        return false;
    }
    /**
     * 握手后
     * @param request
     * @param response
     * @param wsHandler
     * @param exception
     */
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
        System.out.println("握手完成");
    }

}