Commit ac72988f authored by Nidhey Indurkar's avatar Nidhey Indurkar 💻
Browse files

fix(rate_limit): clear per rule and zero remaining while banned

parent 551c2165
Loading
Loading
Loading
Loading
+4 −3
Original line number Diff line number Diff line
@@ -138,9 +138,10 @@ login.

Every rule in the limiter is cleared, matched or not, because a caller
clearing state after a success knows the identifier rather than which rules
happened to match on the way in. Clearing an identifier with no state is a
no-op returning `0`, and a Redis failure fails open like `check`, returning
`0` and leaving the state to expire on its own.
happened to match on the way in. Each rule is deleted in its own call, since
rules do not share a Redis Cluster slot. Clearing an identifier with no state
is a no-op returning `0`, and a Redis failure fails open like `check`,
returning `0` and leaving the state to expire on its own.

## Identifier

+12 −5
Original line number Diff line number Diff line
@@ -154,15 +154,20 @@ module Labkit
      # Fails open like {#check}: an unreachable Redis leaves the state to
      # expire on its own rather than raising into the caller.
      def clear(identifier)
        keys = @rules.flat_map do |rule|
          next [] if rule.action == :skip
        key_groups = @rules.filter_map do |rule|
          next if rule.action == :skip

          [build_redis_key(rule, identifier), build_redis_key(rule, identifier, BAN_KEY_SUFFIX)]
        end

        return 0 if keys.empty?
        return 0 if key_groups.empty?

        @redis.with { |conn| conn.del(*keys) }
        # One DEL per rule. The hash tag holds the rule name, so a rule's own
        # keys share a slot but two rules do not, and Redis Cluster rejects a
        # DEL spanning slots.
        @redis.with do |conn|
          key_groups.sum { |group| conn.del(*group) }
        end
      rescue StandardError => e
        report_error_metrics
        log_error(e, identifier, nil)
@@ -351,7 +356,9 @@ module Labkit
        info = Result::Info.new(
          resolved_limit: resolved_limit, resolved_period: resolved_period,
          count: count,
          remaining: [resolved_limit - count, 0].max,
          # A ban outlives its counter, so count reads 0 once the window has
          # gone. Nothing is remaining while the ban still blocks.
          remaining: banned ? 0 : [resolved_limit - count, 0].max,
          reset_at: Time.now.utc + (banned ? ban_ttl : window_remaining)
        )

+26 −0
Original line number Diff line number Diff line
@@ -377,6 +377,18 @@ RSpec.describe Labkit::RateLimit do
      expect(lim.check({ ip: "1.2.3.4" }).action).to eq(:block)
    end

    it "reports nothing remaining while banned, even once the counter is gone", :aggregate_failures do
      lim = limiter(rules: [ban_rule])

      3.times { lim.check({ ip: "1.2.3.4" }) }
      raw_redis.del(counter_key)

      result = lim.check({ ip: "1.2.3.4" })

      expect(result.info.remaining).to eq(0)
      expect(result.to_response_headers["RateLimit-Remaining"]).to eq("0")
    end

    it "does not count while banned, so a banned caller cannot extend its own window" do
      lim = limiter(rules: [ban_rule])

@@ -458,6 +470,20 @@ RSpec.describe Labkit::RateLimit do
        expect(lim.clear({ ip: "1.2.3.4" })).to eq(1)
        expect(get_count("labkit:rl:{rack_request:plain:ip:1.2.3.4}")).to eq(0)
      end

      # Each rule sits under its own hash tag, so this is the case that has to
      # issue one DEL per rule rather than a single cross-slot one.
      it "clears every rule in a multi-rule limiter", :aggregate_failures do
        lim = limiter(rules: [ban_rule, rule(name: "plain", limit: 5, characteristics: [:ip])])

        3.times { lim.check({ ip: "1.2.3.4" }) }

        expect(lim.clear({ ip: "1.2.3.4" })).to eq(3)

        expect(raw_redis.exists?(ban_key)).to be(false)
        expect(get_count(counter_key)).to eq(0)
        expect(get_count("labkit:rl:{rack_request:plain:ip:1.2.3.4}")).to eq(0)
      end
    end

    # A :log rule with ban_for does the same accounting as the :limit version,