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

fix(rate_limit): tighten ban_for validation, count partial clears

parent ac72988f
Loading
Loading
Loading
Loading
+10 −3
Original line number Diff line number Diff line
@@ -140,8 +140,9 @@ 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. 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.
is a no-op returning `0`. A Redis failure fails open like `check`, leaving the
rest of the state to expire on its own, and the return value counts whatever
was removed before the failure.

## Identifier

@@ -173,7 +174,7 @@ A `Rule` is a `Data.define` value object with the following fields:
| `action`          | What the rule does when it matches. One of `:limit`, `:log`, `:skip`. Default `:limit`. See [Actions](#actions).                                                     |
| `characteristics` | Array of identifier keys whose values are folded into the Redis counter key. Each unique combination gets its own counter.                                           |
| `count_distinct`  | Optional identifier key. When set, the rule counts distinct values seen for that key within the window, backed by a Redis SET rather than a counter. Must not overlap `characteristics`. |
| `ban_for`         | Optional ban duration in seconds. Changes the accounting; `action` still decides who is blocked. Rejected on `:skip`. May be a callable resolved on every check.     |
| `ban_for`         | Optional ban duration, whole seconds, minimum 1. Changes the accounting; `action` still decides who is blocked. Rejected on `:skip`. May be a callable resolved on every check.     |

Making `limit` or `period` callable is the supported pattern for
runtime-tunable thresholds (e.g. feature flags or database-backed settings):
@@ -356,6 +357,12 @@ window, and it stops counting while that ban holds:
  while the next `check` re-bans immediately. Bans are normally longer than the
  window they are counted in.

  The duration is whole seconds and must be at least 1, since the ban is
  written with `SET EX`. A callable is checked again each time it resolves, and
  a value under a second raises rather than reaching Redis. Being an error, it
  fails open like any other, so a rule whose `ban_for` cannot resolve stops
  blocking entirely: watch `errors_total` after changing one.

### Redis keys

Each matched check writes a key shaped:
+20 −3
Original line number Diff line number Diff line
@@ -154,6 +154,10 @@ 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)
        # Counted as we go, so a connection lost mid-loop still reports the keys
        # that did go rather than claiming nothing was cleared.
        removed = 0

        key_groups = @rules.filter_map do |rule|
          next if rule.action == :skip

@@ -166,12 +170,14 @@ module Labkit
        # 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) }
          key_groups.each { |group| removed += conn.del(*group) }
        end

        removed
      rescue StandardError => e
        report_error_metrics
        log_error(e, identifier, nil)
        0
        removed
      end

      private
@@ -322,7 +328,7 @@ module Labkit
      def evaluate_ban_rule(rule, identifier, cost, rule_context)
        resolved_limit = Integer(resolve_value(rule.limit, rule_context))
        resolved_period = Integer(resolve_value(rule.period, rule_context))
        resolved_ban_for = Integer(resolve_value(rule.ban_for, rule_context))
        resolved_ban_for = resolve_ban_for(rule, rule_context)

        count, ttl, ban_ttl = ban_incr_with_ttl(
          build_redis_key(rule, identifier),
@@ -333,6 +339,17 @@ module Labkit
        build_ban_evaluation(rule, resolved_limit, resolved_period, count, ttl, ban_ttl)
      end

      # A callable ban_for is only checked at construction, so it can still hand
      # back something Redis would reject on the SET that writes the ban. Failing
      # here names ban_for rather than surfacing a line number from the script.
      def resolve_ban_for(rule, rule_context)
        raw = resolve_value(rule.ban_for, rule_context)
        seconds = Integer(raw, exception: false)
        raise ArgumentError, "ban_for resolved to #{raw.inspect}, need at least 1 second" if seconds.nil? || seconds < 1

        seconds
      end

      def peek_ban_rule(rule, identifier, rule_context)
        resolved_limit = Integer(resolve_value(rule.limit, rule_context))
        resolved_period = Integer(resolve_value(rule.period, rule_context))
+4 −2
Original line number Diff line number Diff line
@@ -49,9 +49,11 @@ module Labkit
        raise ArgumentError, "ban_for cannot be combined with count_distinct #{count_distinct_sym.inspect}" if count_distinct_sym

        return value if value.respond_to?(:call)
        return value if value.is_a?(Numeric) && value.positive?
        # Whole seconds: the ban is written with SET EX, and anything under a
        # second truncates to 0, which Redis rejects.
        return value if value.is_a?(Numeric) && value >= 1

        raise ArgumentError, "ban_for must be a positive Numeric or a callable, got #{value.inspect}"
        raise ArgumentError, "ban_for must be a Numeric of at least 1 second or a callable, got #{value.inspect}"
      end

      def self.normalize_count_distinct(value, characteristics_arr)
+9 −2
Original line number Diff line number Diff line
@@ -272,12 +272,19 @@ RSpec.describe Labkit::RateLimit::Rule do

    it "raises on a non-positive duration" do
      expect { valid_rule(action: :limit, ban_for: 0) }
        .to raise_error(ArgumentError, /positive Numeric or a callable/)
        .to raise_error(ArgumentError, /at least 1 second/)
    end

    # 0.5 is positive but truncates to 0 on the way to SET EX, which Redis
    # rejects, so the rule would fail open on the check that meant to ban.
    it "raises on a duration under one second" do
      expect { valid_rule(action: :limit, ban_for: 0.5) }
        .to raise_error(ArgumentError, /at least 1 second/)
    end

    it "raises on a non-numeric duration" do
      expect { valid_rule(action: :limit, ban_for: "900") }
        .to raise_error(ArgumentError, /positive Numeric or a callable/)
        .to raise_error(ArgumentError, /at least 1 second/)
    end

    it "raises when combined with count_distinct, which the ban path cannot honour" do
+33 −0
Original line number Diff line number Diff line
@@ -389,6 +389,21 @@ RSpec.describe Labkit::RateLimit do
      expect(result.to_response_headers["RateLimit-Remaining"]).to eq("0")
    end

    # Rule.new validates ban_for, but a callable is only resolved per check, so
    # an unusable duration can still arrive at the point of writing the ban.
    it "fails open without calling Redis when ban_for resolves below one second", :aggregate_failures do
      lim = limiter(rules: [ban_rule(ban_for: ->(_ctx) { 0 })])

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

      expect(result.error?).to be(true)
      expect(raw_redis.exists?(counter_key)).to be(false)
      expect(raw_redis.exists?(ban_key)).to be(false)
      expect(logger).to have_received(:warn)
        .with(hash_including("error_message" => /ban_for resolved to 0/)).at_least(:once)
    end

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

@@ -484,6 +499,24 @@ RSpec.describe Labkit::RateLimit do
        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

      it "reports the keys it did remove when Redis drops part way through", :aggregate_failures do
        rules = [ban_rule, rule(name: "plain", limit: 5, characteristics: [:ip])]
        3.times { limiter(rules: rules).check({ ip: "1.2.3.4" }) }

        dels = 0
        allow(raw_redis).to receive(:del).and_wrap_original do |original, *keys|
          dels += 1
          raise "connection reset" if dels > 1

          original.call(*keys)
        end

        expect(limiter(rules: rules).clear({ ip: "1.2.3.4" })).to eq(2)

        expect(raw_redis).to have_received(:del).twice
        expect(get_count("labkit:rl:{rack_request:plain:ip:1.2.3.4}")).to eq(2)
      end
    end

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