zianxin/pj_service_cloud
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

♻️ Refactoring code. #2358 websocket 调整

Browse Source
This commit is contained in:
lbw committed 2021-02-05 13:34:03 +08:00
1 parent ced4a056e5
commit 5b9d778425
1 file changed
+22 -7
Download Patch File Download Diff File Expand all files Collapse all files
pigx-visual/pigx-oa-platform/src/main/java/com/pig4cloud/pigx/act/config/WebSocketConfig.java
+22 -7
View File
@@ -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 {