[APPSEC-10303] Replace AppSec rate limiter with core rate limiter

Steepfile

@@ -70,7 +70,6 @@ target :datadog do
  ignore 'lib/datadog/core/metrics/metric.rb'
  ignore 'lib/datadog/core/metrics/options.rb'
  ignore 'lib/datadog/core/pin.rb'
- ignore 'lib/datadog/core/rate_limiter.rb'
  # steep fails in this file due to https://github.com/soutaro/steep/issues/1231
  ignore 'lib/datadog/core/remote/tie.rb'
  ignore 'lib/datadog/core/runtime/ext.rb'

lib/datadog/appsec/event.rb

@@ -52,7 +52,7 @@ def record(span, *events)
  # ensure rate limiter is called only when there are events to record
  return if events.empty? || span.nil?
 
- Datadog::AppSec::RateLimiter.limit(:traces) do
+ Datadog::AppSec::RateLimiter.thread_local.limit do
  record_via_span(span, *events)
  end
  end

lib/datadog/appsec/rate_limiter.rb

@@ -1,60 +1,45 @@
 # frozen_string_literal: true
 
+require_relative '../core/rate_limiter'
+
 module Datadog
  module AppSec
- # Simple per-thread rate limiter
- # Since AppSec marks sampling to keep on a security event, this limits the flood of egress traces involving AppSec
+ # Per-thread rate limiter based on token bucket rate limiter.
+ #
+ # Since AppSec marks sampling to keep on a security event, this limits
+ # the flood of egress traces involving AppSec
  class RateLimiter
- def initialize(rate)
- @rate = rate
- @timestamps = []
- end
-
- def limit
- now = Time.now.to_f
-
- loop do
- oldest = @timestamps.first
-
- break if oldest.nil? || now - oldest < 1
-
- @timestamps.shift
- end
-
- @timestamps << now
-
- if (count = @timestamps.count) <= @rate
- yield
- else
- Datadog.logger.debug { "Rate limit hit: #{count}/#{@rate} AppSec traces/second" }
- end
- end
+ THREAD_KEY = :datadog_security_appsec_rate_limiter
 
  class << self
- def limit(name, &block)
- rate_limiter(name).limit(&block)
+ def thread_local
+ rate_limiter = Thread.current.thread_variable_get(THREAD_KEY)
+ return rate_limiter unless rate_limiter.nil?
+
+ Thread.current.thread_variable_set(THREAD_KEY, new(trace_rate_limit))
  end
 
  # reset a rate limiter: used for testing
- def reset!(name)
- Thread.current[:datadog_security_trace_rate_limiter] = nil
+ def reset!
+ Thread.current.thread_variable_set(THREAD_KEY, nil)
  end
 
- protected
-
- def rate_limiter(name)
- case name
- when :traces
- Thread.current[:datadog_security_trace_rate_limiter] ||= RateLimiter.new(trace_rate_limit)
- else
- raise "unsupported rate limiter: #{name.inspect}"
- end
- end
+ private
 
  def trace_rate_limit
  Datadog.configuration.appsec.trace_rate_limit
  end
  end
+
+ def initialize(rate)
+ @rate_limiter = Core::TokenBucket.new(rate)
+ end
+
+ def limit
+ return yield if @rate_limiter.allow?
+
+ Datadog.logger.debug { "Rate limit hit: #{@rate_limiter.current_window_rate} AppSec traces/second" }
+ end
  end
  end
 end

sig/datadog/appsec/rate_limiter.rbs

@@ -1,24 +1,23 @@
 module Datadog
  module AppSec
  class RateLimiter
- type timestamp = ::Float
- type rate = ::Integer
+ @rate_limiter: Datadog::Core::TokenBucket
 
- @rate: ::Integer
- @timestamps: ::Array[timestamp]
+ THREAD_KEY: :datadog_security_appsec_rate_limiter
 
- def initialize: (rate rate) -> void
+ def self.thread_local: () -> RateLimiter
 
- # TODO: return type of limit is return type of block
- def limit: () { () -> untyped } -> untyped
+ def self.reset!: () -> void
 
- def self.limit: (::Symbol name) { () -> untyped } -> untyped
+ private
 
- def self.reset!: (::Symbol name) -> void
+ def self.trace_rate_limit: () -> ::Integer
 
- def self.rate_limiter: (::Symbol name) -> RateLimiter
+ public
 
- def self.trace_rate_limit: () -> rate
+ def initialize: (::Integer rate) -> void
+
+ def limit: () { () -> untyped } -> void
  end
  end
 end

sig/datadog/core/rate_limiter.rbs

@@ -1,34 +1,72 @@
 module Datadog
  module Core
  class RateLimiter
- def allow?: (untyped size) -> nil
- def effective_rate: () -> nil
+ def allow?: (?::Integer size) -> void
+
+ def effective_rate: () -> void
  end
+
  class TokenBucket < RateLimiter
- attr_reader rate: untyped
+ # This should be `::Numeric`, but it's used with `*` method on another
+ # `Numeric` which makes steep fail with an error `Ruby::UnresolvedOverloading`
+ # "Cannot find compatible overloading of method"
+ @rate: ::Float
+
+ @max_tokens: ::Numeric
+
+ @tokens: ::Numeric
+
+ @total_messages: ::Integer
+
+ @conforming_messages: ::Integer
+
+ # This should be `::Integer?` but steep can't see conditional branching
+ # which safe-guards us from errors like "undefined method for NilClass"
+ @prev_conforming_messages: ::Integer
+
+ # This should be `::Integer?` but steep can't see conditional branching
+ # which safe-guards us from errors like "undefined method for NilClass"
+ @prev_total_messages: ::Integer
+
+ # This should be `::Integer?`, but steep can't see conditional branching
+ # which safe-guards us from errors like "undefined method for NilClass"
+ @current_window: ::Numeric
+
+ @last_refill: ::Numeric
 
- attr_reader max_tokens: untyped
- def initialize: (untyped rate, ?untyped max_tokens) -> void
- def allow?: (untyped size) -> untyped
- def effective_rate: () -> (::Float | untyped)
- def current_window_rate: () -> (::Float | untyped)
- def available_tokens: () -> untyped
+ # This should be `::Numeric`, but has to follow `@rate` type definition
+ attr_reader rate: ::Float
+
+ attr_reader max_tokens: ::Numeric
+
+ def initialize: (::Numeric rate, ?::Numeric max_tokens) -> void
+
+ def allow?: (?::Integer size) -> bool
+
+ def effective_rate: () -> ::Float
+
+ def current_window_rate: () -> ::Float
+
+ def available_tokens: () -> ::Numeric
 
  private
 
- def refill_since_last_message: () -> untyped
+ def refill_since_last_message: () -> void
+
+ def refill_tokens: (::Numeric size) -> void
 
- def refill_tokens: (untyped size) -> untyped
+ def increment_total_count: () -> void
 
- def increment_total_count: () -> untyped
+ def increment_conforming_count: () -> void
 
- def increment_conforming_count: () -> untyped
+ def should_allow?: (?::Integer size) -> bool
 
- def should_allow?: (untyped size) -> (false | true)
- def update_rate_counts: (untyped allowed) -> untyped
+ def update_rate_counts: (bool allowed) -> void
  end
+
  class UnlimitedLimiter < RateLimiter
- def allow?: (untyped _) -> true
+ def allow?: (?::Integer _) -> true
+
  def effective_rate: () -> ::Float
  end
  end

spec/datadog/appsec/event_spec.rb

@@ -45,7 +45,7 @@
  describe '.record' do
  before do
  # prevent rate limiter to bias tests
- Datadog::AppSec::RateLimiter.reset!(:traces)
+ Datadog::AppSec::RateLimiter.reset!
  end
 
  let(:options) { {} }
@@ -309,7 +309,7 @@
  end
 
  it 'does not call the rate limiter' do
- expect(Datadog::AppSec::RateLimiter).to_not receive(:limit)
+ expect_any_instance_of(Datadog::AppSec::RateLimiter).to_not receive(:limit)
 
  expect(trace).to_not be nil
  end
@@ -325,14 +325,19 @@
  end
 
  it 'does not call the rate limiter' do
- expect(Datadog::AppSec::RateLimiter).to_not receive(:limit)
+ expect_any_instance_of(Datadog::AppSec::RateLimiter).to_not receive(:limit)
 
  described_class.record(nil, events)
  end
  end
 
  context 'with many traces' do
- let(:rate_limit) { 100 }
+ before do
+ allow(Datadog::Core::Utils::Time).to receive(:get_time).and_return(0)
+ allow(Datadog::AppSec::RateLimiter).to receive(:trace_rate_limit).and_return(rate_limit)
+ end
+
+ let(:rate_limit) { 50 }
  let(:trace_count) { rate_limit * 2 }
 
  let(:traces) do

spec/datadog/appsec/rate_limiter_spec.rb

@@ -0,0 +1,30 @@
+require 'datadog/appsec/spec_helper'
+require 'datadog/appsec/rate_limiter'
+
+RSpec.describe Datadog::AppSec::RateLimiter do
+ before { described_class.reset! }
+
+ describe '#limit' do
+ context 'in different threads' do
+ before { stub_const("#{described_class}::THREAD_KEY", :__spec_instance) }
+
+ it 'creates separate rate limiter per thread' do
+ thread_1 = Thread.new do
+ described_class.thread_local
+ end
+
+ thread_2 = Thread.new do
+ described_class.thread_local
+ end
+
+ thread_1.join
+ thread_2.join
+
+ rate_limiter_1 = thread_1.thread_variable_get(:__spec_instance)
+ rate_limiter_2 = thread_2.thread_variable_get(:__spec_instance)
+
+ expect(rate_limiter_1).not_to be(rate_limiter_2)
+ end
+ end
+ end
+end