Replace AppSec rate limiter with core rate limiter

This allows us to be more precise in throttling outgoing AppSec traces
and removes additional custom logic of rate limiting.

lib/datadog/appsec/rate_limiter.rb

@@ -1,35 +1,14 @@
 # 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
-
  class << self
  def limit(name, &block)
  rate_limiter(name).limit(&block)
@@ -40,12 +19,12 @@ def reset!(name)
  Thread.current[:datadog_security_trace_rate_limiter] = nil
  end
 
- protected
+ private
 
  def rate_limiter(name)
  case name
  when :traces
- Thread.current[:datadog_security_trace_rate_limiter] ||= RateLimiter.new(trace_rate_limit)
+ Thread.current[:datadog_security_trace_rate_limiter] ||= new(trace_rate_limit)
  else
  raise "unsupported rate limiter: #{name.inspect}"
  end
@@ -55,6 +34,16 @@ 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

spec/datadog/appsec/event_spec.rb

@@ -332,7 +332,11 @@
  end
 
  context 'with many traces' do
- let(:rate_limit) { 100 }
+ before do
+ 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
@@ -349,9 +353,11 @@
  end
  end
 
+ # NOTE: Allows 10% excessive traces due to differece in execution time
+ # of the limited block of code on different machines
  it 'rate limits event recording' do
- expect(described_class).to receive(:record_via_span).exactly(rate_limit).times.and_call_original
-
+ expect(described_class).to receive(:record_via_span)
+ .at_least(rate_limit).at_most(rate_limit * 1.1).times.and_call_original
  expect(traces).to have_attributes(count: trace_count)
  end
  end

spec/datadog/appsec/rate_limiter_spec.rb

@@ -0,0 +1,38 @@
+require 'datadog/appsec/spec_helper'
+require 'datadog/appsec/rate_limiter'
+
+RSpec.describe Datadog::AppSec::RateLimiter do
+ before { described_class.reset!(:traces) }
+
+ describe '#limit' do
+ context 'in different threads' do
+ before do
+ allow(described_class).to receive(:trace_rate_limit).and_return(rate_limit)
+ end
+
+ let(:rate_limit) { 50 }
+ let(:spy_1) { spy('thread-1 workload') }
+ let(:spy_2) { spy('thread-2 workload') }
+
+ it 'creates separate rate limiter per thread' do
+ thread_1 = Thread.new do
+ Array.new(rate_limit * 2) do
+ described_class.limit(:traces) { spy_1.call }
+ end
+ end
+
+ thread_2 = Thread.new do
+ Array.new(rate_limit * 2) do
+ described_class.limit(:traces) { spy_2.call }
+ end
+ end
+
+ thread_1.join
+ thread_2.join
+
+ expect(spy_1).to have_received(:call).exactly(rate_limit)
+ expect(spy_2).to have_received(:call).exactly(rate_limit)
+ end
+ end
+ end
+end