From c3a077442a525816736773f9aa7fe7a0bc8c432b Mon Sep 17 00:00:00 2001
From: Steve Hodgkiss <steve@hodgkiss.me>
Date: Fri, 28 Mar 2014 14:53:51 +1100
Subject: [PATCH] Make Rack::Attack a class and Rack::Attack.new return an
 instance of that class, rather than Rack::Attack.

---
 lib/rack/attack.rb             | 59 +++++++++++++++++-----------------
 lib/rack/attack/allow2ban.rb   |  2 +-
 lib/rack/attack/blacklist.rb   |  2 +-
 lib/rack/attack/cache.rb       |  2 +-
 lib/rack/attack/check.rb       |  2 +-
 lib/rack/attack/fail2ban.rb    |  2 +-
 lib/rack/attack/store_proxy.rb |  2 +-
 lib/rack/attack/throttle.rb    |  2 +-
 lib/rack/attack/track.rb       |  2 +-
 lib/rack/attack/version.rb     |  2 +-
 lib/rack/attack/whitelist.rb   |  2 +-
 11 files changed, 39 insertions(+), 40 deletions(-)

diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb
index 91baafd..1ac1c6c 100644
--- a/lib/rack/attack.rb
+++ b/lib/rack/attack.rb
@@ -1,5 +1,6 @@
 require 'rack'
-module Rack::Attack
+
+class Rack::Attack
   autoload :Cache,     'rack/attack/cache'
   autoload :Check,     'rack/attack/check'
   autoload :Throttle,  'rack/attack/throttle'
@@ -35,35 +36,6 @@ module Rack::Attack
     def throttles;  @throttles  ||= {}; end
     def tracks;     @tracks     ||= {}; end
 
-    def new(app)
-      @app = app
-
-      # Set defaults
-      @notifier ||= ActiveSupport::Notifications if defined?(ActiveSupport::Notifications)
-      @blacklisted_response ||= lambda {|env| [403, {}, ["Forbidden\n"]] }
-      @throttled_response   ||= lambda {|env|
-        retry_after = env['rack.attack.match_data'][:period] rescue nil
-        [429, {'Retry-After' => retry_after.to_s}, ["Retry later\n"]]
-      }
-
-      self
-    end
-
-    def call(env)
-      req = Rack::Request.new(env)
-
-      if whitelisted?(req)
-        @app.call(env)
-      elsif blacklisted?(req)
-        blacklisted_response[env]
-      elsif throttled?(req)
-        throttled_response[env]
-      else
-        tracked?(req)
-        @app.call(env)
-      end
-    end
-
     def whitelisted?(req)
       whitelists.any? do |name, whitelist|
         whitelist[req]
@@ -101,4 +73,31 @@ module Rack::Attack
     end
 
   end
+
+  # Set defaults
+  @notifier ||= ActiveSupport::Notifications if defined?(ActiveSupport::Notifications)
+  @blacklisted_response ||= lambda {|env| [403, {}, ["Forbidden\n"]] }
+  @throttled_response   ||= lambda {|env|
+    retry_after = env['rack.attack.match_data'][:period] rescue nil
+    [429, {'Retry-After' => retry_after.to_s}, ["Retry later\n"]]
+  }
+
+  def initialize(app)
+    @app = app
+  end
+
+  def call(env)
+    req = Rack::Request.new(env)
+
+    if self.class.whitelisted?(req)
+      @app.call(env)
+    elsif self.class.blacklisted?(req)
+      self.class.blacklisted_response[env]
+    elsif self.class.throttled?(req)
+      self.class.throttled_response[env]
+    else
+      self.class.tracked?(req)
+      @app.call(env)
+    end
+  end
 end
diff --git a/lib/rack/attack/allow2ban.rb b/lib/rack/attack/allow2ban.rb
index f5772c8..bf8c598 100644
--- a/lib/rack/attack/allow2ban.rb
+++ b/lib/rack/attack/allow2ban.rb
@@ -1,5 +1,5 @@
 module Rack
-  module Attack
+  class Attack
     class Allow2Ban < Fail2Ban
       class << self
         protected
diff --git a/lib/rack/attack/blacklist.rb b/lib/rack/attack/blacklist.rb
index 4d2f666..c3a8341 100644
--- a/lib/rack/attack/blacklist.rb
+++ b/lib/rack/attack/blacklist.rb
@@ -1,5 +1,5 @@
 module Rack
-  module Attack
+  class Attack
     class Blacklist < Check
       def initialize(name, block)
         super
diff --git a/lib/rack/attack/cache.rb b/lib/rack/attack/cache.rb
index 9c54505..e65fa84 100644
--- a/lib/rack/attack/cache.rb
+++ b/lib/rack/attack/cache.rb
@@ -1,5 +1,5 @@
 module Rack
-  module Attack
+  class Attack
     class Cache
 
       attr_accessor :prefix
diff --git a/lib/rack/attack/check.rb b/lib/rack/attack/check.rb
index 79f3cd9..caea75a 100644
--- a/lib/rack/attack/check.rb
+++ b/lib/rack/attack/check.rb
@@ -1,5 +1,5 @@
 module Rack
-  module Attack
+  class Attack
     class Check
       attr_reader :name, :block, :type
       def initialize(name, block)
diff --git a/lib/rack/attack/fail2ban.rb b/lib/rack/attack/fail2ban.rb
index be0cba9..443182d 100644
--- a/lib/rack/attack/fail2ban.rb
+++ b/lib/rack/attack/fail2ban.rb
@@ -1,5 +1,5 @@
 module Rack
-  module Attack
+  class Attack
     class Fail2Ban
       class << self
         def filter(discriminator, options)
diff --git a/lib/rack/attack/store_proxy.rb b/lib/rack/attack/store_proxy.rb
index 6f6a507..4a6d8cc 100644
--- a/lib/rack/attack/store_proxy.rb
+++ b/lib/rack/attack/store_proxy.rb
@@ -1,7 +1,7 @@
 require 'delegate'
 
 module Rack
-  module Attack
+  class Attack
     class StoreProxy
       def self.build(store)
         # RedisStore#increment needs different behavior, so detect that
diff --git a/lib/rack/attack/throttle.rb b/lib/rack/attack/throttle.rb
index fa82266..7ad6267 100644
--- a/lib/rack/attack/throttle.rb
+++ b/lib/rack/attack/throttle.rb
@@ -1,5 +1,5 @@
 module Rack
-  module Attack
+  class Attack
     class Throttle
       MANDATORY_OPTIONS = [:limit, :period]
       attr_reader :name, :limit, :period, :block
diff --git a/lib/rack/attack/track.rb b/lib/rack/attack/track.rb
index c9d9152..3e516d7 100644
--- a/lib/rack/attack/track.rb
+++ b/lib/rack/attack/track.rb
@@ -1,5 +1,5 @@
 module Rack
-  module Attack
+  class Attack
     class Track < Check
       def initialize(name, block)
         super
diff --git a/lib/rack/attack/version.rb b/lib/rack/attack/version.rb
index 8dc08be..c70dd50 100644
--- a/lib/rack/attack/version.rb
+++ b/lib/rack/attack/version.rb
@@ -1,5 +1,5 @@
 module Rack
-  module Attack
+  class Attack
     VERSION = '3.0.0'
   end
 end
diff --git a/lib/rack/attack/whitelist.rb b/lib/rack/attack/whitelist.rb
index cd2699b..604268e 100644
--- a/lib/rack/attack/whitelist.rb
+++ b/lib/rack/attack/whitelist.rb
@@ -1,5 +1,5 @@
 module Rack
-  module Attack
+  class Attack
     class Whitelist < Check
       def initialize(name, block)
         super