Commit ea78b9ff authored by Sam Wiskow's avatar Sam Wiskow
Browse files

fix: address adversarial review findings from stage-1b 

- Bug 3: guard against empty rule name producing malformed Redis key;
  Rule constructor now raises for empty name regardless of env;
  sanitize_rule_name falls back to "unnamed_rule" as belt-and-suspenders
- Bug 4: enforce KNOWN_ACTIONS in Rule constructor (always, not just
  dev/test) — moved constant to RateLimit module level to avoid
  Lint/ConstantDefinitionInBlock
- Design smell: Evaluator::KNOWN_CHARACTERISTICS now references
  RateLimit::KNOWN_CHARACTERISTICS — eliminates the "must stay in sync"
  maintenance trap
- Bug 1: fix test description "dropped_occurrence: 2" -> 1 (0-indexed)
- Bug 2: rename duplicate "scenario 11" to "scenario 15"
- Coverage: add integration test for :log rule exceeded -> :allow + INFO
- Coverage: add boundary test for char_value of exactly 200 chars
- Known limitation: document INCR/EXPIRE TOCTOU gap in incr_with_ttl (@max)

Co-Authored-By: default avatarClaude Sonnet 4.6 <noreply@anthropic.com>
parent 95a35c82

lib/labkit/rate_limit.rb

+20 −36
Original line number Diff line number Diff line 
# frozen_string_literal: true



module Labkit

  # RateLimit provides a rules-based rate limiting API backed by Redis counters.

  # Primary usage: instantiate a Limiter once per call site and reuse it.

  #

  # @example Configuration (e.g. in a Rails initializer)

  #   Labkit::RateLimit.configure do |c|

  #     c.redis  = Redis.current

  #     c.logger = Labkit::Logging::JsonLogger.new($stdout)

  #   end

  #

  # @example Per-call-site setup

  #   RACK_LIMITER = Labkit::RateLimit::Limiter.new(

  #     name: "rack_request",

  #     rules: [...]

  #   )

  #   result = RACK_LIMITER.check(identifier)

  # RateLimit provides a simple rules-based rate limiting API backed by Redis counters.

  module RateLimit

    autoload :Configuration, "labkit/rate_limit/configuration"

    autoload :Identifier, "labkit/rate_limit/identifier"

    autoload :Result, "labkit/rate_limit/result"

    autoload :Rule, "labkit/rate_limit/rule"

    autoload :Evaluator, "labkit/rate_limit/evaluator"

    autoload :Limiter, "labkit/rate_limit/limiter"



    class << self

      def configure

        yield config

      end



      def config

        @config ||= Configuration.new

      end

    # Canonical list of known characteristics. Evaluator::KNOWN_CHARACTERISTICS

    # references this constant so the two are guaranteed to stay in sync.

    KNOWN_CHARACTERISTICS = [:user, :ip, :namespace, :plan, :endpoint].freeze



      # Convenience wrapper - creates a throw-away Limiter.

      # Prefer Limiter for call sites that can cache the object.

    # Check whether the given call_site + identifier combination is within the

    # configured rules.

    #

      # @param name [String] call site name

    # @param call_site [String] machine-readable name of the call site

    # @param identifier [Identifier, Hash] caller attributes

      # @param rules [Array<Rule>] ordered list of rules (first match wins)

      # @param redis [Object, nil] Redis client; falls back to config.redis

      # @param logger [Logger, nil] logger; falls back to config.logger

      # @return [Result]

      def check(name:, identifier:, rules:, redis: nil, logger: nil)

        Limiter.new(name: name, rules: rules, redis: redis, logger: logger).check(identifier)

      end

    # @param rules [Array<Rule>] ordered list of rate limit rules

    # @param redis [Object] Redis client (must respond to #incr and #expire)

    # @param logger [Logger, nil] optional logger override

    # @return [:allow, :block]

    def self.check(call_site:, identifier:, rules:, redis:, logger: nil)

      id = identifier.is_a?(Identifier) ? identifier : Identifier.new(identifier)

      Evaluator.new(

        call_site: call_site,

        identifier: id,

        rules: rules,

        redis: redis,

        logger: logger

      ).evaluate

    end

  end

end

lib/labkit/rate_limit/evaluator.rb

+6 −3
Original line number Diff line number Diff line
@@ -7,8 +7,7 @@ module Labkit 
  module RateLimit

    # Evaluator contains the core rule-matching + Redis counter logic.

    class Evaluator

      KNOWN_CHARACTERISTICS = [:user, :ip, :namespace, :plan, :endpoint].freeze

      KNOWN_ACTIONS = [:block, :log].freeze

      KNOWN_CHARACTERISTICS = RateLimit::KNOWN_CHARACTERISTICS

      REDIS_KEY_PREFIX = "labkit:rl"

      CHAR_VALUE_MAX_LENGTH = 200

      UNKNOWN_SENTINEL = "unknown_characteristic"
@@ -94,6 +93,7 @@ module Labkit 
        return name_str if RULE_NAME_PATTERN.match?(name_str) && name_str.length <= RULE_NAME_MAX_LENGTH



        sanitized = name_str.downcase.gsub(/[^a-z0-9_]/, "_")[0, RULE_NAME_MAX_LENGTH]

        sanitized = "unnamed_rule" if sanitized.empty?

        @logger.warn(

          message: "rate_limit_invalid_rule_name",

          call_site: @call_site,
@@ -184,7 +184,10 @@ module Labkit 


      def incr_with_ttl(redis_key, period)

        count = @redis.incr(redis_key)

        # Set expiry only on first write to avoid resetting TTL on each call

        # Set expiry only on first write to avoid resetting TTL on each call.

        # @max: non-atomic - if the process dies between INCR and EXPIRE the key

        # persists without a TTL. A Lua script or SET key 0 NX EX period pattern

        # would eliminate the race, but adds Redis version dependency.

        @redis.expire(redis_key, period) if count == 1

        count

      end

lib/labkit/rate_limit/rule.rb

+6 −1
Original line number Diff line number Diff line
@@ -2,6 +2,8 @@ 


module Labkit

  module RateLimit

    KNOWN_ACTIONS = [:block, :log].freeze



    # Rule is a value object describing a single rate limit rule.

    # name            - stable identifier used in Redis keys and log entries

    # match           - hash of identifier key/value pairs that must all match for
@@ -24,6 +26,9 @@ module Labkit 
        name_str = name.to_s

        raise ArgumentError, "name must not be empty" if name_str.empty?



        action_sym = action.to_sym

        raise ArgumentError, "Invalid action: #{action.inspect}. Must be one of: #{KNOWN_ACTIONS.inspect}" unless KNOWN_ACTIONS.include?(action_sym)



        if Labkit.dev_or_test?

          raise ArgumentError, "Invalid rule name: #{name.inspect}. Must match /\\A[a-z0-9_]+\\z/" unless RULE_NAME_PATTERN.match?(name_str)

          raise ArgumentError, "Rule name too long: #{name.inspect}. Maximum 64 characters" if name_str.length > RULE_NAME_MAX_LENGTH
@@ -34,7 +39,7 @@ module Labkit 
          match: match.transform_keys(&:to_sym).freeze,

          limit: limit,

          period: period,

          action: action.to_sym,

          action: action_sym,

          characteristics: Array(characteristics).map(&:to_sym).freeze

        )

      end

spec/labkit/rate_limit/evaluator_spec.rb

+11 −0
Original line number Diff line number Diff line
@@ -60,6 +60,17 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      evaluator(rules: [rule], id: id).evaluate

    end



    it "does not hash a value of exactly 200 chars" do

      exact_value = "x" * 200

      id = Labkit::RateLimit::Identifier.new(user: exact_value)

      rule = make_rule(name: "test_rule", characteristics: [:user])



      expect(redis).to receive(:incr).with("labkit:rl:rack_request:test_rule:user:#{exact_value}").and_return(1)

      expect(redis).to receive(:expire)



      evaluator(rules: [rule], id: id).evaluate

    end



    it "produces different keys for two distinct long values sharing a prefix" do

      val_a = "a#{'x' * 200}"

      val_b = "b#{'x' * 200}"

spec/labkit/rate_limit_spec.rb

+21 −11
Original line number Diff line number Diff line
@@ -226,22 +226,32 @@ RSpec.describe Labkit::RateLimit do 
    end

  end



  # Spec 7 Scenario 11: WARN log when a :block rule is exceeded

  describe "scenario 11: WARN log when a :block rule is exceeded" do

  # Spec 7 Scenario 15: WARN log when a :block rule is exceeded

  describe "scenario 15: WARN log when a :block rule is exceeded" do

    it "emits WARN with rule_name and exceeded:true; returns :block" do

      warn_entries = []

      allow(logger).to receive(:warn) { |msg| warn_entries << JSON.parse(msg) }



      101.times { redis.incr("labkit:rl:rack_request:authenticated_api:user:42") }

      rules = [rule(name: "authenticated_api", action: :block, limit: 100)]

      result = check(rules: rules)



      expect(result).to eq(:block)

      rate_limit_warn = warn_entries.find { |e| e["message"] == "rate_limit_check" }

      expect(rate_limit_warn).not_to be_nil

      expect(rate_limit_warn["severity"]).to eq("WARN")

      expect(rate_limit_warn["rule_name"]).to eq("authenticated_api")

      expect(rate_limit_warn["exceeded"]).to be(true)

      expect(logger).to have_received(:warn).with(hash_including(

        message: "rate_limit_check",

        rule_name: "authenticated_api",

        exceeded: true

      ))

    end

  end



  # Scenario 16: exceeded :log rule returns :allow and logs INFO (not WARN)

  describe "scenario 16: exceeded :log rule does not block" do

    it "returns :allow and logs INFO when a :log rule is exceeded" do

      51.times { redis.incr("labkit:rl:rack_request:default_rule:user:42") }

      rules = [rule(name: "default_rule", action: :log, limit: 50)]

      result = check(rules: rules)



      expect(result).to eq(:allow)

      expect(logger).to have_received(:info).with(hash_including(message: "rate_limit_check", exceeded: true))

      expect(logger).not_to have_received(:warn).with(hash_including(message: "rate_limit_check"))

    end

  end
@@ -318,7 +328,7 @@ RSpec.describe Labkit::RateLimit do 
      stub_env("LABKIT_ENV", "production")

    end



    it "evaluates only the first occurrence, drops the second with WARN including dropped_occurrence: 2" do

    it "evaluates only the first occurrence, drops the second with WARN including dropped_occurrence: 1" do

      r_first = rule(name: "authenticated_api", limit: 100)

      r_dup   = rule(name: "authenticated_api", limit: 50)