Move individual proxy classes to separate files

lib/rack/attack.rb

@@ -2,16 +2,18 @@ require 'rack'
 require 'forwardable'
 
 class Rack::Attack
-  autoload :Cache,     'rack/attack/cache'
+  autoload :Cache,           'rack/attack/cache'
-  autoload :Check,     'rack/attack/check'
+  autoload :Check,           'rack/attack/check'
-  autoload :Throttle,  'rack/attack/throttle'
+  autoload :Throttle,        'rack/attack/throttle'
-  autoload :Whitelist, 'rack/attack/whitelist'
+  autoload :Whitelist,       'rack/attack/whitelist'
-  autoload :Blacklist, 'rack/attack/blacklist'
+  autoload :Blacklist,       'rack/attack/blacklist'
-  autoload :Track,     'rack/attack/track'
+  autoload :Track,           'rack/attack/track'
-  autoload :StoreProxy,'rack/attack/store_proxy'
+  autoload :StoreProxy,      'rack/attack/store_proxy'
-  autoload :Fail2Ban,  'rack/attack/fail2ban'
+  autoload :DalliProxy,      'rack/attack/store_proxy/dalli_proxy'
-  autoload :Allow2Ban, 'rack/attack/allow2ban'
+  autoload :RedisStoreProxy, 'rack/attack/store_proxy/redis_store_proxy'
-  autoload :Request,   'rack/attack/request'
+  autoload :Fail2Ban,        'rack/attack/fail2ban'
+  autoload :Allow2Ban,       'rack/attack/allow2ban'
+  autoload :Request,         'rack/attack/request'
 
   class << self
 

lib/rack/attack/store_proxy.rb

@@ -1,8 +1,8 @@
-require 'delegate'
-
 module Rack
   class Attack
-    class StoreProxy
+    module StoreProxy
+      PROXIES = [DalliProxy, RedisStoreProxy]
+
       def self.build(store)
         # RedisStore#increment needs different behavior, so detect that
         # (method has an arity of 2; must call #expire separately
@@ -12,102 +12,11 @@ module Rack
           store = store.instance_variable_get(:@data)
         end
 
-        if defined?(::Redis::Store) && store.is_a?(::Redis::Store)
+        klass = PROXIES.find { |proxy| proxy.handle?(store) }
-          RedisStoreProxy.new(store)
+
-        elsif defined?(::Dalli) && store.is_a?(::Dalli::Client)
+        klass ? klass.new(store) : store
-          DalliProxy.new(store)
-        elsif defined?(::ConnectionPool) && store.is_a?(::ConnectionPool)
-          store.with do |conn|
-            if conn.is_a?(::Dalli::Client)
-              DalliProxy.new(store)
-            else
-              raise NotImplementedError
-            end
-          end
-        else
-          store
-        end
       end
 
-      class RedisStoreProxy < SimpleDelegator
-        def initialize(store)
-          super(store)
-        end
-
-        def read(key)
-          self.get(key)
-          rescue Redis::BaseError
-            nil
-        end
-
-        def write(key, value, options={})
-          if (expires_in = options[:expires_in])
-            self.setex(key, expires_in, value)
-          else
-            self.set(key, value)
-          end
-          rescue Redis::BaseError
-            nil
-        end
-
-        def increment(key, amount, options={})
-          count = nil
-          self.pipelined do
-            count = self.incrby(key, amount)
-            self.expire(key, options[:expires_in]) if options[:expires_in]
-          end
-          count.value if count
-          rescue Redis::BaseError
-            nil
-        end
-
-      end
-
-      class DalliProxy < SimpleDelegator
-        def initialize(client)
-          super(client)
-          stub_with_method_if_missing
-        end
-
-        def read(key)
-          with do |client|
-            client.get(key)
-          end
-        rescue Dalli::DalliError
-        end
-
-        def write(key, value, options={})
-          with do |client|
-            client.set(key, value, options.fetch(:expires_in, 0), raw: true)
-          end
-        rescue Dalli::DalliError
-        end
-
-        def increment(key, amount, options={})
-          with do |client|
-            client.incr(key, amount, options.fetch(:expires_in, 0), amount)
-          end
-        rescue Dalli::DalliError
-        end
-
-        def delete(key)
-          with do |client|
-            client.delete(key)
-          end
-        rescue Dalli::DalliError
-        end
-
-        private
-
-        def stub_with_method_if_missing
-          unless __getobj__.respond_to?(:with)
-            class << self
-              def with; yield __getobj__; end
-            end
-          end
-        end
-
-      end
     end
   end
 end

lib/rack/attack/store_proxy/dalli_proxy.rb

@@ -0,0 +1,65 @@
+require 'delegate'
+
+module Rack
+  class Attack
+    module StoreProxy
+      class DalliProxy < SimpleDelegator
+        def self.handle?(store)
+          return false unless defined?(::Dalli)
+
+          # Consider extracting to a separate Connection Pool proxy to reduce
+          # code here and handle clients other than Dalli.
+          if defined?(::ConnectionPool) && store.is_a?(::ConnectionPool)
+            store.with { |conn| conn.is_a?(::Dalli::Client) }
+          else
+            store.is_a?(::Dalli::Client)
+          end
+        end
+
+        def initialize(client)
+          super(client)
+          stub_with_if_missing
+        end
+
+        def read(key)
+          with do |client|
+            client.get(key)
+          end
+        rescue Dalli::DalliError
+        end
+
+        def write(key, value, options={})
+          with do |client|
+            client.set(key, value, options.fetch(:expires_in, 0), raw: true)
+          end
+        rescue Dalli::DalliError
+        end
+
+        def increment(key, amount, options={})
+          with do |client|
+            client.incr(key, amount, options.fetch(:expires_in, 0), amount)
+          end
+        rescue Dalli::DalliError
+        end
+
+        def delete(key)
+          with do |client|
+            client.delete(key)
+          end
+        rescue Dalli::DalliError
+        end
+
+        private
+
+        def stub_with_if_missing
+          unless __getobj__.respond_to?(:with)
+            class << self
+              def with; yield __getobj__; end
+            end
+          end
+        end
+
+      end
+    end
+  end
+end

lib/rack/attack/store_proxy/redis_store_proxy.rb

@@ -0,0 +1,42 @@
+require 'delegate'
+
+module Rack
+  class Attack
+    module StoreProxy
+      class RedisStoreProxy < SimpleDelegator
+        def self.handle?(store)
+          defined?(::Redis::Store) && store.is_a?(::Redis::Store)
+        end
+
+        def initialize(store)
+          super(store)
+        end
+
+        def read(key)
+          self.get(key)
+        rescue Redis::BaseError
+        end
+
+        def write(key, value, options={})
+          if (expires_in = options[:expires_in])
+            self.setex(key, expires_in, value)
+          else
+            self.set(key, value)
+          end
+        rescue Redis::BaseError
+        end
+
+        def increment(key, amount, options={})
+          count = nil
+          self.pipelined do
+            count = self.incrby(key, amount)
+            self.expire(key, options[:expires_in]) if options[:expires_in]
+          end
+          count.value if count
+        rescue Redis::BaseError
+        end
+
+      end
+    end
+  end
+end

spec/integration/rack_attack_cache_spec.rb

@@ -83,12 +83,4 @@ describe Rack::Attack::Cache do
     end
 
   end
-
-  describe "given an older Dalli::Client" do
-    it "should stub #with" do
-      proxy = Rack::Attack::StoreProxy::DalliProxy.new(Class.new)
-      proxy.with {} # will not raise an error
-    end
-  end
-
 end

spec/rack_attack_dalli_proxy_spec.rb

@@ -0,0 +1,10 @@
+require_relative 'spec_helper'
+
+describe Rack::Attack::StoreProxy::DalliProxy do
+
+  it 'should stub Dalli::Client#with on older clients' do
+    proxy = Rack::Attack::StoreProxy::DalliProxy.new(Class.new)
+    proxy.with {} # will not raise an error
+  end
+
+end