Merge pull request #128 from stanhu/support-reset-throttles

Provide the ability to reset Fail2Ban count and ban flag

lib/rack/attack/cache.rb

@@ -15,10 +15,7 @@ module Rack
       end
 
       def count(unprefixed_key, period)
-        epoch_time = Time.now.to_i
-        # Add 1 to expires_in to avoid timing error: http://git.io/i1PHXA
-        expires_in = period - (epoch_time % period) + 1
-        key = "#{prefix}:#{(epoch_time/period).to_i}:#{unprefixed_key}"
+        key, expires_in = key_and_expiry(unprefixed_key, period)
         do_count(key, expires_in)
       end
 
@@ -30,7 +27,24 @@ module Rack
         store.write("#{prefix}:#{unprefixed_key}", value, :expires_in => expires_in)
       end
 
+      def reset_count(unprefixed_key, period)
+        key, _ = key_and_expiry(unprefixed_key, period)
+        store.delete(key)
+      end
+
+      def delete(unprefixed_key)
+        store.delete("#{prefix}:#{unprefixed_key}")
+      end
+
       private
+
+      def key_and_expiry(unprefixed_key, period)
+        epoch_time = Time.now.to_i
+        # Add 1 to expires_in to avoid timing error: http://git.io/i1PHXA
+        expires_in = (period - (epoch_time % period) + 1).to_i
+        ["#{prefix}:#{(epoch_time / period).to_i}:#{unprefixed_key}", expires_in]
+      end
+
       def do_count(key, expires_in)
         result = store.increment(key, 1, :expires_in => expires_in)
 

lib/rack/attack/fail2ban.rb

@@ -15,6 +15,17 @@ module Rack
           end
         end
 
+        def reset(discriminator, options)
+          findtime = options[:findtime] or raise ArgumentError, "Must pass findtime option"
+          cache.reset_count("#{key_prefix}:count:#{discriminator}", findtime)
+          # Clear ban flag just in case it's there
+          cache.delete("#{key_prefix}:ban:#{discriminator}")
+        end
+
+        def banned?(discriminator)
+          cache.read("#{key_prefix}:ban:#{discriminator}") ? true : false
+        end
+
         protected
         def key_prefix
           'fail2ban'
@@ -35,10 +46,6 @@ module Rack
           cache.write("#{key_prefix}:ban:#{discriminator}", 1, bantime)
         end
 
-        def banned?(discriminator)
-          cache.read("#{key_prefix}:ban:#{discriminator}")
-        end
-
         def cache
           Rack::Attack.cache
         end

spec/fail2ban_spec.rb

@@ -58,7 +58,27 @@ describe 'Rack::Attack.Fail2Ban' do
           key = "rack::attack:fail2ban:ban:1.2.3.4"
           @cache.store.read(key).must_equal 1
         end
+      end
 
+      describe 'reset after success' do
+        before do
+          get '/?test=OMGHAX', {}, 'REMOTE_ADDR' => '1.2.3.4'
+          Rack::Attack::Fail2Ban.reset('1.2.3.4', @f2b_options)
+          get '/', {}, 'REMOTE_ADDR' => '1.2.3.4'
+        end
+
+        it 'succeeds' do
+          last_response.status.must_equal 200
+        end
+
+        it 'resets fail count' do
+          key = "rack::attack:#{Time.now.to_i/@findtime}:fail2ban:count:1.2.3.4"
+          @cache.store.read(key).must_equal nil
+        end
+
+        it 'IP is not banned' do
+          Rack::Attack::Fail2Ban.banned?('1.2.3.4').must_equal false
+        end
       end
     end
   end

spec/integration/rack_attack_cache_spec.rb

@@ -89,6 +89,27 @@ describe Rack::Attack::Cache do
           @cache.read('cache-test-key').must_equal nil
         end
       end
+
+      describe "cache#delete" do
+        it "must delete the value" do
+          @cache.write("cache-test-key", "foobar", 1)
+          store.read(@key).must_equal "foobar"
+          @cache.delete('cache-test-key')
+          store.read(@key).must_be :nil?
+        end
+      end
+
+      describe "reset_count" do
+        it "must delete the value" do
+          period = 1.minute
+          unprefixed_key = 'cache-test-key'
+          @cache.count(unprefixed_key, period)
+          period_key, _ = @cache.send(:key_and_expiry, 'cache-test-key', period)
+          store.read(period_key).to_i.must_equal 1
+          @cache.reset_count(unprefixed_key, period)
+          store.read(period_key).must_equal nil
+        end
+      end
     end
 
   end