SecurityContextHolder 自定义 SecurityContext 存储策略

306 阅读2分钟

采用TransmittableThreadLocal作为线程本地变量,TransmittableThreadLocal是阿里巴巴开源的一个线程本地变量,它是ThreadLocal的一个增强版,可以在线程池等多线程环境下使用,解决了ThreadLocal在多线程环境下的一些问题。

主要源码

// SecurityContextHolder
private static void initializeStrategy() {
  if (MODE_PRE_INITIALIZED.equals(strategyName)) {
    Assert.state(strategy != null, "When using " + MODE_PRE_INITIALIZED
        + ", setContextHolderStrategy must be called with the fully constructed strategy");
    return;
  }
  // 默认策略:采用ThreadLocal,有缺陷子线程中获取不到认证用户
  if (!StringUtils.hasText(strategyName)) {
    // Set default
    strategyName = MODE_THREADLOCAL;
  }
  if (strategyName.equals(MODE_THREADLOCAL)) {
    strategy = new ThreadLocalSecurityContextHolderStrategy();
    return;

  // 采用InheritableThreadLocal 子线程也可以获取到认证用户信息,但是存在线程池等使用实效的场景
  // 在创建Thread时,才会将父线程中的inheritableThreadLocals复制给新创建Thread的inheritableThreadLocals。
  // 但是在线程池中,业务线程(父线程)只是将任务对象(实现了Runnable或者Callable的对象)加入到任务队列中,
  // 并不是去创建线程池中的线程,因此线程池中线程也就获取不到业务线程中的上下文信息。
  if (strategyName.equals(MODE_INHERITABLETHREADLOCAL)) {
    strategy = new InheritableThreadLocalSecurityContextHolderStrategy();
    return;
  }
  // 存储到一个静态变量中
  if (strategyName.equals(MODE_GLOBAL)) {
    strategy = new GlobalSecurityContextHolderStrategy();
    return;
  }
  // Try to load a custom strategy 通过策略名称反射得到对应策略对象
  try {
    Class<?> clazz = Class.forName(strategyName);
    Constructor<?> customStrategy = clazz.getConstructor();
    strategy = (SecurityContextHolderStrategy) customStrategy.newInstance();
  }
  catch (Exception ex) {
    ReflectionUtils.handleReflectionException(ex);
  }
}

自定义存储策略

参考 ThreadLocalSecurityContextHolderStrategy 实现类实现,采用 TransmittableThreadLocal 作为线程本地变量,使其在多线程环境中能正常传递数据

/**
 * 基于 TransmittableThreadLocal 实现的 Security Context 持有者策略
 * 目的是,避免 @Async 等异步执行时,原生 ThreadLocal 的丢失问题
 *
 * @author LGC
 */
public class TransmittableThreadLocalSecurityContextHolderStrategy implements SecurityContextHolderStrategy {

    /**
     * 使用 TransmittableThreadLocal 作为上下文
     */
    private static final ThreadLocal<SecurityContext> CONTEXT_HOLDER = new TransmittableThreadLocal<>();

    @Override
    public void clearContext() {
        CONTEXT_HOLDER.remove();
    }

    @Override
    public SecurityContext getContext() {
        SecurityContext ctx = CONTEXT_HOLDER.get();
        if (ctx == null) {
            ctx = createEmptyContext();
            CONTEXT_HOLDER.set(ctx);
        }
        return ctx;
    }

    @Override
    public void setContext(SecurityContext context) {
        Assert.notNull(context, "Only non-null SecurityContext instances are permitted");
        CONTEXT_HOLDER.set(context);
    }

    @Override
    public SecurityContext createEmptyContext() {
        return new SecurityContextImpl();
    }

}

配置自定义策略

/**
* 更改SecurityContextHolder 的策略, 采用@Bean注入MethodInvokingBean实现,当然也可以其它方式
* <p>
* MethodInvokingBean实现了InitializingBean接口
* 所以在bean初始化过程中会执行afterPropertiesSet方法,在方法中会执行TargetClass对应setTargetMethod方法
* 可查看MethodInvokingBean相关源码
*/
@Bean
public MethodInvokingBean methodInvokingBean() {
 MethodInvokingBean methodInvokingBean = new MethodInvokingBean();
 methodInvokingBean.setTargetClass(SecurityContextHolder.class);
 methodInvokingBean.setTargetMethod("setStrategyName");
 methodInvokingBean.setArguments(TransmittableThreadLocalSecurityContextHolderStrategy.class.getName());
    return methodInvokingBean;
}

测试是否生效

需要用户已登录,正确传递 token,调用获取登录用户信息接口,看是否打印自定义策略名

// 获取登录用户信息
@GetMapping("/api/user/info")
public Result<LoginUser> userIfo() {
 log.info("SecurityContextHolder 策略名:{}",SecurityContextHolder.getContextHolderStrategy().getClass().getName());
 SecurityContext context = SecurityContextHolder.getContext();
 Authentication authentication = context.getAuthentication();
 LoginUser loginUser = (LoginUser) authentication.getPrincipal();
 return Result.success(loginUser);
}
// 打印输出自定义策略名:
// SecurityContextHolder 策略名:cn.xicode.cloud.lab2.security.config.TransmittableThreadLocalSecurityContextHolderStrategy