From 5b9d77842577ac662221fd990463e19c465a2a4f Mon Sep 17 00:00:00 2001
From: lbw <wangiegie@163.com>
Date: Fri, 5 Feb 2021 13:34:03 +0800
Subject: [PATCH] =?UTF-8?q?:recycle:=20Refactoring=20code.=20#2358=20=20we?=
 =?UTF-8?q?bsocket=20=E8=B0=83=E6=95=B4?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../pigx/act/config/WebSocketConfig.java      | 29 ++++++++++++++-----
 1 file changed, 22 insertions(+), 7 deletions(-)

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 {