Commit 251bb2a0 authored by Sam Wiskow's avatar Sam Wiskow
Browse files

feat: add Labkit::RateLimit identifier and rules API (Stage 1a) 

Adds the foundational rate limiting interface for Phase 2 of GitLab's
unified rate limiting architecture:

- `Labkit::RateLimit.check(call_site:, identifier:, rules:)` evaluates
  an array of rules against a request identifier and returns :allow or
  :block
- `Identifier` value object carrying request context as key-value pairs
  (user, ip, namespace, plan, endpoint) with endpoint normalization
- `Rule` value object with match conditions, limit, period, action, and
  characteristics
- Per-rule independent Redis counters with key format:
  `labkit:rl:{call_site}:{rule_index}:{char_key}:{char_value}`
- SHA-256 hashing for char_value > 200 chars to prevent prefix collisions
- Fail-open on Redis errors with structured WARN logging
- Characteristic validation: raises in dev/test, sentinel + WARN in prod
- call_site validation: raises in dev/test, sanitizes in prod
- Structured per-rule JSON logging with redis_key for on-call triage

39 specs covering all acceptance criteria from
gl-infra/production-engineering#28784

Relates to gitlab-com/gl-infra#2021
parent 10231435

lib/gitlab-labkit.rb

+1 −0
Original line number Diff line number Diff line
@@ -18,6 +18,7 @@ module Labkit 
  autoload :Metrics, "labkit/metrics"

  autoload :Middleware, "labkit/middleware"

  autoload :Fields, "labkit/fields"

  autoload :RateLimit, "labkit/rate_limit"



  # Publishers to publish notifications whenever a HTTP reqeust is made.

  # A broadcasted notification's payload in topic "request.external_http" includes:

lib/labkit/rate_limit.rb

0 → 100644
+32 −0
Original line number Diff line number Diff line 
# frozen_string_literal: true



module Labkit

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

  module RateLimit

    autoload :Identifier, "labkit/rate_limit/identifier"

    autoload :Rule, "labkit/rate_limit/rule"

    autoload :Evaluator, "labkit/rate_limit/evaluator"



    KNOWN_CHARACTERISTICS = Evaluator::KNOWN_CHARACTERISTICS



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

    # configured rules.

    #

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

    # @param identifier [Identifier, Hash] caller attributes

    # @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

0 → 100644
+173 −0
Original line number Diff line number Diff line 
# frozen_string_literal: true



require "digest"

require "json"

require "logger"



module Labkit

  module RateLimit

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

    class Evaluator

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

      REDIS_KEY_PREFIX = "labkit:rl"

      CHAR_VALUE_MAX_LENGTH = 200

      UNKNOWN_SENTINEL = "unknown_characteristic"

      CALL_SITE_PATTERN = /\A[a-z0-9_]+\z/



      def initialize(call_site:, identifier:, rules:, redis:, logger: nil)

        @call_site = call_site

        @identifier = identifier

        @rules = rules

        @redis = redis

        @logger = logger || build_default_logger

      end



      def evaluate

        validate_call_site!



        aggregate = :allow



        @rules.each_with_index do |rule, index|

          next unless rule_matches?(rule, @identifier)



          begin

            result = evaluate_rule(rule, index)

            aggregate = :block if result == :block

          rescue ArgumentError

            raise

          rescue => e

            log_redis_error(e)

            return :allow

          end

        end



        aggregate

      end



      private



      def validate_call_site!

        return if CALL_SITE_PATTERN.match?(@call_site)



        if dev_or_test?

          raise ArgumentError, "Invalid call_site: #{@call_site.inspect}. Must match /\\A[a-z0-9_]+\\z/"

        else

          sanitized = @call_site.gsub(/[^a-z0-9_]/, "_")

          @logger.warn(JSON.generate(

            severity: "WARN",

            message: "rate_limit_invalid_call_site",

            call_site: @call_site,

            sanitized: sanitized

          ))

          @call_site = sanitized

        end

      end



      def rule_matches?(rule, identifier)

        rule.match.all? do |key, value|

          identifier[key] == value

        end

      end



      def evaluate_rule(rule, index)

        exceeded = false



        rule.characteristics.each do |char|

          char_value = resolve_characteristic(char, @identifier)

          redis_key = build_redis_key(@call_site, index, char, char_value)



          count = increment_counter(redis_key, rule.period)

          rule_exceeded = count > rule.limit



          exceeded = true if rule_exceeded



          log_rule(rule, index, count, redis_key, rule_exceeded)

        end



        (exceeded && rule.action == :block) ? :block : :allow

      end



      def resolve_characteristic(char, identifier)

        unless KNOWN_CHARACTERISTICS.include?(char)

          if dev_or_test?

            raise ArgumentError, "Unknown characteristic: #{char.inspect}. Known: #{KNOWN_CHARACTERISTICS.inspect}"

          else

            @logger.warn(JSON.generate(

              severity: "WARN",

              message: "rate_limit_unknown_characteristic",

              characteristic: char

            ))

            return UNKNOWN_SENTINEL

          end

        end



        value = identifier[char]



        # Normalize endpoint: strip query string

        value = Identifier.normalize_endpoint(value) if char == :endpoint



        value.to_s

      end



      def build_redis_key(call_site, rule_index, char, char_value)

        safe_value = encode_char_value(char_value.to_s)

        "#{REDIS_KEY_PREFIX}:#{call_site}:#{rule_index}:#{char}:#{safe_value}"

      end



      def encode_char_value(value)

        if value.length > CHAR_VALUE_MAX_LENGTH

          Digest::SHA256.hexdigest(value)

        else

          value

        end

      end



      def increment_counter(redis_key, period)

        count = @redis.incr(redis_key)

        # Set expiry only on first write (count == 1) to avoid resetting TTL on each call

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

        count

      end



      def log_rule(rule, index, count, redis_key, exceeded)

        entry = {

          severity: "INFO",

          message: "rate_limit_check",

          call_site: @call_site,

          rule_index: index,

          action: rule.action.to_s,

          limit: rule.limit,

          period: rule.period,

          count: count,

          matched: true,

          exceeded: exceeded,

          identifier: @identifier.to_h,

          redis_key: redis_key

        }

        @logger.info(JSON.generate(entry))

      end



      def log_redis_error(error)

        entry = {

          severity: "WARN",

          message: "rate_limit_redis_error",

          call_site: @call_site,

          error: error.class.to_s,

          result: "allow"

        }

        @logger.warn(JSON.generate(entry))

      end



      def dev_or_test?

        env = ENV["LABKIT_ENV"]

        env == "test" || env == "development"

      end



      def build_default_logger

        logger = Logger.new($stdout)

        logger.formatter = proc { |_sev, _dt, _prog, msg| "#{msg}\n" }

        logger

      end

    end

  end

end

lib/labkit/rate_limit/identifier.rb

0 → 100644
+36 −0
Original line number Diff line number Diff line 
# frozen_string_literal: true



module Labkit

  module RateLimit

    # Identifier is a value object wrapping a hash of key-value pairs that

    # describe the caller (e.g. user, ip, endpoint).

    class Identifier

      attr_reader :attributes



      def initialize(attributes = {})

        @attributes = attributes.transform_keys(&:to_sym).freeze

      end



      # Return the value for a characteristic key.

      def [](key)

        @attributes[key.to_sym]

      end



      # Serialize to a plain Hash suitable for JSON logging.

      def to_h

        @attributes.transform_keys(&:to_s)

      end



      def ==(other)

        other.is_a?(Identifier) && other.attributes == @attributes

      end



      # Normalize an endpoint value: strip query string.

      def self.normalize_endpoint(value)

        return value unless value.is_a?(String)



        value.split("?", 2).first

      end

    end

  end

end

lib/labkit/rate_limit/rule.rb

0 → 100644
+18 −0
Original line number Diff line number Diff line 
# frozen_string_literal: true



module Labkit

  module RateLimit

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

    Rule = Data.define(:match, :limit, :period, :action, :characteristics) do

      def initialize(match:, limit:, period:, action:, characteristics:)

        super(

          match: match.transform_keys(&:to_sym).freeze,

          limit: limit,

          period: period,

          action: action.to_sym,

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

        )

      end

    end

  end

end