throttle_machines 0.0.0 → 0.1.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- checksums.yaml +4 -4
- data/MIT-LICENSE +20 -0
- data/README.md +187 -13
- data/Rakefile +12 -0
- data/lib/throttle_machines/async_limiter.rb +134 -0
- data/lib/throttle_machines/clock.rb +41 -0
- data/lib/throttle_machines/control.rb +95 -0
- data/lib/throttle_machines/controller_helpers.rb +79 -0
- data/lib/throttle_machines/dependency_error.rb +6 -0
- data/lib/throttle_machines/engine.rb +25 -0
- data/lib/throttle_machines/hedged_request.rb +137 -0
- data/lib/throttle_machines/instrumentation.rb +162 -0
- data/lib/throttle_machines/limiter.rb +167 -0
- data/lib/throttle_machines/middleware.rb +90 -0
- data/lib/throttle_machines/rack_middleware/allow2_ban.rb +62 -0
- data/lib/throttle_machines/rack_middleware/blocklist.rb +27 -0
- data/lib/throttle_machines/rack_middleware/configuration.rb +103 -0
- data/lib/throttle_machines/rack_middleware/fail2_ban.rb +87 -0
- data/lib/throttle_machines/rack_middleware/request.rb +12 -0
- data/lib/throttle_machines/rack_middleware/safelist.rb +27 -0
- data/lib/throttle_machines/rack_middleware/throttle.rb +95 -0
- data/lib/throttle_machines/rack_middleware/track.rb +51 -0
- data/lib/throttle_machines/rack_middleware.rb +87 -0
- data/lib/throttle_machines/storage/base.rb +93 -0
- data/lib/throttle_machines/storage/memory.rb +374 -0
- data/lib/throttle_machines/storage/null.rb +90 -0
- data/lib/throttle_machines/storage/redis.rb +451 -0
- data/lib/throttle_machines/throttled_error.rb +14 -0
- data/lib/throttle_machines/version.rb +5 -0
- data/lib/throttle_machines.rb +134 -5
- metadata +105 -9
- data/LICENSE.txt +0 -21
@@ -0,0 +1,137 @@
|
|
1
|
+
# frozen_string_literal: true
|
2
|
+
|
3
|
+
require 'concurrent-ruby'
|
4
|
+
|
5
|
+
module ThrottleMachines
|
6
|
+
# Hedged Request - Multi-path Navigation System
|
7
|
+
#
|
8
|
+
# Like sending scout ships on multiple routes to find the fastest path.
|
9
|
+
# The first ship to reach the destination wins, others are recalled.
|
10
|
+
#
|
11
|
+
# Reduces latency by racing multiple backends/attempts with staggered delays.
|
12
|
+
#
|
13
|
+
# Example:
|
14
|
+
# hedged = ThrottleMachines::HedgedRequest.new(
|
15
|
+
# delay: 0.05, # 50ms between attempts
|
16
|
+
# max_attempts: 3
|
17
|
+
# )
|
18
|
+
#
|
19
|
+
# result = hedged.run do |attempt|
|
20
|
+
# case attempt
|
21
|
+
# when 0 then primary_backend.get(key)
|
22
|
+
# when 1 then secondary_backend.get(key)
|
23
|
+
# when 2 then tertiary_backend.get(key)
|
24
|
+
# end
|
25
|
+
# end
|
26
|
+
class HedgedRequest
|
27
|
+
attr_reader :delay, :max_attempts, :timeout
|
28
|
+
|
29
|
+
def initialize(delay: 0.05, max_attempts: 2, timeout: nil)
|
30
|
+
@delay = delay
|
31
|
+
@max_attempts = max_attempts
|
32
|
+
@timeout = timeout
|
33
|
+
@executor = Concurrent::ThreadPoolExecutor.new(
|
34
|
+
min_threads: 1,
|
35
|
+
max_threads: max_attempts,
|
36
|
+
max_queue: max_attempts,
|
37
|
+
fallback_policy: :caller_runs
|
38
|
+
)
|
39
|
+
end
|
40
|
+
|
41
|
+
# Run hedged request with automatic cancellation of slower attempts
|
42
|
+
def run(&block)
|
43
|
+
raise ArgumentError, 'Block required' unless block
|
44
|
+
|
45
|
+
# Generate a unique request ID for tracking
|
46
|
+
request_id = "#{object_id}-#{Time.now.to_f}"
|
47
|
+
|
48
|
+
# Instrument the start of the hedged request
|
49
|
+
Instrumentation.hedged_request_started(request_id, attempts: @max_attempts)
|
50
|
+
|
51
|
+
# Use Concurrent::Promises for better async handling
|
52
|
+
futures = []
|
53
|
+
first_result = Concurrent::Promises.resolvable_future
|
54
|
+
start_time = Time.now.to_f
|
55
|
+
|
56
|
+
@max_attempts.times do |attempt|
|
57
|
+
# Schedule with delay
|
58
|
+
future = if attempt.zero?
|
59
|
+
Concurrent::Promises.future { yield(attempt) }
|
60
|
+
else
|
61
|
+
Concurrent::Promises.schedule(@delay * attempt) { yield(attempt) }
|
62
|
+
end
|
63
|
+
|
64
|
+
# Race to resolve first_result
|
65
|
+
future.then do |result|
|
66
|
+
if !first_result.resolved? && first_result.fulfill(result)
|
67
|
+
# This attempt won the race
|
68
|
+
duration = Time.now.to_f - start_time
|
69
|
+
Instrumentation.hedged_request_winner(request_id, attempt: attempt, duration: duration)
|
70
|
+
end
|
71
|
+
result
|
72
|
+
end.rescue do |error|
|
73
|
+
# Only reject if this was the last attempt and nothing succeeded
|
74
|
+
first_result.reject(error) if attempt == @max_attempts - 1 && !first_result.resolved?
|
75
|
+
end
|
76
|
+
|
77
|
+
futures << future
|
78
|
+
end
|
79
|
+
|
80
|
+
# Wait with optional timeout
|
81
|
+
if @timeout
|
82
|
+
# Use any_resolved_future with timeout
|
83
|
+
timeout_future = Concurrent::Promises.schedule(@timeout) do
|
84
|
+
raise TimeoutError, "Hedged request timed out after #{@timeout}s"
|
85
|
+
end
|
86
|
+
|
87
|
+
Concurrent::Promises.any_resolved_future(first_result, timeout_future).value!
|
88
|
+
else
|
89
|
+
first_result.value!
|
90
|
+
end
|
91
|
+
ensure
|
92
|
+
# Cancel pending futures
|
93
|
+
futures.each { |f| f.cancel if f.pending? }
|
94
|
+
end
|
95
|
+
|
96
|
+
# Run async version
|
97
|
+
def run_async(&)
|
98
|
+
Concurrent::Promises.future { run(&) }
|
99
|
+
end
|
100
|
+
|
101
|
+
# Shutdown the executor
|
102
|
+
def shutdown
|
103
|
+
@executor.shutdown
|
104
|
+
@executor.wait_for_termination(5)
|
105
|
+
end
|
106
|
+
end
|
107
|
+
|
108
|
+
# Hedged request with circuit breaker integration
|
109
|
+
class HedgedBreaker
|
110
|
+
def initialize(breakers, delay: 0.05)
|
111
|
+
@breakers = Array(breakers)
|
112
|
+
@hedged = HedgedRequest.new(
|
113
|
+
delay: delay,
|
114
|
+
max_attempts: @breakers.size
|
115
|
+
)
|
116
|
+
end
|
117
|
+
|
118
|
+
def run(&)
|
119
|
+
@hedged.run do |attempt|
|
120
|
+
breaker = @breakers[attempt]
|
121
|
+
next if breaker.nil?
|
122
|
+
|
123
|
+
breaker.call(&)
|
124
|
+
end
|
125
|
+
end
|
126
|
+
end
|
127
|
+
|
128
|
+
# Convenience method
|
129
|
+
def self.hedged_request(**, &)
|
130
|
+
hedged = HedgedRequest.new(**)
|
131
|
+
begin
|
132
|
+
hedged.run(&)
|
133
|
+
ensure
|
134
|
+
hedged.shutdown
|
135
|
+
end
|
136
|
+
end
|
137
|
+
end
|
@@ -0,0 +1,162 @@
|
|
1
|
+
# frozen_string_literal: true
|
2
|
+
|
3
|
+
module ThrottleMachines
|
4
|
+
# Instrumentation module for emitting events via ActiveSupport::Notifications
|
5
|
+
module Instrumentation
|
6
|
+
class << self
|
7
|
+
attr_writer :enabled, :backend
|
8
|
+
|
9
|
+
def enabled
|
10
|
+
@enabled = true if @enabled.nil?
|
11
|
+
@enabled
|
12
|
+
end
|
13
|
+
|
14
|
+
def backend
|
15
|
+
@backend ||= if defined?(ActiveSupport::Notifications)
|
16
|
+
ActiveSupport::Notifications
|
17
|
+
else
|
18
|
+
NullBackend.new
|
19
|
+
end
|
20
|
+
end
|
21
|
+
|
22
|
+
def instrument(event_name, payload = {}, &)
|
23
|
+
if !enabled || backend.nil?
|
24
|
+
return yield if block_given?
|
25
|
+
|
26
|
+
return
|
27
|
+
end
|
28
|
+
|
29
|
+
full_event_name = "#{event_name}.throttle_machines"
|
30
|
+
backend.instrument(full_event_name, payload, &)
|
31
|
+
end
|
32
|
+
|
33
|
+
# Convenience methods for common events
|
34
|
+
|
35
|
+
# Rate limiter events
|
36
|
+
def rate_limit_checked(limiter, allowed:, remaining: nil)
|
37
|
+
payload = {
|
38
|
+
key: limiter.key,
|
39
|
+
limit: limiter.limit,
|
40
|
+
period: limiter.period,
|
41
|
+
algorithm: limiter.algorithm,
|
42
|
+
allowed: allowed,
|
43
|
+
remaining: remaining
|
44
|
+
}
|
45
|
+
instrument('rate_limit.checked', payload)
|
46
|
+
end
|
47
|
+
|
48
|
+
def rate_limit_allowed(limiter, remaining: nil)
|
49
|
+
payload = {
|
50
|
+
key: limiter.key,
|
51
|
+
limit: limiter.limit,
|
52
|
+
period: limiter.period,
|
53
|
+
algorithm: limiter.algorithm,
|
54
|
+
remaining: remaining
|
55
|
+
}
|
56
|
+
instrument('rate_limit.allowed', payload)
|
57
|
+
end
|
58
|
+
|
59
|
+
def rate_limit_throttled(limiter, retry_after: nil)
|
60
|
+
payload = {
|
61
|
+
key: limiter.key,
|
62
|
+
limit: limiter.limit,
|
63
|
+
period: limiter.period,
|
64
|
+
algorithm: limiter.algorithm,
|
65
|
+
retry_after: retry_after
|
66
|
+
}
|
67
|
+
instrument('rate_limit.throttled', payload)
|
68
|
+
end
|
69
|
+
|
70
|
+
# Circuit breaker events
|
71
|
+
def circuit_opened(breaker, failure_count:)
|
72
|
+
payload = {
|
73
|
+
key: breaker.key,
|
74
|
+
failure_threshold: breaker.failure_threshold,
|
75
|
+
timeout: breaker.timeout,
|
76
|
+
failure_count: failure_count
|
77
|
+
}
|
78
|
+
instrument('circuit_breaker.opened', payload)
|
79
|
+
end
|
80
|
+
|
81
|
+
def circuit_closed(breaker)
|
82
|
+
payload = {
|
83
|
+
key: breaker.key,
|
84
|
+
failure_threshold: breaker.failure_threshold,
|
85
|
+
timeout: breaker.timeout
|
86
|
+
}
|
87
|
+
instrument('circuit_breaker.closed', payload)
|
88
|
+
end
|
89
|
+
|
90
|
+
def circuit_half_opened(breaker)
|
91
|
+
payload = {
|
92
|
+
key: breaker.key,
|
93
|
+
failure_threshold: breaker.failure_threshold,
|
94
|
+
timeout: breaker.timeout,
|
95
|
+
half_open_requests: breaker.half_open_requests
|
96
|
+
}
|
97
|
+
instrument('circuit_breaker.half_opened', payload)
|
98
|
+
end
|
99
|
+
|
100
|
+
def circuit_success(breaker)
|
101
|
+
payload = {
|
102
|
+
key: breaker.key,
|
103
|
+
state: breaker.state
|
104
|
+
}
|
105
|
+
instrument('circuit_breaker.success', payload)
|
106
|
+
end
|
107
|
+
|
108
|
+
def circuit_failure(breaker, error: nil)
|
109
|
+
payload = {
|
110
|
+
key: breaker.key,
|
111
|
+
state: breaker.state,
|
112
|
+
error_class: error&.class&.name,
|
113
|
+
error_message: error&.message
|
114
|
+
}
|
115
|
+
instrument('circuit_breaker.failure', payload)
|
116
|
+
end
|
117
|
+
|
118
|
+
def circuit_rejected(breaker)
|
119
|
+
payload = {
|
120
|
+
key: breaker.key,
|
121
|
+
failure_threshold: breaker.failure_threshold,
|
122
|
+
timeout: breaker.timeout
|
123
|
+
}
|
124
|
+
instrument('circuit_breaker.rejected', payload)
|
125
|
+
end
|
126
|
+
|
127
|
+
# Cascade events
|
128
|
+
def cascade_triggered(primary_key, cascaded_key)
|
129
|
+
payload = {
|
130
|
+
primary_key: primary_key,
|
131
|
+
cascaded_key: cascaded_key
|
132
|
+
}
|
133
|
+
instrument('cascade.triggered', payload)
|
134
|
+
end
|
135
|
+
|
136
|
+
# Hedged request events
|
137
|
+
def hedged_request_started(request_id, attempts:)
|
138
|
+
payload = {
|
139
|
+
request_id: request_id,
|
140
|
+
max_attempts: attempts
|
141
|
+
}
|
142
|
+
instrument('hedged_request.started', payload)
|
143
|
+
end
|
144
|
+
|
145
|
+
def hedged_request_winner(request_id, attempt:, duration:)
|
146
|
+
payload = {
|
147
|
+
request_id: request_id,
|
148
|
+
winning_attempt: attempt,
|
149
|
+
duration: duration
|
150
|
+
}
|
151
|
+
instrument('hedged_request.winner', payload)
|
152
|
+
end
|
153
|
+
end
|
154
|
+
|
155
|
+
# Null backend for when ActiveSupport::Notifications is not available
|
156
|
+
class NullBackend
|
157
|
+
def instrument(_name, _payload = {})
|
158
|
+
yield if block_given?
|
159
|
+
end
|
160
|
+
end
|
161
|
+
end
|
162
|
+
end
|
@@ -0,0 +1,167 @@
|
|
1
|
+
# frozen_string_literal: true
|
2
|
+
|
3
|
+
module ThrottleMachines
|
4
|
+
class Limiter
|
5
|
+
attr_reader :key, :limit, :period, :algorithm, :storage
|
6
|
+
|
7
|
+
def initialize(key, limit:, period:, algorithm: :fixed_window, storage: nil)
|
8
|
+
@key = key
|
9
|
+
@limit = limit
|
10
|
+
@period = period
|
11
|
+
@algorithm = algorithm
|
12
|
+
@storage = storage || ThrottleMachines.storage
|
13
|
+
end
|
14
|
+
|
15
|
+
def allow?
|
16
|
+
allowed = case @algorithm
|
17
|
+
when :fixed_window
|
18
|
+
# Don't increment here, just check
|
19
|
+
count = @storage.get_counter(@key, @period)
|
20
|
+
count < @limit
|
21
|
+
when :gcra
|
22
|
+
result = @storage.peek_gcra_limit(
|
23
|
+
@key,
|
24
|
+
@period.to_f / @limit, # emission_interval
|
25
|
+
0 # delay_tolerance (no burst)
|
26
|
+
)
|
27
|
+
result[:allowed]
|
28
|
+
when :token_bucket
|
29
|
+
result = @storage.peek_token_bucket(
|
30
|
+
@key,
|
31
|
+
@limit, # capacity
|
32
|
+
@limit.to_f / @period # refill_rate
|
33
|
+
)
|
34
|
+
result[:allowed]
|
35
|
+
else
|
36
|
+
raise ArgumentError, "Unknown algorithm: #{@algorithm}"
|
37
|
+
end
|
38
|
+
|
39
|
+
# Instrument the check
|
40
|
+
Instrumentation.rate_limit_checked(self, allowed: allowed, remaining: nil)
|
41
|
+
|
42
|
+
allowed
|
43
|
+
end
|
44
|
+
|
45
|
+
def throttle!
|
46
|
+
case @algorithm
|
47
|
+
when :fixed_window
|
48
|
+
# Increment and check atomically
|
49
|
+
count = @storage.increment_counter(@key, @period)
|
50
|
+
if count > @limit
|
51
|
+
Instrumentation.rate_limit_throttled(self, retry_after: retry_after)
|
52
|
+
raise ThrottledError, self
|
53
|
+
end
|
54
|
+
when :gcra
|
55
|
+
result = @storage.check_gcra_limit(
|
56
|
+
@key,
|
57
|
+
@period.to_f / @limit, # emission_interval
|
58
|
+
0, # delay_tolerance (no burst)
|
59
|
+
(@period * 2).to_i # ttl
|
60
|
+
)
|
61
|
+
unless result[:allowed]
|
62
|
+
Instrumentation.rate_limit_throttled(self, retry_after: retry_after)
|
63
|
+
raise ThrottledError, self
|
64
|
+
end
|
65
|
+
when :token_bucket
|
66
|
+
result = @storage.check_token_bucket(
|
67
|
+
@key,
|
68
|
+
@limit, # capacity
|
69
|
+
@limit.to_f / @period, # refill_rate
|
70
|
+
(@period * 2).to_i # ttl
|
71
|
+
)
|
72
|
+
unless result[:allowed]
|
73
|
+
Instrumentation.rate_limit_throttled(self, retry_after: retry_after)
|
74
|
+
raise ThrottledError, self
|
75
|
+
end
|
76
|
+
else
|
77
|
+
raise ArgumentError, "Unknown algorithm: #{@algorithm}"
|
78
|
+
end
|
79
|
+
|
80
|
+
# If we get here, the request was allowed
|
81
|
+
# Calculate remaining after the operation
|
82
|
+
remaining_count = begin
|
83
|
+
remaining
|
84
|
+
rescue StandardError
|
85
|
+
nil
|
86
|
+
end
|
87
|
+
Instrumentation.rate_limit_allowed(self, remaining: remaining_count)
|
88
|
+
|
89
|
+
yield if block_given?
|
90
|
+
end
|
91
|
+
|
92
|
+
def reset!
|
93
|
+
case @algorithm
|
94
|
+
when :fixed_window
|
95
|
+
@storage.reset_counter(@key, @period)
|
96
|
+
else
|
97
|
+
@storage.clear("#{@key}*")
|
98
|
+
end
|
99
|
+
end
|
100
|
+
|
101
|
+
def remaining
|
102
|
+
case @algorithm
|
103
|
+
when :fixed_window
|
104
|
+
count = @storage.get_counter(@key, @period)
|
105
|
+
[@limit - count, 0].max
|
106
|
+
when :gcra
|
107
|
+
# GCRA doesn't have a simple "remaining" count
|
108
|
+
# Just return 1 or 0 based on current state
|
109
|
+
result = @storage.peek_gcra_limit(
|
110
|
+
@key,
|
111
|
+
@period.to_f / @limit,
|
112
|
+
0
|
113
|
+
)
|
114
|
+
result[:allowed] ? 1 : 0
|
115
|
+
when :token_bucket
|
116
|
+
result = @storage.peek_token_bucket(
|
117
|
+
@key,
|
118
|
+
@limit,
|
119
|
+
@limit.to_f / @period
|
120
|
+
)
|
121
|
+
result[:tokens_remaining].to_i
|
122
|
+
else
|
123
|
+
0
|
124
|
+
end
|
125
|
+
end
|
126
|
+
|
127
|
+
def retry_after
|
128
|
+
case @algorithm
|
129
|
+
when :fixed_window
|
130
|
+
count = @storage.get_counter(@key, @period)
|
131
|
+
if count >= @limit
|
132
|
+
# Return the actual time remaining in the current window
|
133
|
+
@storage.get_counter_ttl(@key, @period)
|
134
|
+
else
|
135
|
+
0
|
136
|
+
end
|
137
|
+
when :gcra
|
138
|
+
result = @storage.peek_gcra_limit(
|
139
|
+
@key,
|
140
|
+
@period.to_f / @limit,
|
141
|
+
0
|
142
|
+
)
|
143
|
+
result[:retry_after]
|
144
|
+
when :token_bucket
|
145
|
+
result = @storage.peek_token_bucket(
|
146
|
+
@key,
|
147
|
+
@limit,
|
148
|
+
@limit.to_f / @period
|
149
|
+
)
|
150
|
+
result[:retry_after]
|
151
|
+
else
|
152
|
+
0
|
153
|
+
end
|
154
|
+
end
|
155
|
+
|
156
|
+
def to_h
|
157
|
+
{
|
158
|
+
key: @key,
|
159
|
+
limit: @limit,
|
160
|
+
period: @period,
|
161
|
+
algorithm: @algorithm,
|
162
|
+
remaining: remaining,
|
163
|
+
retry_after: retry_after
|
164
|
+
}
|
165
|
+
end
|
166
|
+
end
|
167
|
+
end
|
@@ -0,0 +1,90 @@
|
|
1
|
+
# frozen_string_literal: true
|
2
|
+
|
3
|
+
module ThrottleMachines
|
4
|
+
class Middleware
|
5
|
+
def initialize(app, store: nil, &config_block)
|
6
|
+
@app = app
|
7
|
+
@store = store || ThrottleMachines.configuration.store
|
8
|
+
@rules = []
|
9
|
+
|
10
|
+
instance_eval(&config_block) if config_block
|
11
|
+
end
|
12
|
+
|
13
|
+
def call(env)
|
14
|
+
request = Rack::Request.new(env)
|
15
|
+
|
16
|
+
@rules.each do |rule|
|
17
|
+
next unless rule[:matcher].call(request)
|
18
|
+
|
19
|
+
key = rule[:key_generator].call(request)
|
20
|
+
limiter = ThrottleMachines.limiter(
|
21
|
+
key,
|
22
|
+
limit: rule[:limit],
|
23
|
+
period: rule[:period],
|
24
|
+
algorithm: rule[:algorithm]
|
25
|
+
)
|
26
|
+
|
27
|
+
return rate_limit_response(limiter) unless limiter.allow?
|
28
|
+
end
|
29
|
+
|
30
|
+
@app.call(env)
|
31
|
+
rescue ThrottledError => e
|
32
|
+
rate_limit_response(e.limiter)
|
33
|
+
end
|
34
|
+
|
35
|
+
def throttle(path, limit:, period:, by: :ip, algorithm: :gcra)
|
36
|
+
@rules << {
|
37
|
+
matcher: build_matcher(path),
|
38
|
+
key_generator: build_key_generator(by),
|
39
|
+
limit: limit,
|
40
|
+
period: period,
|
41
|
+
algorithm: algorithm
|
42
|
+
}
|
43
|
+
end
|
44
|
+
|
45
|
+
private
|
46
|
+
|
47
|
+
def build_matcher(path)
|
48
|
+
case path
|
49
|
+
when String
|
50
|
+
->(request) { request.path == path }
|
51
|
+
when Regexp
|
52
|
+
->(request) { request.path =~ path }
|
53
|
+
when Proc
|
54
|
+
path
|
55
|
+
else
|
56
|
+
->(_request) { true }
|
57
|
+
end
|
58
|
+
end
|
59
|
+
|
60
|
+
def build_key_generator(by)
|
61
|
+
case by
|
62
|
+
when :ip
|
63
|
+
->(request) { "ip:#{request.ip}" }
|
64
|
+
when Symbol
|
65
|
+
->(request) { "#{by}:#{request.env['rack.session']&.dig(by)}" }
|
66
|
+
when Proc
|
67
|
+
by
|
68
|
+
else
|
69
|
+
->(_request) { by.to_s }
|
70
|
+
end
|
71
|
+
end
|
72
|
+
|
73
|
+
def rate_limit_response(limiter)
|
74
|
+
headers = {
|
75
|
+
'Content-Type' => 'application/json',
|
76
|
+
'X-RateLimit-Limit' => limiter.limit.to_s,
|
77
|
+
'X-RateLimit-Remaining' => limiter.remaining.to_s,
|
78
|
+
'X-RateLimit-Reset' => (Time.now.to_i + limiter.retry_after).to_s,
|
79
|
+
'Retry-After' => limiter.retry_after.ceil.to_s
|
80
|
+
}
|
81
|
+
|
82
|
+
body = JSON.generate({
|
83
|
+
error: 'Rate limit exceeded',
|
84
|
+
retry_after: limiter.retry_after
|
85
|
+
})
|
86
|
+
|
87
|
+
[429, headers, [body]]
|
88
|
+
end
|
89
|
+
end
|
90
|
+
end
|
@@ -0,0 +1,62 @@
|
|
1
|
+
# frozen_string_literal: true
|
2
|
+
|
3
|
+
module ThrottleMachines
|
4
|
+
class RackMiddleware
|
5
|
+
class Allow2Ban
|
6
|
+
attr_reader :name, :maxretry, :findtime, :bantime, :block
|
7
|
+
|
8
|
+
def initialize(name, options, &block)
|
9
|
+
@name = name
|
10
|
+
@block = block
|
11
|
+
|
12
|
+
@maxretry = options[:maxretry] || 5
|
13
|
+
@findtime = options[:findtime] || 60
|
14
|
+
@bantime = options[:bantime] || 300
|
15
|
+
end
|
16
|
+
|
17
|
+
def matched_by?(request)
|
18
|
+
discriminator = discriminator_for(request)
|
19
|
+
return false unless discriminator
|
20
|
+
|
21
|
+
# Allow2Ban resets fail2ban counters on successful requests
|
22
|
+
# We'll track successful requests and reset the breaker when threshold is met
|
23
|
+
|
24
|
+
success_key = "allow2ban:#{@name}:#{discriminator}"
|
25
|
+
fail_key = "fail2ban:#{@name}:#{discriminator}"
|
26
|
+
|
27
|
+
# Count successful requests
|
28
|
+
success_limiter = ThrottleMachines.limiter(
|
29
|
+
success_key,
|
30
|
+
limit: @maxretry,
|
31
|
+
period: @findtime,
|
32
|
+
algorithm: :fixed_window
|
33
|
+
)
|
34
|
+
|
35
|
+
# Check if we've had enough successful requests
|
36
|
+
if success_limiter.remaining.zero?
|
37
|
+
# Reset the fail2ban breaker
|
38
|
+
storage = ThrottleMachines.storage
|
39
|
+
storage.reset_breaker(fail_key)
|
40
|
+
|
41
|
+
# Reset our own counter
|
42
|
+
storage.reset_counter(success_key, @findtime)
|
43
|
+
else
|
44
|
+
# Increment success counter
|
45
|
+
begin
|
46
|
+
success_limiter.throttle!
|
47
|
+
rescue ThrottledError
|
48
|
+
# We've hit the limit, which triggers the reset above
|
49
|
+
end
|
50
|
+
end
|
51
|
+
|
52
|
+
false # Allow2Ban never blocks directly
|
53
|
+
end
|
54
|
+
|
55
|
+
private
|
56
|
+
|
57
|
+
def discriminator_for(request)
|
58
|
+
@block.call(request)
|
59
|
+
end
|
60
|
+
end
|
61
|
+
end
|
62
|
+
end
|
@@ -0,0 +1,27 @@
|
|
1
|
+
# frozen_string_literal: true
|
2
|
+
|
3
|
+
module ThrottleMachines
|
4
|
+
class RackMiddleware
|
5
|
+
class Blocklist
|
6
|
+
attr_reader :name, :block
|
7
|
+
|
8
|
+
def initialize(name, &block)
|
9
|
+
@name = name
|
10
|
+
@block = block
|
11
|
+
end
|
12
|
+
|
13
|
+
def matched_by?(request)
|
14
|
+
return false unless @block
|
15
|
+
|
16
|
+
if @block.call(request)
|
17
|
+
request.env['rack.attack.matched'] = @name
|
18
|
+
request.env['rack.attack.match_type'] = :blocklist
|
19
|
+
ThrottleMachines::RackMiddleware.instrument(request)
|
20
|
+
true
|
21
|
+
else
|
22
|
+
false
|
23
|
+
end
|
24
|
+
end
|
25
|
+
end
|
26
|
+
end
|
27
|
+
end
|