diff --git a/src/main/java/com/youlai/system/SystemApplication.java b/src/main/java/com/youlai/system/SystemApplication.java index 94bc3f42..7d844380 100644 --- a/src/main/java/com/youlai/system/SystemApplication.java +++ b/src/main/java/com/youlai/system/SystemApplication.java @@ -3,9 +3,11 @@ package com.youlai.system; import org.springframework.boot.SpringApplication; import org.springframework.boot.autoconfigure.SpringBootApplication; import org.springframework.boot.context.properties.ConfigurationPropertiesScan; +import org.springframework.scheduling.annotation.EnableScheduling; @SpringBootApplication @ConfigurationPropertiesScan +@EnableScheduling public class SystemApplication { public static void main(String[] args) { SpringApplication.run(SystemApplication.class, args); diff --git a/src/main/java/com/youlai/system/config/WebSocketConfig.java b/src/main/java/com/youlai/system/config/WebSocketConfig.java index 427d20fc..a614c944 100644 --- a/src/main/java/com/youlai/system/config/WebSocketConfig.java +++ b/src/main/java/com/youlai/system/config/WebSocketConfig.java @@ -4,8 +4,12 @@ import cn.hutool.core.util.StrUtil; import cn.hutool.jwt.JWTPayload; import cn.hutool.jwt.JWTUtil; import com.youlai.system.common.constant.SecurityConstants; +import com.youlai.system.service.WebsocketService; +import groovy.lang.Lazy; +import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.jetbrains.annotations.NotNull; +import org.springframework.beans.factory.annotation.Autowired; import org.springframework.context.annotation.Configuration; import org.springframework.http.HttpHeaders; import org.springframework.messaging.Message; @@ -32,6 +36,8 @@ import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerCo @Slf4j public class WebSocketConfig implements WebSocketMessageBrokerConfigurer { + @Autowired + private WebsocketService websocketService; /** * 注册一个端点,客户端通过这个端点进行连接 */ @@ -74,27 +80,24 @@ public class WebSocketConfig implements WebSocketMessageBrokerConfigurer { @Override public Message preSend(@NotNull Message message, @NotNull MessageChannel channel) { StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class); - // 如果是连接请求(CONNECT 命令),从请求头中取出 token 并设置到认证信息中 - if (accessor != null && StompCommand.CONNECT.equals(accessor.getCommand())) { - // 从连接头中提取授权令牌 - String bearerToken = accessor.getFirstNativeHeader(HttpHeaders.AUTHORIZATION); - - // 验证令牌格式并提取用户信息 - if (StrUtil.isNotBlank(bearerToken) && bearerToken.startsWith("Bearer ")) { - try { - // 移除 "Bearer " 前缀,从令牌中提取用户信息(username), 并设置到认证信息中 + if (accessor != null) { + if (StompCommand.CONNECT.equals(accessor.getCommand())) { + String bearerToken = accessor.getFirstNativeHeader(HttpHeaders.AUTHORIZATION); + if (StrUtil.isNotBlank(bearerToken) && bearerToken.startsWith("Bearer ")) { bearerToken = bearerToken.substring(SecurityConstants.JWT_TOKEN_PREFIX.length()); String username = JWTUtil.parseToken(bearerToken).getPayloads().getStr(JWTPayload.SUBJECT); if (StrUtil.isNotBlank(username)) { accessor.setUser(() -> username); - return message; + websocketService.addUser(username); } - } catch (Exception e) { - log.error("Failed to process authentication token.", e); + } + } else if (StompCommand.DISCONNECT.equals(accessor.getCommand())) { + if (accessor.getUser() != null) { + String username = accessor.getUser().getName(); + websocketService.removeUser(username); } } } - // 不是连接请求,直接放行 return ChannelInterceptor.super.preSend(message, channel); } }); diff --git a/src/main/java/com/youlai/system/service/WebsocketService.java b/src/main/java/com/youlai/system/service/WebsocketService.java new file mode 100644 index 00000000..daf64158 --- /dev/null +++ b/src/main/java/com/youlai/system/service/WebsocketService.java @@ -0,0 +1,10 @@ +package com.youlai.system.service; + +public interface WebsocketService { + + void addUser(String username); + + void removeUser(String username) ; + + int getOnlineUserCount() ; +} diff --git a/src/main/java/com/youlai/system/service/impl/WebsocketServiceImpl.java b/src/main/java/com/youlai/system/service/impl/WebsocketServiceImpl.java new file mode 100644 index 00000000..613251c7 --- /dev/null +++ b/src/main/java/com/youlai/system/service/impl/WebsocketServiceImpl.java @@ -0,0 +1,44 @@ +package com.youlai.system.service.impl; + +import com.youlai.system.service.WebsocketService; +import groovy.lang.Lazy; +import lombok.RequiredArgsConstructor; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.messaging.simp.SimpMessagingTemplate; +import org.springframework.scheduling.annotation.Scheduled; +import org.springframework.stereotype.Service; + +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; + +@Service +public class WebsocketServiceImpl implements WebsocketService { + + @Lazy + @Autowired + private SimpMessagingTemplate messagingTemplate; + + private static final Set onlineUsers = ConcurrentHashMap.newKeySet(); + + @Override + public void addUser(String username) { + onlineUsers.add(username); + } + + @Override + public void removeUser(String username) { + onlineUsers.remove(username); + } + + @Override + public int getOnlineUserCount() { + return onlineUsers.size(); + } + + @Scheduled(fixedRate = 5000) + public void sendOnlineUserCount() { + int onlineUserCount = this.getOnlineUserCount(); + messagingTemplate.convertAndSend("/topic/onlineUserCount", onlineUserCount); + } + +} diff --git a/src/main/resources/application.yml b/src/main/resources/application.yml index 913e9673..335ea46b 100644 --- a/src/main/resources/application.yml +++ b/src/main/resources/application.yml @@ -1,4 +1,6 @@ spring: + main: + allow-circular-references: true application: name: youlai-boot profiles: