Merge branch 'rate-limit/stage-1a' into 'master'

  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:

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

    # Defined independently to avoid forcing eager load of Evaluator at module load time.

    # Must stay in sync with Evaluator::KNOWN_CHARACTERISTICS.

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

    # 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

# frozen_string_literal: true

require "openssl"

require "labkit/logging/json_logger"

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

      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!

        evaluate_rules

      rescue ArgumentError

        raise

      rescue StandardError => e

        # Intentionally broad: fail-open applies to any unexpected error (network,

        # timeout, OOM, etc.), not only Redis protocol errors.

        log_evaluate_error(e)

        :allow

      end

      private

      def evaluate_rules

        aggregate = :allow

        @rules.each_with_index do |rule, index|

          next unless rule_matches?(rule, @identifier)

          result = evaluate_rule(rule, index)

          aggregate = :block if result == :block

        end

        aggregate

      end

      def validate_call_site!

        return if CALL_SITE_PATTERN.match?(@call_site)

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

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

        @logger.warn(

          message: "rate_limit_invalid_call_site",

          call_site: @call_site,

          sanitized: sanitized

        )

        @call_site = sanitized

      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)

          if char_value.nil?

            log_skipped_characteristic(rule, index, char)

            next

          end

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

          count = incr_with_ttl(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)

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

          @logger.warn(

            message: "rate_limit_unknown_characteristic",

            characteristic: char

          )

          return UNKNOWN_SENTINEL

        end

        value = identifier[char]

        # Normalize endpoint: strip query string

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

        # Treat nil and empty-string the same: anonymous traffic must not collide on a shared bucket.

        return nil if value.nil? || value.to_s.empty?

        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

          OpenSSL::Digest::SHA256.hexdigest(value)

        else

          value

        end

      end

      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

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

        count

      end

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

        @logger.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

        )

      end

      def log_skipped_characteristic(rule, index, char)

        @logger.info(

          message: "rate_limit_check",

          call_site: @call_site,

          rule_index: index,

          action: rule.action.to_s,

          limit: rule.limit,

          period: rule.period,

          characteristic: char,

          matched: true,

          skipped: true,

          identifier: @identifier.to_h

        )

      end

      def log_evaluate_error(error)

        @logger.warn(

          message: "rate_limit_redis_error",

          call_site: @call_site,

          error: error.class.to_s,

          result: "allow"

        )

      end

      def dev_or_test?

        # Memoized: ENV access is not free under concurrency.

        return @dev_or_test unless @dev_or_test.nil?

        env = ENV.fetch("LABKIT_ENV", nil)

        @dev_or_test = env == "test" || env == "development"

      end

      def build_default_logger

        Labkit::Logging::JsonLogger.new($stdout)

      end

    end

  end

end

# 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

      # Normalize an endpoint value: strip query string.

      def self.normalize_endpoint(value)

        return value unless value.is_a?(String)

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

      end

      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

    end

  end

end

# 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(limit:, period:, characteristics:, match: {}, action: :block)

        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