初始化

xiangpei

2025-04-18 ccadf9480d4e6a9dcc227a2a0b1f9ae0612e36fd

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
package com.monkeylessey.websocket.handler;
 
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.monkeylessey.framework.utils.TokenUtil;
import com.monkeylessey.sys.domain.vo.SysUserVO;
import com.monkeylessey.websocket.Message;
import com.monkeylessey.websocket.NettyConnect;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.codec.http.websocketx.*;
import io.netty.util.AttributeKey;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
 
import java.util.Objects;
 
@Slf4j
@ChannelHandler.Sharable
@Component
@RequiredArgsConstructor
public class WebSocketHandler extends SimpleChannelInboundHandler<WebSocketFrame> {
 
    private final TokenUtil tokenUtil;
 
 
    @Override
    protected void channelRead0(ChannelHandlerContext ctx, WebSocketFrame msg) {
        this.handleWebSocketFrame(ctx, msg);
    }
 
    @Override
    public void channelActive(ChannelHandlerContext ctx) {
        if (!NettyConnect.getChannelGroup().contains(ctx.channel())) {
            NettyConnect.getChannelGroup().add(ctx.channel());
        }
    }
 
    @Override
    public void channelInactive(ChannelHandlerContext ctx) {
        log.info("有客户端断开连接了");
        AttributeKey<Integer> userIdKey = AttributeKey.valueOf("userId");
        Integer userId = ctx.channel().attr(userIdKey).get();
        if (Objects.nonNull(userId)) {
            NettyConnect.getUserChannelMap().remove(userId);
        }
        NettyConnect.getChannelGroup().remove(ctx.channel());
    }
 
    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        ctx.channel().close();
    }
 
    // 处理ws数据
    private void handleWebSocketFrame(ChannelHandlerContext ctx, WebSocketFrame frame) {
        // 处理关闭连接
        if (frame instanceof CloseWebSocketFrame) {
            NettyConnect.getChannelGroup().remove(ctx.channel());
            ctx.close();
            return;
        }
        if (frame instanceof TextWebSocketFrame) {
            if ("ping".equals(((TextWebSocketFrame) frame).text())) {
                ctx.channel().writeAndFlush(new PongWebSocketFrame(frame.content().retain()));
                return;
            }
            // 处理token，只有客户端连接的时候才会发送socket消息
            String token = ((TextWebSocketFrame) frame).text();
            try {
                Message message = new ObjectMapper().readValue(token, Message.class);
                // 验证token并将用户ID存入到连接中
                this.handleToken(message.getToken(), ctx);
            } catch (JsonProcessingException e) {
                log.error("消息格式错误");
            }
        } else if (frame instanceof BinaryWebSocketFrame) {
            // 处理二进制消息
            // ...
        } else if (frame instanceof PingWebSocketFrame) {
            // 客户端也是netty实现的就可以用这个来ping、peng
            // 处理 Ping 消息
            // 收到 Ping 消息，回应一个 Pong 消息（表明我还活着）
            ctx.channel().writeAndFlush(new PongWebSocketFrame(frame.content().retain()));
        } else if (frame instanceof PongWebSocketFrame) {
            // 处理 Pong 消息
            // pong消息如果没有特定需求，不用处理
        } else if (frame instanceof ContinuationWebSocketFrame) {
            // 处理连续帧消息（比较大的数据，分片）
            // ...
        }
    }
 
    // 处理token
    private String handleToken(String token, ChannelHandlerContext ctx) {
        if (!StringUtils.hasText(token)) {
            NettyConnect.getChannelGroup().remove(ctx.channel());
            throw new RuntimeException("非法的访问凭证");
        }
        // 获取 userId 参数
        SysUserVO currentUserInfo = null;
        try {
            currentUserInfo = tokenUtil.getCurrentUserInfo(token);
        } catch (JsonProcessingException e) {
            NettyConnect.getChannelGroup().remove(ctx.channel());
            throw new RuntimeException("非法的访问凭证");
        }
        if (Objects.isNull(currentUserInfo)) {
            NettyConnect.getChannelGroup().remove(ctx.channel());
            throw new RuntimeException("用户不存在");
        }
 
        String userId = currentUserInfo.getId();
        AttributeKey<String> userIdKey = AttributeKey.valueOf("userId");
        ctx.channel().attr(userIdKey).set(userId);
        NettyConnect.getUserChannelMap().put(userId, ctx.channel());
 
        return userId;
    }
 
 
}