Spring多租户应用中WebSocket消息的租户关联方案咨询

思路一：通过握手拦截器+消息通道拦截器复用Thread Local租户上下文 #

这个思路和你现有Web Filter的逻辑完全对齐，核心是把握手阶段拿到的租户信息，在每个WebSocket消息处理前注入Thread Local，处理完成后清理，确保业务代码无需修改就能复用现有租户上下文逻辑。

1. 自定义HandshakeInterceptor存储租户信息
   在握手阶段解析主机名得到租户ID，存入WebSocketSession的属性中：

public class TenantHandshakeInterceptor implements HandshakeInterceptor {
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        // 复用你现有Filter里的主机名解析逻辑
        String host = request.getHeaders().getHost().getHostName();
        String tenantId = TenantResolver.resolveFromHost(host); // 替换成你的解析方法
        attributes.put("tenantId", tenantId);
        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {}
}

2. 自定义ChannelInterceptor注入Thread Local
   实现ChannelInterceptor，在消息处理前从WebSocketSession中取出租户ID，放到Thread Local的租户上下文，处理完成后清理避免内存泄漏：

public class TenantChannelInterceptor implements ChannelInterceptor {
    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        WebSocketSession session = (WebSocketSession) message.getHeaders().get(SimpMessageHeaderAccessor.SESSION_HEADER);
        if (session != null) {
            String tenantId = (String) session.getAttributes().get("tenantId");
            TenantContext.setCurrentTenant(tenantId); // 复用你现有Thread Local上下文类
        }
        return message;
    }

    @Override
    public void afterSendCompletion(Message<?> message, MessageChannel channel, boolean sent, Exception ex) {
        TenantContext.clear();
    }
}

3. 配置WebSocket注册拦截器
   在WebSocket配置类中注册这两个拦截器：

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {
    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        registration.interceptors(new TenantChannelInterceptor());
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/ws")
                .addInterceptors(new TenantHandshakeInterceptor())
                .withSockJS();
    }

    // 其他消息代理配置...
}

完成后，你的@MessageMapping方法就能像处理Web请求一样，直接通过TenantContext.getCurrentTenant()获取租户ID。

思路二：自定义参数解析器，直接在@MessageMapping方法中注入租户ID #

如果不想依赖Thread Local，可以自定义参数解析器，让租户ID直接作为方法参数注入，业务代码更直观。

1. 自定义租户参数注解

@Target(ElementType.PARAMETER)
@Retention(RetentionPolicy.RUNTIME)
public @interface TenantId {}

2. 实现HandlerMethodArgumentResolver
   从WebSocketSession中取出租户ID，绑定到标注了@TenantId的参数上：

public class TenantArgumentResolver implements HandlerMethodArgumentResolver {
    @Override
    public boolean supportsParameter(MethodParameter parameter) {
        return parameter.hasParameterAnnotation(TenantId.class);
    }

    @Override
    public Object resolveArgument(MethodParameter parameter, ModelAndViewContainer mavContainer, NativeWebRequest webRequest, WebDataBinderFactory binderFactory) throws Exception {
        WebSocketSession session = webRequest.getAttribute(WebSocketSession.class.getName(), NativeWebRequest.SCOPE_REQUEST);
        return session.getAttributes().get("tenantId");
    }
}

3. 注册参数解析器
   在WebSocket配置类中添加解析器：

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {
    // ... 其他配置

    @Override
    public void addArgumentResolvers(List<HandlerMethodArgumentResolver> argumentResolvers) {
        argumentResolvers.add(new TenantArgumentResolver());
    }
}

之后在@MessageMapping方法中就能直接使用：

@MessageMapping("/chat/send")
public void handleChatMessage(@Payload ChatMessage message, @TenantId String tenantId) {
    // 直接使用tenantId处理业务逻辑
}

思路三：结合Spring Security上下文（若已集成） #

如果你的应用已经集成Spring Security，可以在握手阶段把租户ID存入Authentication对象的Details或Principal中，之后在消息处理时从SecurityContext直接获取：

// 在HandshakeInterceptor的beforeHandshake方法中
Authentication auth = SecurityContextHolder.getContext().getAuthentication();
if (auth != null) {
    // 假设你用了自定义UserDetails
    ((CustomUserDetails) auth.getPrincipal()).setTenantId(tenantId);
}

然后在@MessageMapping方法中通过SecurityContextHolder.getContext().getAuthentication()获取租户信息，适合已经依赖Security的场景。

内容的提问来源于stack exchange，提问作者tbeernot