diff --git a/pigx-visual/pigx-oa-platform/src/main/java/com/pig4cloud/pigx/act/config/WebSocketConfig.java b/pigx-visual/pigx-oa-platform/src/main/java/com/pig4cloud/pigx/act/config/WebSocketConfig.java index 4cece5cc0..dbe1bd7e8 100644 --- a/pigx-visual/pigx-oa-platform/src/main/java/com/pig4cloud/pigx/act/config/WebSocketConfig.java +++ b/pigx-visual/pigx-oa-platform/src/main/java/com/pig4cloud/pigx/act/config/WebSocketConfig.java @@ -2,7 +2,9 @@ package com.pig4cloud.pigx.act.config; import cn.hutool.core.util.ObjectUtil; import cn.hutool.core.util.StrUtil; -import lombok.AllArgsConstructor; +import com.pig4cloud.pigx.common.core.constant.CommonConstants; +import com.pig4cloud.pigx.common.data.tenant.TenantContextHolder; +import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.context.annotation.Configuration; import org.springframework.messaging.Message; @@ -15,7 +17,7 @@ import org.springframework.messaging.support.ChannelInterceptor; import org.springframework.messaging.support.MessageHeaderAccessor; import org.springframework.security.core.context.SecurityContextHolder; import org.springframework.security.oauth2.provider.OAuth2Authentication; -import org.springframework.security.oauth2.provider.token.RemoteTokenServices; +import org.springframework.security.oauth2.provider.token.ResourceServerTokenServices; import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker; import org.springframework.web.socket.config.annotation.StompEndpointRegistry; import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer; @@ -27,15 +29,15 @@ import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerCo */ @Slf4j @Configuration -@AllArgsConstructor +@RequiredArgsConstructor @EnableWebSocketMessageBroker public class WebSocketConfig implements WebSocketMessageBrokerConfigurer { - private RemoteTokenServices tokenService; + private final ResourceServerTokenServices tokenService; @Override public void registerStompEndpoints(StompEndpointRegistry registry) { - registry.addEndpoint("/ws").setAllowedOrigins("*").withSockJS(); + registry.addEndpoint("/ws").setAllowedOriginPatterns("*").withSockJS(); } @Override @@ -49,18 +51,31 @@ public class WebSocketConfig implements WebSocketMessageBrokerConfigurer { @Override public Message preSend(Message message, MessageChannel channel) { StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class); + // 判断是否首次连接请求 if (StompCommand.CONNECT.equals(accessor.getCommand())) { + + // 处理租户 + String tenant = accessor.getFirstNativeHeader(CommonConstants.TENANT_ID); + if (StrUtil.isNotBlank(tenant)) { + TenantContextHolder.setTenantId(Integer.parseInt(tenant)); + + } + else { + TenantContextHolder.setTenantId(CommonConstants.TENANT_ID_1); + } + + // 验证令牌信息 String tokens = accessor.getFirstNativeHeader("Authorization"); log.info("webSocket token is {}", tokens); if (StrUtil.isBlank(tokens)) { return null; } - // 验证令牌信息 + OAuth2Authentication auth2Authentication = tokenService.loadAuthentication(tokens.split(" ")[1]); if (ObjectUtil.isNotNull(auth2Authentication)) { SecurityContextHolder.getContext().setAuthentication(auth2Authentication); - accessor.setUser(() -> auth2Authentication.getName()); + accessor.setUser(auth2Authentication::getName); return message; } else {