增加 spring websocket 用例

spring/spring-websocket/src/main/java/com/heibaiying/websocket/CustomHandshakeInterceptor.java

@@ -5,24 +5,27 @@ import org.springframework.http.server.ServerHttpResponse;
 import org.springframework.web.socket.WebSocketHandler;
 import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;
 
+import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.util.Map;
 
 /**
  * @author : heibaiying
- * @description : 握手拦截器
+ * @description : 可以按照需求实现权限拦截等功能
  */
 public class CustomHandshakeInterceptor extends HttpSessionHandshakeInterceptor {
 
-    /*@Override
-    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
-        InetSocketAddress remoteAddress = request.getRemoteAddress();
-        System.out.println(remoteAddress);
-        return true;
-    }*/
-
     @Override
     public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
+        InetSocketAddress remoteAddress = request.getRemoteAddress();
+        InetAddress address = remoteAddress.getAddress();
+        System.out.println(address);
+        /*
+         * 最后需要要显示调用父类方法，父类的beforeHandshake方法
+         * 把ServerHttpRequest 中session中对应的值拷贝到WebSocketSession中。
+         * 如果我们没有实现这个方法，我们在最后的handler处理中 是拿不到 session中的值
+         * 作为测试 可以注释掉下面这一行 可以发现自定义处理器中session的username总是为空
+         */
         return super.beforeHandshake(request, response, wsHandler, attributes);
     }
 }

spring/spring-websocket/src/main/java/com/heibaiying/websocket/CustomerHandler.java

@@ -1,21 +1,61 @@
 package com.heibaiying.websocket;
 
+import com.heibaiying.constant.Constant;
+import org.springframework.web.socket.CloseStatus;
 import org.springframework.web.socket.TextMessage;
+import org.springframework.web.socket.WebSocketMessage;
 import org.springframework.web.socket.WebSocketSession;
 import org.springframework.web.socket.handler.TextWebSocketHandler;
 
-import java.time.LocalDateTime;
+import java.io.IOException;
+import java.net.URLDecoder;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * @author : heibaiying
- * @description : 自定义消息处理器
+ * @description : 自定义消息处理类
  */
 public class CustomerHandler extends TextWebSocketHandler {
 
+    private Map<String, WebSocketSession> nameAndSession = new ConcurrentHashMap<>();
+
+    // 建立连接时候触发
+    @Override
+    public void afterConnectionEstablished(WebSocketSession session)  {
+        String username = getNameFromSession(session);
+        nameAndSession.putIfAbsent(username, session);
+    }
+
+
+    // 关闭连接时候触发
+    @Override
+    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
+        String username = getNameFromSession(session);
+        nameAndSession.remove(username);
+    }
+
+    // 处理消息
     @Override
     protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
-        String stringMessage = new String(message.asBytes());
-        System.out.println("服务端收到消息：" + stringMessage);
-        session.sendMessage(new TextMessage(stringMessage+LocalDateTime.now()));
+        // 防止中文乱码
+        String msg = URLDecoder.decode(message.getPayload(), "utf-8");
+        String username = getNameFromSession(session);
+        // 简单模拟群发消息
+        TextMessage reply = new TextMessage(username + " : " + msg);
+        nameAndSession.forEach((s, webSocketSession)
+                -> {
+            try {
+                webSocketSession.sendMessage(reply);
+            } catch (IOException e) {
+                e.printStackTrace();
+            }
+        });
+    }
+
+
+    private String getNameFromSession(WebSocketSession session) {
+        Map<String, Object> attributes = session.getAttributes();
+        return (String) attributes.get(Constant.USER_NAME);
     }
 }