Merge branch 'rate-limit/stage-1b' into 'master' (8eea3315) · Commits · GitLab.org / ruby / gems / labkit-ruby

lib/labkit/rate_limit/limiter.rb

+59 −7

Original line number	Diff line number	Diff line
		@@ -20,14 +20,15 @@ module Labkit
		NAME_PATTERN = /\A[a-z0-9_]+\z/

		def initialize(name:, rules:, redis: nil, logger: nil)
		resolved_logger = logger \|\| RateLimit.config.logger \|\| Labkit::Logging::JsonLogger.new($stdout)
		validated_name = validate_name!(name, resolved_logger)
		@logger = logger \|\| RateLimit.config.logger \|\| Labkit::Logging::JsonLogger.new($stdout)
		validated_name = validate_name!(name)
		@name = validated_name

		@evaluator = Evaluator.new(
		name: validated_name,
		rules: rules,
		rules: prepare_rules(rules),
		redis: redis \|\| RateLimit.config.redis,
		logger: resolved_logger
		logger: @logger
		)
		end

		@@ -35,19 +36,70 @@ module Labkit
		# @return [Result]
		def check(identifier)
		id = identifier.is_a?(Identifier) ? identifier : Identifier.new(identifier)
		@evaluator.check(id)
		result = @evaluator.check(id)

		if result.exceeded? && result.action == :block
		@logger.warn(
		message: "rate_limit_check",
		name: @name,
		rule_name: result.rule.name,
		exceeded: true,
		severity: "WARN"
		)
		end

		result
		end

		private

		def validate_name!(name, logger)
		def validate_name!(name)
		raise ArgumentError, "name must be a non-empty String" unless name.is_a?(String) && !name.empty?
		return name if NAME_PATTERN.match?(name)

		raise ArgumentError, "Invalid name: #{name.inspect}. Must match /\\A[a-z0-9_]+\\z/" if Labkit.dev_or_test?

		sanitized = name.gsub(/[^a-z0-9_]/, "_")
		logger.warn(message: "rate_limit_invalid_name", name: name, sanitized: sanitized)
		@logger.warn(message: "rate_limit_invalid_name", name: name, sanitized: sanitized)
		sanitized
		end

		# Validates and deduplicates rule names before passing rules to Evaluator.
		# In dev/test: raises on invalid format or duplicate names.
		# In production: sanitizes invalid names (WARN) and drops duplicates (WARN, first wins).
		# Returns an array of rules with sanitized names.
		def prepare_rules(rules)
		seen = {}
		rules.each_with_index.filter_map do \|rule, idx\|
		sanitized = sanitize_rule_name(rule.name)

		if seen.key?(sanitized)
		raise ArgumentError, "Duplicate rule name #{sanitized.inspect} at index #{idx}" if Labkit.dev_or_test?

		@logger.warn(
		message: "rate_limit_duplicate_rule_name",
		name: sanitized,
		dropped_occurrence: idx
		)
		next nil
		end

		seen[sanitized] = true
		sanitized == rule.name ? rule : rule.with(name: sanitized) # rubocop:disable CodeReuse/ActiveRecord
		end
		end

		def sanitize_rule_name(name)
		s = name.to_s
		return s if RULE_NAME_PATTERN.match?(s) && s.length <= RULE_NAME_MAX_LENGTH

		sanitized = s.downcase.gsub(/[^a-z0-9_]/, "_")[0, RULE_NAME_MAX_LENGTH]
		sanitized = "unnamed_rule" if sanitized.empty?
		@logger.warn(
		message: "rate_limit_invalid_rule_name",
		original_name: s,
		sanitized_name: sanitized
		)
		sanitized
		end
		end

lib/labkit/rate_limit/rule.rb

+23 −2

Original line number	Diff line number	Diff line
		@@ -2,6 +2,10 @@

		module Labkit
		module RateLimit
		KNOWN_ACTIONS = [:block, :log].freeze
		RULE_NAME_PATTERN = /\A[a-z0-9_]+\z/
		RULE_NAME_MAX_LENGTH = 64

		# 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
		@@ -10,14 +14,31 @@ module Labkit
		# period - window in seconds; may be a callable (resolved per check)
		# action - :block (enforce) or :log (count and log, but do not block)
		# characteristics - identifier keys used to build the compound Redis counter key
		#
		# +name+ must be a lowercase alphanumeric-and-underscore string of at most 64
		# characters. It is used as the middle segment of every Redis counter key for
		# this rule, so changing a rule's name mid-window abandons its in-flight counters.
		Rule = Data.define(:name, :match, :limit, :period, :action, :characteristics) do
		def initialize(name:, limit:, period:, characteristics:, match: {}, action: :block)
		raise ArgumentError, "name must be a String or Symbol, got #{name.class}" unless name.is_a?(String) \|\| name.is_a?(Symbol)

		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
		end

		super(
		name: name.to_s.tr(":", "_"),
		name: name_str.freeze,
		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/limiter_spec.rb

+111 −0

Original line number	Diff line number	Diff line
		@@ -106,4 +106,115 @@ RSpec.describe Labkit::RateLimit::Limiter do
		lim.check(id)
		end
		end

		describe "Scenario A: evaluator is reused across checks" do
		it "returns the same evaluator object on repeated calls" do
		lim = limiter
		evaluator = lim.instance_variable_get(:@evaluator)
		lim.check({ user: 1 })
		lim.check({ user: 2 })
		expect(lim.instance_variable_get(:@evaluator)).to equal(evaluator)
		end
		end

		describe "Scenario B: exceeded :block rule emits rate_limit_check WARN" do
		it "logs warn with rule_name and exceeded: true" do
		allow(redis).to receive(:incr).and_return(101)
		lim = limiter
		lim.check({ user: 42 })
		expect(logger).to have_received(:warn).with(
		hash_including(
		message: "rate_limit_check",
		rule_name: "default",
		exceeded: true
		)
		)
		end
		end

		describe "Scenario C: exceeded :log rule does not emit rate_limit_check WARN" do
		it "does not call warn with rate_limit_check message" do
		allow(redis).to receive(:incr).and_return(101)
		r = rule(action: :log)
		lim = limiter(rules: [r])
		lim.check({ user: 42 })
		expect(logger).not_to have_received(:warn).with(hash_including(message: "rate_limit_check"))
		end
		end

		describe "prepare_rules: rule name deduplication and sanitization" do
		context "when in dev/test environment" do
		it "raises ArgumentError for duplicate rule names" do
		r1 = rule(name: "api_user")
		r2 = rule(name: "api_user")
		expect { limiter(rules: [r1, r2]) }
		.to raise_error(ArgumentError, /Duplicate rule name."api_user".index 1/)
		end
		end

		context "when in production environment" do
		before do
		stub_env("RAILS_ENV", "production")
		end

		it "drops second occurrence of duplicate rule name with WARN" do
		r1 = rule(name: "api_user")
		r2 = rule(name: "api_user")
		limiter(rules: [r1, r2])
		expect(logger).to have_received(:warn).with(
		hash_including(
		message: "rate_limit_duplicate_rule_name",
		name: "api_user",
		dropped_occurrence: 1
		)
		)
		end

		it "sanitizes invalid rule name and emits WARN with original and sanitized names" do
		r = Labkit::RateLimit::Rule.new(
		name: "Bad Name!", limit: 100, period: 60, characteristics: [:user]
		)
		limiter(rules: [r])
		expect(logger).to have_received(:warn).with(
		hash_including(
		message: "rate_limit_invalid_rule_name",
		original_name: "Bad Name!",
		sanitized_name: "bad_name_"
		)
		)
		end

		it "truncates rule name exceeding 64 characters with WARN" do
		long_name = "a" * 65
		r = Labkit::RateLimit::Rule.new(
		name: long_name, limit: 100, period: 60, characteristics: [:user]
		)
		limiter(rules: [r])
		expect(logger).to have_received(:warn).with(
		hash_including(
		message: "rate_limit_invalid_rule_name",
		original_name: long_name,
		sanitized_name: "a" * 64
		)
		)
		end

		it "treats post-sanitize collision as duplicate and drops second rule with WARN" do
		r1 = Labkit::RateLimit::Rule.new(
		name: "bad name", limit: 100, period: 60, characteristics: [:user]
		)
		r2 = Labkit::RateLimit::Rule.new(
		name: "bad!name", limit: 100, period: 60, characteristics: [:user]
		)
		limiter(rules: [r1, r2])
		expect(logger).to have_received(:warn).with(
		hash_including(
		message: "rate_limit_duplicate_rule_name",
		name: "bad_name",
		dropped_occurrence: 1
		)
		)
		end
		end
		end
		end

spec/labkit/rate_limit/rule_spec.rb

0 → 100644

+143 −0

Original line number	Diff line number	Diff line
		# frozen_string_literal: true

		require "spec_helper"

		RSpec.describe Labkit::RateLimit::Rule do
		include StubENV

		before do
		stub_env("RAILS_ENV", "test")
		end

		def valid_rule(name: "api_user", limit: 100, period: 60, characteristics: [:user], **rest)
		described_class.new(name: name, limit: limit, period: period, characteristics: characteristics, **rest)
		end

		describe "name type validation (always enforced)" do
		it "raises when name is nil" do
		expect { valid_rule(name: nil) }
		.to raise_error(ArgumentError, /name must be a String or Symbol/)
		end

		it "raises when name is an Integer" do
		expect { valid_rule(name: 42) }
		.to raise_error(ArgumentError, /name must be a String or Symbol/)
		end

		it "accepts a Symbol name and converts it to String" do
		rule = valid_rule(name: :api_user)
		expect(rule.name).to eq("api_user")
		end

		it "raises when name is an empty string" do
		expect { valid_rule(name: "") }.to raise_error(ArgumentError, /must not be empty/)
		end

		it "raises when Symbol resolves to empty string" do
		expect { valid_rule(name: :"") }.to raise_error(ArgumentError, /must not be empty/)
		end
		end

		describe "name format and length validation" do
		context "when in dev/test environment" do
		it "raises on names with invalid characters" do
		expect { valid_rule(name: "Bad Name!") }
		.to raise_error(ArgumentError, /Invalid rule name/)
		end

		it "raises on names with uppercase letters" do
		expect { valid_rule(name: "ApiUser") }
		.to raise_error(ArgumentError, /Invalid rule name/)
		end

		it "raises on names with hyphens" do
		expect { valid_rule(name: "api-user") }
		.to raise_error(ArgumentError, /Invalid rule name/)
		end

		it "raises when name exceeds 64 characters" do
		expect { valid_rule(name: "a" * 65) }
		.to raise_error(ArgumentError, /Rule name too long/)
		end

		it "accepts a name at exactly 64 characters" do
		rule = valid_rule(name: "a" * 64)
		expect(rule.name.length).to eq(64)
		end
		end

		context "when in production environment" do
		before do
		stub_env("RAILS_ENV", "production")
		end

		it "does not raise on invalid format (deferred to Limiter sanitization)" do
		expect { valid_rule(name: "Bad Name!") }.not_to raise_error
		end

		it "does not raise when name exceeds 64 characters" do
		expect { valid_rule(name: "a" * 65) }.not_to raise_error
		end
		end
		end

		describe "action validation (always enforced)" do
		it "accepts :block" do
		expect(valid_rule(action: :block).action).to eq(:block)
		end

		it "accepts :log" do
		expect(valid_rule(action: :log).action).to eq(:log)
		end

		it "accepts string action and coerces to symbol" do
		expect(valid_rule(action: "block").action).to eq(:block)
		end

		it "raises on unknown action" do
		expect { valid_rule(action: :deny) }
		.to raise_error(ArgumentError, /Invalid action/)
		end
		end

		describe "match and characteristics normalization" do
		it "symbolizes match keys" do
		rule = valid_rule(match: { "endpoint" => "/api/v4" })
		expect(rule.match).to eq({ endpoint: "/api/v4" })
		end

		it "symbolizes characteristics" do
		rule = valid_rule(characteristics: %w[user ip])
		expect(rule.characteristics).to eq(%i[user ip])
		end

		it "wraps a single characteristic in an array" do
		rule = valid_rule(characteristics: :user)
		expect(rule.characteristics).to eq([:user])
		end

		it "defaults match to empty hash" do
		rule = valid_rule
		expect(rule.match).to eq({})
		end

		it "defaults action to :block" do
		rule = valid_rule
		expect(rule.action).to eq(:block)
		end
		end

		describe "immutability" do
		it "freezes the name" do
		expect(valid_rule.name).to be_frozen
		end

		it "freezes the match hash" do
		expect(valid_rule.match).to be_frozen
		end

		it "freezes the characteristics array" do
		expect(valid_rule.characteristics).to be_frozen
		end
		end
		end

spec/labkit/rate_limit_spec.rb

+20 −0

Original line number	Diff line number	Diff line
		@@ -254,4 +254,24 @@ RSpec.describe Labkit::RateLimit do
		expect(result.matched?).to be(true)
		end
		end

		describe "Scenario Q extension: rule name stability across array reordering" do
		it "accumulates counters by rule name regardless of rule array order" do
		r_a = rule(name: "rule_a", match: {}, characteristics: [:user], limit: 100)
		r_b = rule(name: "rule_b", match: { ip: "1.2.3.4" }, characteristics: [:ip], limit: 100)

		limiter(rules: [r_a, r_b]).check({ user: 42, ip: "1.2.3.4" })
		limiter(rules: [r_b, r_a]).check({ user: 42, ip: "1.2.3.4" })

		expect(redis.get("labkit:rl:rack_request:rule_a:user:42")).to eq(1)
		expect(redis.get("labkit:rl:rack_request:rule_b:ip:1.2.3.4")).to eq(1)
		end

		it "never writes integer-index key segments" do
		r = rule(name: "named_rule", characteristics: [:user])
		limiter(rules: [r]).check({ user: 42 })

		expect(redis.get("labkit:rl:rack_request:0:user:42")).to eq(0)
		end
		end
		end