diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb index a2342f5..740df4f 100644 --- a/lib/rack/attack.rb +++ b/lib/rack/attack.rb @@ -14,6 +14,7 @@ module Rack class Error < StandardError; end class MisconfiguredStoreError < Error; end class MissingStoreError < Error; end + class IncompatibleStoreError < Error; end autoload :Check, 'rack/attack/check' autoload :Throttle, 'rack/attack/throttle' @@ -53,6 +54,10 @@ module Rack @configuration.clear_configuration end + def reset! + cache.reset! + end + extend Forwardable def_delegators( :@configuration, diff --git a/lib/rack/attack/cache.rb b/lib/rack/attack/cache.rb index cfa2efa..0defa64 100644 --- a/lib/rack/attack/cache.rb +++ b/lib/rack/attack/cache.rb @@ -41,6 +41,17 @@ module Rack store.delete("#{prefix}:#{unprefixed_key}") end + def reset! + if store.respond_to?(:delete_matched) + store.delete_matched("#{prefix}*") + else + raise( + Rack::Attack::IncompatibleStoreError, + "Configured store #{store.class.name} doesn't respond to #delete_matched method" + ) + end + end + private def key_and_expiry(unprefixed_key, period) diff --git a/lib/rack/attack/store_proxy/redis_proxy.rb b/lib/rack/attack/store_proxy/redis_proxy.rb index d4e6f3a..e51b3e9 100644 --- a/lib/rack/attack/store_proxy/redis_proxy.rb +++ b/lib/rack/attack/store_proxy/redis_proxy.rb @@ -43,6 +43,19 @@ module Rack rescuing { del(key) } end + def delete_matched(matcher, _options = nil) + cursor = "0" + + rescuing do + # Fetch keys in batches using SCAN to avoid blocking the Redis server. + loop do + cursor, keys = scan(cursor, match: matcher, count: 1000) + del(*keys) unless keys.empty? + break if cursor == "0" + end + end + end + private def rescuing diff --git a/spec/rack_attack_spec.rb b/spec/rack_attack_spec.rb index 647a08d..d9bab4d 100644 --- a/spec/rack_attack_spec.rb +++ b/spec/rack_attack_spec.rb @@ -99,4 +99,26 @@ describe 'Rack::Attack' do end end end + + describe 'reset!' do + it 'raises an error when is not supported by cache store' do + Rack::Attack.cache.store = Class.new + assert_raises(Rack::Attack::IncompatibleStoreError) do + Rack::Attack.reset! + end + end + + if defined?(Redis) + it 'should delete rack attack keys' do + redis = Redis.new + redis.set('key', 'value') + redis.set("#{Rack::Attack.cache.prefix}::key", 'value') + Rack::Attack.cache.store = redis + Rack::Attack.reset! + + _(redis.get('key')).must_equal 'value' + _(redis.get("#{Rack::Attack.cache.prefix}::key")).must_be_nil + end + end + end end