Websocket消息发送工具类调整

soc-modules/soc-modules-mediator/src/main/java/com/xunmei/mediator/websocket/utils/WebSocketUtils.java

@@ -6,9 +6,9 @@ import com.xunmei.common.core.utils.JacksonUtils;
 import com.xunmei.common.redis.utils.RedisUtils;
 import com.xunmei.mediator.websocket.dto.WebSocketMessageDto;
 import com.xunmei.mediator.websocket.dto.WebsocketResult;
-import com.xunmei.mediator.websocket.dto.WebsocketUrlInfo;
 import com.xunmei.mediator.websocket.holder.WebSocketSessionHolder;
 import com.xunmei.system.api.util.LogUtils;
+import io.netty.util.internal.StringUtil;
 import lombok.AccessLevel;
 import lombok.NoArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -22,8 +22,6 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Consumer;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
 
 import static com.xunmei.mediator.websocket.constant.WebSocketConstants.WEB_SOCKET_TOPIC;
 
@@ -136,46 +134,59 @@ public class WebSocketUtils {
         try {
             for (WebSocketSession session : USER_SESSION_MAP.values()) {
                 session.sendMessage(new TextMessage(message));
-                String hostAddress = session.getRemoteAddress().getAddress().getHostAddress();
-                LogUtils.WEBSOCKET_MSG.info("消息广播成功,发送目标ip:{},消息内容：{}", hostAddress, message);
+                LogUtils.WEBSOCKET_MSG.info("消息广播成功,发送目标ip:{},消息内容：{}", getIp(session), message);
             }
         } catch (IOException e) {
-            log.error("消息广播失败:{}", e);
+            LogUtils.WEBSOCKET_MSG.error("消息广播失败,{}", e);
         }
     }
 
     public static boolean sendMessage(String token, String message) {
         final Map<String, WebSocketSession> USER_SESSION_MAP = WebSocketSessionHolder.map();
+        WebSocketSession webSocketSession = USER_SESSION_MAP.get(token);
         try {
-            WebSocketSession webSocketSession = USER_SESSION_MAP.get(token);
             if (null == webSocketSession) {
                 LogUtils.WEBSOCKET_MSG.error("消息广播失败,未找到对应在线的iot服务,iot服务:{},消息内容:{}", token, message);
                 return false;
             }
             webSocketSession.sendMessage(new TextMessage(message));
-            String hostAddress = webSocketSession.getRemoteAddress().getAddress().getHostAddress();
-            LogUtils.WEBSOCKET_MSG.info("消息广播成功,发送目标ip:{},消息内容：{}", hostAddress, message);
+
+            LogUtils.WEBSOCKET_MSG.info("消息广播成功,发送目标ip:{},消息内容：{}", getIp(webSocketSession), message);
             return true;
         } catch (IOException e) {
-            log.error("消息广播失败:{}", e);
+            LogUtils.WEBSOCKET_MSG.error("消息发送失败,ip:{},异常内容:{}", getIp(webSocketSession), e);
             return false;
         }
     }
 
     public static boolean sendMessage(WebSocketSession webSocketSession, WebsocketResult message) {
-        final Map<String, WebSocketSession> USER_SESSION_MAP = WebSocketSessionHolder.map();
         try {
             if (webSocketSession == null || !webSocketSession.isOpen()) {
                 return false;
             }
             final String msg = JacksonUtils.toJSONString(message);
             webSocketSession.sendMessage(new TextMessage(msg));
-            String hostAddress = webSocketSession.getRemoteAddress().getAddress().getHostAddress();
-            LogUtils.WEBSOCKET_MSG.info("消息广播成功,发送目标ip:{},消息内容：{}", hostAddress, msg);
+            LogUtils.WEBSOCKET_MSG.info("消息发送成功,发送目标ip:{},消息内容：{}", getIp(webSocketSession), msg);
             return true;
         } catch (IOException e) {
-            log.error("消息广播失败:{}", e);
+            LogUtils.WEBSOCKET_MSG.error("消息发送失败,ip:{},异常内容:{}", getIp(webSocketSession), e);
             return false;
         }
     }
+
+    public static String getIp(WebSocketSession webSocketSession) {
+        String ip = StringUtil.EMPTY_STRING;
+        if (webSocketSession == null || !webSocketSession.isOpen()) {
+            return ip;
+        }
+        final List<String> list = webSocketSession.getHandshakeHeaders().get("x-forwarded-for");
+        if (ObjectUtil.isNotEmpty(list)) {
+            ip = list.get(0);
+            if (ip.contains(",")) {
+                ip = ip.split(",")[0];
+            }
+            return ip;
+        }
+        return ip;
+    }
 }