throttle_machines 0.0.0 → 0.1.1

data/lib/throttle_machines/rack_middleware/configuration.rb

@@ -0,0 +1,103 @@
+# frozen_string_literal: true
+
+require 'ipaddr'
+
+module ThrottleMachines
+  class RackMiddleware
+    class Configuration
+      DEFAULT_BLOCKLISTED_RESPONDER = ->(_req) { [403, { 'content-type' => 'text/plain' }, ["Forbidden\n"]] }
+
+      DEFAULT_THROTTLED_RESPONDER = lambda do |req|
+        match_data = req.env['rack.attack.match_data']
+        retry_after = match_data[:retry_after] || 60
+
+        [429, { 'content-type' => 'text/plain', 'retry-after' => retry_after.to_s }, ["Retry later\n"]]
+      end
+
+      attr_reader :throttles, :tracks, :safelists, :blocklists
+      attr_accessor :throttled_responder, :blocklisted_responder
+
+      def initialize
+        @throttles = {}
+        @tracks = {}
+        @safelists = {}
+        @blocklists = {}
+        @anonymous_safelists = []
+        @anonymous_blocklists = []
+        @fail2bans = {}
+        @allow2bans = {}
+
+        @throttled_responder = DEFAULT_THROTTLED_RESPONDER
+        @blocklisted_responder = DEFAULT_BLOCKLISTED_RESPONDER
+      end
+
+      # DSL Methods
+      def throttle(name, options = {}, &)
+        @throttles[name] = Throttle.new(name, options, &)
+      end
+
+      def track(name, options = {}, &)
+        @tracks[name] = Track.new(name, options, &)
+      end
+
+      def safelist(name = nil, &)
+        if name
+          @safelists[name] = Safelist.new(name, &)
+        else
+          @anonymous_safelists << Safelist.new(nil, &)
+        end
+      end
+
+      def blocklist(name = nil, &)
+        if name
+          @blocklists[name] = Blocklist.new(name, &)
+        else
+          @anonymous_blocklists << Blocklist.new(nil, &)
+        end
+      end
+
+      def safelist_ip(ip_address)
+        @anonymous_safelists << Safelist.new(nil) { |req| req.ip == ip_address }
+      end
+
+      def blocklist_ip(ip_address)
+        @anonymous_blocklists << Blocklist.new(nil) { |req| req.ip == ip_address }
+      end
+
+      def fail2ban(name, options = {}, &)
+        @fail2bans[name] = Fail2Ban.new(name, options, &)
+      end
+
+      def allow2ban(name, options = {}, &)
+        @allow2bans[name] = Allow2Ban.new(name, options, &)
+      end
+
+      # Check methods
+      def safelisted?(request)
+        @anonymous_safelists.any? { |safelist| safelist.matched_by?(request) } ||
+          @safelists.values.any? { |safelist| safelist.matched_by?(request) }
+      end
+
+      def blocklisted?(request)
+        # Check explicit blocklists
+        return true if @anonymous_blocklists.any? { |blocklist| blocklist.matched_by?(request) }
+        return true if @blocklists.values.any? { |blocklist| blocklist.matched_by?(request) }
+
+        # Check fail2bans
+        @fail2bans.values.any? { |fail2ban| fail2ban.banned?(request) }
+      end
+
+      def throttled?(request)
+        # Process allow2bans first (they can reset fail2ban counters)
+        @allow2bans.each_value { |allow2ban| allow2ban.matched_by?(request) }
+
+        # Check throttles
+        @throttles.values.any? { |throttle| throttle.matched_by?(request) }
+      end
+
+      def tracked?(request)
+        @tracks.each_value { |track| track.matched_by?(request) }
+      end
+    end
+  end
+end

data/lib/throttle_machines/rack_middleware/fail2_ban.rb

@@ -0,0 +1,87 @@
+# frozen_string_literal: true
+
+module ThrottleMachines
+  class RackMiddleware
+    class Fail2Ban
+      attr_reader :name, :maxretry, :findtime, :bantime, :block
+
+      def initialize(name, options, &block)
+        @name = name
+        @block = block
+
+        @maxretry = options[:maxretry] || 5
+        @findtime = options[:findtime] || 60
+        @bantime = options[:bantime] || 300
+      end
+
+      def banned?(request)
+        discriminator = discriminator_for(request)
+        return false unless discriminator
+
+        key = "fail2ban:#{@name}:#{discriminator}"
+
+        # Use circuit breaker to track failures
+        breaker = ThrottleMachines::Breaker.new(
+          key,
+          failure_threshold: @maxretry,
+          timeout: @bantime,
+          storage: ThrottleMachines.storage
+        )
+
+        # Check if circuit is open (banned)
+        if breaker.open?
+          # Get breaker state for instrumentation
+          state = breaker.to_h
+
+          request.env['rack.attack.matched'] = @name
+          request.env['rack.attack.match_type'] = :fail2ban
+          request.env['rack.attack.match_discriminator'] = discriminator
+          request.env['rack.attack.match_data'] = {
+            discriminator: discriminator,
+            maxretry: @maxretry,
+            findtime: @findtime,
+            bantime: @bantime,
+            failures: state[:failure_count],
+            time_until_unban: state[:time_until_retry]
+          }
+
+          ThrottleMachines::RackMiddleware.instrument(request)
+          true
+        else
+          false
+        end
+      end
+
+      def count(request)
+        discriminator = discriminator_for(request)
+        return unless discriminator
+
+        key = "fail2ban:#{@name}:#{discriminator}"
+
+        # Use the breaker to record a failure if block returns true
+        return unless yield
+
+        breaker = ThrottleMachines::Breaker.new(
+          key,
+          failure_threshold: @maxretry,
+          timeout: @bantime,
+          storage: ThrottleMachines.storage
+        )
+
+        # Record failure by trying to call through the breaker
+        # and letting it fail
+        begin
+          breaker.call { raise 'Fail2Ban failure' }
+        rescue StandardError
+          # Expected - this records the failure
+        end
+      end
+
+      private
+
+      def discriminator_for(request)
+        @block.call(request)
+      end
+    end
+  end
+end

data/lib/throttle_machines/rack_middleware/request.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module ThrottleMachines
+  class RackMiddleware
+    # Rack 3 Request wrapper
+    class Request < ::Rack::Request
+      def user_agent
+        @env['HTTP_USER_AGENT']
+      end
+    end
+  end
+end

data/lib/throttle_machines/rack_middleware/safelist.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module ThrottleMachines
+  class RackMiddleware
+    class Safelist
+      attr_reader :name, :block
+
+      def initialize(name, &block)
+        @name = name
+        @block = block
+      end
+
+      def matched_by?(request)
+        return false unless @block
+
+        if @block.call(request)
+          request.env['rack.attack.matched'] = @name
+          request.env['rack.attack.match_type'] = :safelist
+          ThrottleMachines::RackMiddleware.instrument(request)
+          true
+        else
+          false
+        end
+      end
+    end
+  end
+end

data/lib/throttle_machines/rack_middleware/throttle.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+module ThrottleMachines
+  class RackMiddleware
+    class Throttle
+      attr_reader :name, :limit, :period, :block, :algorithm
+
+      def initialize(name, options, &block)
+        @name = name
+        @block = block
+
+        raise ArgumentError, 'Must pass :limit option' unless options[:limit]
+        raise ArgumentError, 'Must pass :period option' unless options[:period]
+
+        @limit = options[:limit]
+        @period = options[:period].to_i
+        @algorithm = options[:algorithm] || :fixed_window
+      end
+
+      def matched_by?(request)
+        discriminator = discriminator_for(request)
+
+        return false unless discriminator
+
+        key = "#{@name}:#{discriminator}"
+        current_limit = limit_for(request)
+        current_period = period_for(request)
+
+        # Use ThrottleMachines limiter
+        limiter = ThrottleMachines.limiter(
+          key,
+          limit: current_limit,
+          period: current_period,
+          algorithm: @algorithm
+        )
+
+        # Try to consume the request atomically
+        throttled = false
+
+        begin
+          # This will either succeed and consume a request, or raise ThrottledError
+          limiter.throttle!
+
+          # If we get here, request was allowed
+        rescue ThrottledError
+          # Request was throttled
+          throttled = true
+        end
+
+        # Get current state for instrumentation
+        data = {
+          discriminator: discriminator,
+          count: current_limit - limiter.remaining,
+          period: current_period,
+          limit: current_limit,
+          retry_after: limiter.retry_after
+        }
+
+        annotate_request_with_throttle_data(request, data)
+
+        if throttled
+          annotate_request_with_matched_data(request, data)
+          ThrottleMachines::RackMiddleware.instrument(request)
+        end
+
+        throttled
+      end
+
+      private
+
+      def discriminator_for(request)
+        @block.call(request)
+      end
+
+      def limit_for(request)
+        @limit.respond_to?(:call) ? @limit.call(request) : @limit
+      end
+
+      def period_for(request)
+        @period.respond_to?(:call) ? @period.call(request) : @period
+      end
+
+      def annotate_request_with_throttle_data(request, data)
+        (request.env['rack.attack.throttle_data'] ||= {})[@name] = data
+      end
+
+      def annotate_request_with_matched_data(request, data)
+        request.env['rack.attack.matched'] = @name
+        request.env['rack.attack.match_discriminator'] = data[:discriminator]
+        request.env['rack.attack.match_type'] = :throttle
+        request.env['rack.attack.match_data'] = data
+      end
+    end
+  end
+end

data/lib/throttle_machines/rack_middleware/track.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module ThrottleMachines
+  class RackMiddleware
+    class Track
+      attr_reader :name, :block, :limit, :period
+
+      def initialize(name, options = {}, &block)
+        @name = name
+        @block = block
+        @limit = options[:limit]
+        @period = options[:period]
+      end
+
+      def matched_by?(request)
+        discriminator = @block.call(request)
+        return false unless discriminator
+
+        # Track is just instrumentation without blocking
+        data = {
+          discriminator: discriminator
+        }
+
+        # If limit and period are provided, track the count
+        if @limit && @period
+          key = "track:#{@name}:#{discriminator}"
+          limiter = ThrottleMachines.limiter(
+            key,
+            limit: @limit,
+            period: @period,
+            algorithm: :fixed_window
+          )
+
+          # Just check, don't consume
+          data[:count] = @limit - limiter.remaining
+          data[:limit] = @limit
+          data[:period] = @period
+        end
+
+        request.env['rack.attack.matched'] = @name
+        request.env['rack.attack.match_type'] = :track
+        request.env['rack.attack.match_discriminator'] = discriminator
+        request.env['rack.attack.match_data'] = data
+
+        ThrottleMachines::RackMiddleware.instrument(request)
+
+        false # Track never blocks
+      end
+    end
+  end
+end

data/lib/throttle_machines/rack_middleware.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+require 'rack'
+require 'forwardable'
+
+module ThrottleMachines
+  # Advanced Rack middleware for rate limiting and request filtering
+  class RackMiddleware
+    class << self
+      extend Forwardable
+
+      attr_accessor :enabled, :notifier
+      attr_reader :configuration
+
+      def_delegators :@configuration,
+                     :throttle,
+                     :track,
+                     :safelist,
+                     :blocklist,
+                     :blocklist_ip,
+                     :safelist_ip,
+                     :fail2ban,
+                     :allow2ban,
+                     :throttled_responder,
+                     :throttled_responder=,
+                     :blocklisted_responder,
+                     :blocklisted_responder=,
+                     :throttles,
+                     :tracks,
+                     :safelists,
+                     :blocklists
+
+      def configure(&)
+        @configuration ||= Configuration.new
+        @configuration.instance_eval(&) if block
+      end
+
+      # rubocop:disable Rails/Delegate -- Ruby 3.4 compatibility issue with delegate
+      def reset!
+        ThrottleMachines.reset!
+      end
+      # rubocop:enable Rails/Delegate
+
+      def clear!
+        @configuration = Configuration.new
+      end
+
+      # Instrument for ActiveSupport::Notifications compatibility
+      def instrument(request)
+        return unless notifier
+
+        event_type = request.env['rack.attack.match_type']
+        notifier.instrument("#{event_type}.throttle_machines", request: request)
+      end
+    end
+
+    # Set defaults
+    @enabled = true
+    @notifier = ActiveSupport::Notifications
+    @configuration = Configuration.new
+
+    def initialize(app)
+      @app = app
+    end
+
+    def call(env)
+      return @app.call(env) if !self.class.enabled || env['rack.attack.called']
+
+      env['rack.attack.called'] = true
+      request = Request.new(env)
+
+      # Always use the current class-level configuration
+      configuration = self.class.configuration
+
+      if configuration.safelisted?(request)
+        @app.call(env)
+      elsif configuration.blocklisted?(request)
+        configuration.blocklisted_responder.call(request)
+      elsif configuration.throttled?(request)
+        configuration.throttled_responder.call(request)
+      else
+        configuration.tracked?(request)
+        @app.call(env)
+      end
+    end
+  end
+end
+
+# No alias - users should use the explicit name to avoid confusion

data/lib/throttle_machines/storage/base.rb

@@ -0,0 +1,93 @@
+# frozen_string_literal: true
+
+module ThrottleMachines
+  module Storage
+    class Base
+      def initialize(options = {})
+        @options = options
+      end
+
+      # Rate limiting operations
+      def increment_counter(key, window, amount = 1)
+        raise NotImplementedError
+      end
+
+      def get_counter(key, window)
+        raise NotImplementedError
+      end
+
+      def get_counter_ttl(key, window)
+        raise NotImplementedError
+      end
+
+      def reset_counter(key, window)
+        raise NotImplementedError
+      end
+
+      # GCRA operations (atomic)
+      def check_gcra_limit(key, emission_interval, delay_tolerance, ttl)
+        raise NotImplementedError
+      end
+
+      def peek_gcra_limit(key, emission_interval, delay_tolerance)
+        raise NotImplementedError
+      end
+
+      # Token bucket operations (atomic)
+      def check_token_bucket(key, capacity, refill_rate, ttl)
+        raise NotImplementedError
+      end
+
+      def peek_token_bucket(key, capacity, refill_rate)
+        raise NotImplementedError
+      end
+
+      # Circuit breaker operations
+      def get_breaker_state(key)
+        raise NotImplementedError
+      end
+
+      def record_breaker_success(key, timeout, half_open_requests = 1)
+        raise NotImplementedError
+      end
+
+      def record_breaker_failure(key, threshold, timeout)
+        raise NotImplementedError
+      end
+
+      def trip_breaker(key, timeout)
+        raise NotImplementedError
+      end
+
+      def reset_breaker(key)
+        raise NotImplementedError
+      end
+
+      # Utility operations
+      def clear(pattern = nil)
+        raise NotImplementedError
+      end
+
+      def healthy?
+        raise NotImplementedError
+      end
+
+      def with_timeout(timeout, &)
+        Timeout.timeout(timeout, &)
+      rescue Timeout::Error
+        nil
+      end
+
+      protected
+
+      def current_time
+        # Use monotonic time for consistency with BreakerMachines
+        ThrottleMachines.monotonic_time
+      end
+
+      def monotonic_time
+        ThrottleMachines.monotonic_time
+      end
+    end
+  end
+end