【MS】spring security oauth2之RedisTokenStore添加限流参数保存

spring security oauth2之RedisTokenStore添加限流参数保存

前言

在之前写的《网关 springcloud-gateway 基于Token限流》中，从网关中获取请求的token ，根据token 去redis 中获取限流参数，那么这个参数是在哪里保存的呢？

没错，在用户认证的时候，就把限流参数和token一起保存到redis了

直接上代码

1、创建CustomRedisTokenStore 类也就是 把原来RedisTokenStore的代码复制过来

添加两个方法

• setUserLimitLevel(OAuth2AccessToken token, OAuth2Authentication authentication) 保存限流参数

• updateUserLimitLevel(StoreUser storeUser,String token) 更新限流参数

• 在removeAccessToken() 方法中也把自己添加的那个也一起remove掉就行了

public class CustomRedisTokenStore implements TokenStore {

    ....代码省略
    
    
	public void setUserLimitLevel(OAuth2AccessToken token, OAuth2Authentication authentication){
		StoreUser storeUser = new StoreUser();
		String grantType =  authentication.getOAuth2Request().getGrantType();
		switch (grantType){
			case SecurityConstants.CLIENT_CREDENTIALS:
				storeUser.setLimitLevel(MapUtil.getInt(token.getAdditionalInformation(),SecurityConstants.LIMIT_LEVEL));
				storeUser.setUserId(authentication.getOAuth2Request().getClientId());
				storeUser.setUserName(authentication.getOAuth2Request().getClientId());
				break;
			default:
				CustomUserDetailsUser customUser = (CustomUserDetailsUser) authentication.getPrincipal();
				storeUser = StoreUser.builder().userId(customUser.getUserId()).limitLevel(customUser.getLimitLevel()).userName(customUser.getUsername()).build();
				break;
		}
		byte[] serializedAuthUser = serialize(storeUser);
		byte[] authUserKey = serializeKey(AUTH_USER + token.getValue());
		RedisConnection conn = getConnection();
		try {
			if (springDataRedis_2_0) {
				try {
					this.redisConnectionSet_2_0.invoke(conn, authUserKey, serializedAuthUser);
				} catch (Exception ex) {
					throw new RuntimeException(ex);
				}
			} else {
				conn.set(authUserKey, serializedAuthUser);
			}

			if (token.getExpiration() != null) {
				int seconds = token.getExpiresIn();
				conn.expire(authUserKey, seconds);
			}
			conn.closePipeline();
		}finally {
			conn.close();
		}

	}

	public void updateUserLimitLevel(StoreUser storeUser,String token){
		byte[] serializedAuthUser = serialize(storeUser);
		byte[] authUserKey = serializeKey(AUTH_USER + token);
		RedisConnection conn = getConnection();
		try{
			conn.openPipeline();
			conn.set(authUserKey, serializedAuthUser);
			conn.closePipeline();
		}finally {
			conn.close();
		}
	}
    
    ....代码省略
}

2、配置bean 让自己写的CustomRedisTokenStore 生效

    /**
     * 指定tokenStore 为 redis
     * @return
     */
    @Bean
    public TokenStore tokenStore() {
        CustomRedisTokenStore tokenStore = new CustomRedisTokenStore(redisConnectionFactory);
        tokenStore.setPrefix(SecurityConstants.MS_OAUTH_PREFIX);
        return tokenStore;
    }

3、在用户验证后，把从数据库查询到的限流参数 附加到 security 的User中

• 新建一个类CustomUserDetailsUser 继承 security 的 user类
• 在 UserDetailsService 中的 loadUserByUsername方法 把一些自定义的数据放到 CustomUserDetailsUser 对象中并返回

/**
 * @Description //TODO 认证后返回的用户信息$
 * @Date 20:49
 * @Author yzcheng90@qq.com
 **/
public class CustomUserDetailsUser extends User {

    @Getter
    @Setter
    public Integer userId;
    /**
     * 限流等级
     */
    @Getter
    @Setter
    private Integer limitLevel;

    public CustomUserDetailsUser(Integer userId,Integer limitLevel, String username, String password, boolean enabled, boolean accountNonExpired, boolean credentialsNonExpired, boolean accountNonLocked, Collection<? extends GrantedAuthority> authorities) {
        super(username, password, enabled, accountNonExpired, credentialsNonExpired, accountNonLocked, authorities);
        this.userId = userId;
        this.limitLevel = limitLevel;
    }
}

4、重写 TokenEnhancer 接口，把放在CustomUserDetailsUser中的一些自定义的参数 放到 OAuth2Authentication 中,这样在TokenStore中就可以获取到参数 然后保存到redis了

/**
 * @author czx
 * @title: CustomTokenEnhancer
 * @projectName ms
 * @description: TODO
 * @date 2019/7/2610:12
 */
public class CustomTokenEnhancer implements TokenEnhancer {

    private CustomClientDetailsService customClientDetailsService;

    public CustomTokenEnhancer(CustomClientDetailsService customClientDetailsService){
        this.customClientDetailsService = customClientDetailsService;
    }

    @Override
    public OAuth2AccessToken enhance(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
        if (SecurityConstants.CLIENT_CREDENTIALS.equals(authentication.getOAuth2Request().getGrantType())) {
            ClientDetails clientDetails = customClientDetailsService.loadClientByClientId((String) authentication.getPrincipal());
            ((DefaultOAuth2AccessToken) accessToken).setAdditionalInformation(clientDetails.getAdditionalInformation());
            return accessToken;
        }

        final Map<String, Object> additionalInfo = new HashMap<>(8);
        CustomUserDetailsUser customUserDetailsUser = (CustomUserDetailsUser) authentication.getUserAuthentication().getPrincipal();
        additionalInfo.put(SecurityConstants.USER_ID, customUserDetailsUser.getUserId());
        additionalInfo.put(SecurityConstants.LIMIT_LEVEL, customUserDetailsUser.getLimitLevel());
        additionalInfo.put(SecurityConstants.USER_NAME, customUserDetailsUser.getUsername());
        ((DefaultOAuth2AccessToken) accessToken).setAdditionalInformation(additionalInfo);
        return accessToken;
    }
}

5、配置bean,让重写的CustomTokenEnhancer生效

    /**
     * token 属性附加。
     * @return TokenEnhancer
     */
    @Bean
    public TokenEnhancer tokenEnhancer() {
        return new CustomTokenEnhancer(customClientDetailsService());
    }

到此就完成了！

源码传送门：https://github.com/yzcheng90/ms