websocket基础

pom.xml

@@ -34,6 +34,10 @@
             <groupId>com.kingtom.kirin</groupId>
             <artifactId>kirin_api_core</artifactId>
         </dependency>
+        <dependency>
+            <groupId>org.springframework.boot</groupId>
+            <artifactId>spring-boot-starter-websocket</artifactId>
+        </dependency>
     </dependencies>
 
     <build>

src/main/java/com/kingtom/shengtai/app/message/ws/WebSocketConfig.java

@@ -0,0 +1,37 @@
+package com.kingtom.shengtai.app.message.ws;
+
+import org.springframework.context.annotation.Bean;
+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;
+
+/**
+ * 应用模块名称</p>
+ * 代码描述</p>
+ * Copyright: Copyright (C) 2023 , Inc. All rights reserved. <p>
+ * Company: 成都诚唐科技有限责任公司</p>
+ *
+ * @author wany
+ * @since 2023/12/26
+ */
+@EnableWebSocket
+@Configuration
+public class WebSocketConfig implements WebSocketConfigurer{
+
+    @Bean
+    public WsHandler handler(){
+        return new WsHandler();
+    }
+
+    @Bean
+    public WsHandshakeInterceptor interceptor(){
+        return new WsHandshakeInterceptor();
+    }
+
+    @Override
+    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry){
+        registry.addHandler(handler(), "/ws").setAllowedOrigins("*").addInterceptors(interceptor());
+    }
+
+}

src/main/java/com/kingtom/shengtai/app/message/ws/WsHandler.java

@@ -0,0 +1,61 @@
+package com.kingtom.shengtai.app.message.ws;
+
+import java.time.LocalDateTime;
+
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.web.socket.BinaryMessage;
+import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.TextMessage;
+import org.springframework.web.socket.WebSocketSession;
+import org.springframework.web.socket.handler.AbstractWebSocketHandler;
+
+/**
+ * 应用模块名称</p>
+ * 代码描述</p>
+ * Copyright: Copyright (C) 2023 , Inc. All rights reserved. <p>
+ * Company: 成都诚唐科技有限责任公司</p>
+ *
+ * @author wany
+ * @since 2023/12/26
+ */
+@Slf4j
+public class WsHandler extends AbstractWebSocketHandler{
+
+    @Override
+    public void afterConnectionEstablished(WebSocketSession session) throws Exception{
+        log.info("建立ws连接");
+        String key = session.getAttributes().get("userId").toString();
+        WsSessionManager.add(key, session);
+    }
+
+    @Override
+    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception{
+        log.info("发送文本消息");
+        // 获得客户端传来的消息
+        String payload = message.getPayload();
+        log.info("server 接收到消息 " + payload);
+        session.sendMessage(new TextMessage("server 发送给的消息 " + payload + "，发送时间:" + LocalDateTime.now()));
+    }
+
+    @Override
+    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception{
+        log.info("发送二进制消息");
+        session.sendMessage(message);
+    }
+
+    @Override
+    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception{
+        log.error("异常处理");
+        String key = session.getAttributes().get("userId").toString();
+        WsSessionManager.removeAndClose(key);
+    }
+
+    @Override
+    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception{
+        log.info("关闭ws连接");
+        String key = session.getAttributes().get("userId").toString();
+        WsSessionManager.removeAndClose(key);
+    }
+
+}
+

src/main/java/com/kingtom/shengtai/app/message/ws/WsHandshakeInterceptor.java

@@ -0,0 +1,44 @@
+package com.kingtom.shengtai.app.message.ws;
+
+import java.util.Map;
+
+import com.kingtom.kirin.api.security.SaSecurityUtils;
+import org.springframework.http.server.ServerHttpRequest;
+import org.springframework.http.server.ServerHttpResponse;
+import org.springframework.http.server.ServletServerHttpRequest;
+import org.springframework.util.StringUtils;
+import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.socket.server.HandshakeInterceptor;
+
+/**
+ * 应用模块名称</p>
+ * 代码描述</p>
+ * Copyright: Copyright (C) 2023 , Inc. All rights reserved. <p>
+ * Company: 成都诚唐科技有限责任公司</p>
+ *
+ * @author wany
+ * @since 2023/12/26
+ */
+public class WsHandshakeInterceptor implements HandshakeInterceptor{
+
+    @Override
+    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
+            Map<String, Object> attributes) throws Exception{
+        //将用户id放入socket处理器的会话(WebSocketSession)中
+        ServletServerHttpRequest serverHttpRequest = (ServletServerHttpRequest) request;
+        //获取参数
+        String userId = SaSecurityUtils.findSessionData().getUserId();
+        //String userId = serverHttpRequest.getServletRequest().getParameter("userId");
+        if(StringUtils.isEmpty(userId)){
+            return false;
+        }
+        attributes.put("userId", userId);
+        return true;
+    }
+
+    @Override
+    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
+            Exception exception){
+    }
+
+}

src/main/java/com/kingtom/shengtai/app/message/ws/WsSessionManager.java

@@ -0,0 +1,78 @@
+package com.kingtom.shengtai.app.message.ws;
+
+import java.io.IOException;
+import java.util.concurrent.ConcurrentHashMap;
+
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.web.socket.WebSocketSession;
+
+/**
+ * 应用模块名称</p>
+ * 代码描述</p>
+ * Copyright: Copyright (C) 2023 , Inc. All rights reserved. <p>
+ * Company: 成都诚唐科技有限责任公司</p>
+ *
+ * @author wany
+ * @since 2023/12/26
+ */
+@Slf4j
+
+public class WsSessionManager{
+
+    /**
+     * 保存连接 session 的地方
+     */
+    public 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);
+    }
+
+}
+