IM/src/main/java/com/zhentao/information/handler/WebSocketHandler.java

package com.zhentao.information.handler;

import com.alibaba.fastjson.JSON;
import com.zhentao.information.entity.Message;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket消息处理器
 */
@Slf4j
@Component
public class WebSocketHandler extends SimpleChannelInboundHandler<TextWebSocketFrame> {

    // 用户ID和Channel的映射关系
    private static final Map<String, ChannelHandlerContext> USER_CHANNEL_MAP = new ConcurrentHashMap<>();

    /**
     * 获取用户Channel映射
     * @return 用户Channel映射
     */
    public Map<String, ChannelHandlerContext> getUserChannelMap() {
        return USER_CHANNEL_MAP;
    }

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, TextWebSocketFrame msg) {
        String message = msg.text();
        Message messageObj = JSON.parseObject(message, Message.class);

        // 存储用户连接
        USER_CHANNEL_MAP.put(messageObj.getFromUserId(), ctx);

        // 获取接收者的Channel
        ChannelHandlerContext toUserCtx = USER_CHANNEL_MAP.get(messageObj.getToUserId());
        if (toUserCtx != null) {
            // 发送消息给接收者
            toUserCtx.channel().writeAndFlush(new TextWebSocketFrame(JSON.toJSONString(messageObj)));
            log.info("消息已发送给用户: {}, 内容: {}", messageObj.getToUserId(), messageObj.getContent());
        } else {
            log.info("用户 {} 不在线", messageObj.getToUserId());
        }
    }

    @Override
    public void handlerRemoved(ChannelHandlerContext ctx) {
        // 用户断开连接时，移除映射关系
        USER_CHANNEL_MAP.entrySet().removeIf(entry -> entry.getValue().equals(ctx));
    }
}