Spring Security OAuth 2.0授權(quán)服務(wù)器結(jié)合Redis實(shí)現(xiàn)獲取accessToken速率限制

概述

在生產(chǎn)環(huán)境中脐区，我們通常頒發(fā)給OAuth2客戶端有效期較長(zhǎng)的token蚪腐，但是授權(quán)服務(wù)無(wú)從知曉O(shè)Auth2客戶端服務(wù)是否頻繁獲取token，便于我們主動(dòng)控制token的頒發(fā)泽台，減少數(shù)據(jù)庫(kù)操作，本文我們將結(jié)合Redis實(shí)現(xiàn)滑動(dòng)窗口算法限制速率解決此問(wèn)題瓜贾。

先決條件

• java 8+
• Redis
• Lua

授權(quán)服務(wù)器

本節(jié)中我們將使用Spring Authorization Server 搭建一個(gè)簡(jiǎn)單的授權(quán)服務(wù)器悍抑，并通過(guò)擴(kuò)展OAuth2TokenCustomizer實(shí)現(xiàn)access_token的速率限制。

Maven依賴

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-security</artifactId>
            <version>2.6.7</version>
        </dependency>

        <dependency>
            <groupId>org.springframework.security</groupId>
            <artifactId>spring-security-oauth2-authorization-server</artifactId>
            <version>0.3.1</version>
        </dependency>

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
            <version>2.6.7</version>
        </dependency>

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-data-redis</artifactId>
            <version>2.6.7</version>
        </dependency>

配置

首先添加spring.redis配置連接本地Redis服務(wù)：

server:
  port: 8080

spring:
  redis:
    host: localhost
    database: 0
    port: 6379
    password: 123456
    timeout: 1800
    lettuce:
      pool:
        max-active: 20
        max-wait: 60
        max-idle: 5
        min-idle: 0
      shutdown-timeout: 100

接下來(lái)我們需要注冊(cè)一個(gè)OAuth2客戶端开呐，聲明客戶端如下：

    @Bean
    public RegisteredClientRepository registeredClientRepository() {
        RegisteredClient registeredClient = RegisteredClient.withId(UUID.randomUUID().toString())
                .clientId("relive-client")
                .clientSecret("{noop}relive-client")
                .clientAuthenticationMethods(s -> {
                    s.add(ClientAuthenticationMethod.CLIENT_SECRET_POST);
                    s.add(ClientAuthenticationMethod.CLIENT_SECRET_BASIC);
                })
                .authorizationGrantType(AuthorizationGrantType.CLIENT_CREDENTIALS)
                .redirectUri("http://127.0.0.1:8070/login/oauth2/code/messaging-client-model")
                .scope("message.read")
                .clientSettings(ClientSettings.builder()
                        .requireAuthorizationConsent(false)
                        .requireProofKey(false)
                        .build())
                .tokenSettings(TokenSettings.builder()
                        .accessTokenFormat(OAuth2TokenFormat.SELF_CONTAINED)
                        .idTokenSignatureAlgorithm(SignatureAlgorithm.RS256)
                        .accessTokenTimeToLive(Duration.ofSeconds(30 * 60))
                        .refreshTokenTimeToLive(Duration.ofSeconds(60 * 60))
                        .reuseRefreshTokens(true)
                        .setting("accessTokenLimitTimeSeconds", 5 * 60)
                        .setting("accessTokenLimitRate", 3)
                        .build())
                .build();

        return new InMemoryRegisteredClientRepository(registeredClient);
    }

上述OAuth2客戶端信息如下：

• clientId: relive-client
• clientSecret: relive-client
• clientAuthenticationMethod: client_secret_post,client_secret_basic
• authorizationGrantType: client_credentials
• redirectUri: http://127.0.0.1:8070/login/oauth2/code/messaging-client-model
• scope: message.read

特別注意：我們額外添加了兩個(gè)參數(shù)用于控制AccessToken的速率限制，accessTokenLimitTimeSeconds訪問(wèn)限制時(shí)間规求，accessTokenLimitRate訪問(wèn)限制次數(shù)筐付。
此外，我們?yōu)閱蝹€(gè)客戶端添加限制參數(shù)阻肿，由此可以針對(duì)不同OAuth2客戶端設(shè)置不同的速率限制或者取消瓦戚。

使用Spring Authorization Server提供的授權(quán)服務(wù)默認(rèn)配置，并將未認(rèn)證的授權(quán)請(qǐng)求重定向到登錄頁(yè)面：

    @Bean
    @Order(Ordered.HIGHEST_PRECEDENCE)
    public SecurityFilterChain authorizationServerSecurityFilterChain(HttpSecurity http) throws Exception {
        OAuth2AuthorizationServerConfiguration.applyDefaultSecurity(http);
        return http.exceptionHandling(exceptions -> exceptions.
                authenticationEntryPoint(new LoginUrlAuthenticationEntryPoint("/login"))).build();
    }

其余常規(guī)配置本文將不再贅述丛塌，您可以參考以往文章或從文末鏈接中獲取源碼较解。

接下來(lái)我們將利用Redis sorted set數(shù)據(jù)結(jié)構(gòu)實(shí)現(xiàn)滑動(dòng)窗口算法用于access_token速率限制畜疾，我們將利用Lua腳本保證Redis操作的原子性，節(jié)省網(wǎng)絡(luò)開銷印衔。

redis.replicate_commands()

local key = KEYS[1]

local windowSize = tonumber(ARGV[1])
local rate = tonumber(ARGV[2])
local now = tonumber(redis.call("TIME")[1])

redis.call("zadd", key, now, now)
local start = math.max(0, now - windowSize)

local requestRate = tonumber(redis.call("zcount", key, start, now))

local result = true
if requestRate > rate then
  result = false
end

redis.call("zremrangebyscore", key, "-inf", "("..start)

return result

上述Lua腳本遵循以下步驟：

• 將當(dāng)前時(shí)間（秒）作為value和score 添加進(jìn)有序集合(sorted set)中
• 計(jì)算窗口長(zhǎng)度啡捶，統(tǒng)計(jì)窗口中成員總數(shù)，該總數(shù)表示該窗口長(zhǎng)度中已請(qǐng)求次數(shù)
• 判斷請(qǐng)求次數(shù)是否超過(guò)閾值
• 移除已失效成員

RedisAccessTokenLimiter從TokenSettings獲取參數(shù)accessTokenLimitTimeSeconds奸焙，accessTokenLimitRate瞎暑，由RedisTemplate執(zhí)行Lua腳本，并傳遞參數(shù)信息与帆。

@Slf4j
public class RedisAccessTokenLimiter implements AccessTokenLimiter {
    private static final String ACCESS_TOKEN_LIMIT_TIME_SECONDS = "accessTokenLimitTimeSeconds";
    private static final String ACCESS_TOKEN_LIMIT_RATE = "accessTokenLimitRate";
    private final RedisTemplate<String, Object> redisTemplate;
    private final RedisScript<Boolean> script;

    public RedisAccessTokenLimiter(RedisTemplate<String, Object> redisTemplate, RedisScript<Boolean> script) {
        Assert.notNull(redisTemplate, "redisTemplate can not be null");
        Assert.notNull(script, "script can not be null");
        this.redisTemplate = redisTemplate;
        this.script = script;
    }


    @Override
    public boolean isAllowed(RegisteredClient registeredClient) {

        TokenSettings tokenSettings = registeredClient.getTokenSettings();
        if (tokenSettings == null || tokenSettings.getSetting(ACCESS_TOKEN_LIMIT_TIME_SECONDS) == null ||
                tokenSettings.getSetting(ACCESS_TOKEN_LIMIT_RATE) == null) {
            return true;
        }
        int accessTokenLimitTimeSeconds = tokenSettings.getSetting(ACCESS_TOKEN_LIMIT_TIME_SECONDS);

        int accessTokenLimitRate = tokenSettings.getSetting(ACCESS_TOKEN_LIMIT_RATE);

        String clientId = registeredClient.getClientId();

        try {
            List<String> keys = getKeys(clientId);

            return redisTemplate.execute(this.script, keys, accessTokenLimitTimeSeconds, accessTokenLimitRate);
        } catch (Exception e) {
            /*
             * 我們不希望硬依賴 Redis 來(lái)允許訪問(wèn)了赌。 確保設(shè)置
             * 一個(gè)警報(bào)，知道發(fā)生了許多次玄糟。
             */
            log.error("Error determining if user allowed from redis", e);
        }
        return true;
    }

    static List<String> getKeys(String id) {
        // 在key周圍使用 `{}` 以使用 Redis Key hash tag
        // 這允許使用 redis 集群
        String prefix = "access_token_rate_limiter.{" + id;

        String key = prefix + "}.client";
        return Arrays.asList(key);
    }

}

已知OAuth2TokenCustomizer提供了自定義OAuth2Token的屬性的能力勿她，但是在本示例中我們將使用OAuth2TokenCustomizer作為擴(kuò)展點(diǎn)，使用AccessTokenLimiter提供了速率限制阵翎，當(dāng)請(qǐng)求超過(guò)閾值時(shí)逢并，將拋出OAuth2AuthenticationException異常。

public class AccessTokenRestrictionCustomizer implements OAuth2TokenCustomizer<JwtEncodingContext> {
    private static final String DEFAULT_ERROR_URI = "https://datatracker.ietf.org/doc/html/rfc6749#section-4.1.2.1";
    private final AccessTokenLimiter tokenLimiter;

    public AccessTokenRestrictionCustomizer(AccessTokenLimiter tokenLimiter) {
        Assert.notNull(tokenLimiter, "accessTokenLimiter can not be null");
        this.tokenLimiter = tokenLimiter;
    }

    /**
     * 通過(guò){@link AccessTokenLimiter} 為OAuth2 客戶端模式訪問(wèn)令牌添加訪問(wèn)限制
     *
     * @param context
     */
    @Override
    public void customize(JwtEncodingContext context) {
        if (AuthorizationGrantType.CLIENT_CREDENTIALS.equals(context.getAuthorizationGrantType())) {
            RegisteredClient registeredClient = context.getRegisteredClient();
            if (registeredClient == null) {
                OAuth2Error error = new OAuth2Error(OAuth2ErrorCodes.INVALID_REQUEST, "OAuth 2.0 Parameter: " + OAuth2ParameterNames.CLIENT_ID, DEFAULT_ERROR_URI);
                throw new OAuth2AuthenticationException(error);
            }


            boolean requiresGenerateToken = this.tokenLimiter.isAllowed(registeredClient);
            if (!requiresGenerateToken) {
                OAuth2Error error = new OAuth2Error(OAuth2ErrorCodes.ACCESS_DENIED,
                        "The token generation fails, and the same client is prohibited from repeatedly obtaining the token within a short period of time.", null);
                throw new OAuth2AuthenticationException(error);
            }
        }

    }
}

注意：上述示例中我們使用OAuth 2.0 客戶端模式贮喧。

測(cè)試

本示例中我們限制access_token請(qǐng)求5分鐘響應(yīng)3次筒狠，我們將使用以下單元測(cè)試簡(jiǎn)單測(cè)試。

    @Test
    public void authorizationWhenObtainingTheAccessTokenSucceeds() throws Exception {
        MultiValueMap<String, String> parameters = new LinkedMultiValueMap<>();
        parameters.set(OAuth2ParameterNames.GRANT_TYPE, AuthorizationGrantType.CLIENT_CREDENTIALS.getValue());
        parameters.set(OAuth2ParameterNames.CLIENT_ID, "relive-client");
        parameters.set(OAuth2ParameterNames.CLIENT_SECRET, "relive-client");
        this.mockMvc.perform(post("/oauth2/token")
                .params(parameters))
                .andExpect(status().is2xxSuccessful());


    }

    @Test
    public void authorizationWhenTokenAccessRestrictionIsTriggeredThrowOAuth2AuthenticationException() throws Exception {
        MultiValueMap<String, String> parameters = new LinkedMultiValueMap<>();
        parameters.set(OAuth2ParameterNames.GRANT_TYPE, AuthorizationGrantType.CLIENT_CREDENTIALS.getValue());
        parameters.set(OAuth2ParameterNames.CLIENT_ID, "relive-client");
        parameters.set(OAuth2ParameterNames.CLIENT_SECRET, "relive-client");
        this.mockMvc.perform(post("/oauth2/token")
                .params(parameters))
                .andExpect(status().isBadRequest())
                .andExpect(result -> assertEquals("{\"error_description\":\"The token generation fails, and the same client is prohibited from repeatedly obtaining the token within a short period of time.\",\"error\":\"access_denied\"}", result.getResponse().getContentAsString()));
    }

結(jié)論

可能有人會(huì)有疑問(wèn)箱沦，一般服務(wù)都會(huì)由網(wǎng)關(guān)限流辩恼，為什么使用本示例中方式。當(dāng)然谓形，從實(shí)現(xiàn)上并不妨礙我們?cè)诰W(wǎng)關(guān)中進(jìn)行限制灶伊，這只是一個(gè)選擇問(wèn)題。后續(xù)文章中我將會(huì)介紹如何通過(guò)Spring Cloud Gateway結(jié)合授權(quán)服務(wù)對(duì)OAuth2客戶端進(jìn)行速率限制寒跳。

與往常一樣聘萨，本文中使用的源代碼可在 GitHub 上獲得。