From c90893f5105d5458a0ccd61088315e0b9aa5fca2 Mon Sep 17 00:00:00 2001
From: fatkodima <fatkodima123@gmail.com>
Date: Sat, 12 Oct 2019 18:03:18 +0300
Subject: [PATCH] Move from proxies to well-defined adapters [skip ci]

---
 lib/rack/attack.rb                            |  3 --
 lib/rack/attack/cache.rb                      | 32 +++++--------
 .../active_support_redis_store_adapter.rb}    | 26 +++++-----
 .../dalli_adapter.rb}                         | 26 +++++-----
 .../store_adapters/mem_cache_store_adapter.rb | 31 ++++++++++++
 .../redis_adapter.rb}                         | 24 +++++-----
 .../redis_cache_store_adapter.rb}             | 26 +++++-----
 .../redis_store_adapter.rb}                   | 12 ++---
 .../attack/store_adapters/store_adapter.rb    | 47 +++++++++++++++++++
 lib/rack/attack/store_proxy.rb                | 21 ---------
 .../store_proxy/mem_cache_store_proxy.rb      | 21 ---------
 11 files changed, 144 insertions(+), 125 deletions(-)
 rename lib/rack/attack/{store_proxy/active_support_redis_store_proxy.rb => store_adapters/active_support_redis_store_adapter.rb} (52%)
 rename lib/rack/attack/{store_proxy/dalli_proxy.rb => store_adapters/dalli_adapter.rb} (77%)
 create mode 100644 lib/rack/attack/store_adapters/mem_cache_store_adapter.rb
 rename lib/rack/attack/{store_proxy/redis_proxy.rb => store_adapters/redis_adapter.rb} (65%)
 rename lib/rack/attack/{store_proxy/redis_cache_store_proxy.rb => store_adapters/redis_cache_store_adapter.rb} (59%)
 rename lib/rack/attack/{store_proxy/redis_store_proxy.rb => store_adapters/redis_store_adapter.rb} (59%)
 create mode 100644 lib/rack/attack/store_adapters/store_adapter.rb
 delete mode 100644 lib/rack/attack/store_proxy.rb
 delete mode 100644 lib/rack/attack/store_proxy/mem_cache_store_proxy.rb

diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb
index c8d4ad1a..a18a8773 100644
--- a/lib/rack/attack.rb
+++ b/lib/rack/attack.rb
@@ -10,9 +10,6 @@
 
 module Rack
   class Attack
-    class MisconfiguredStoreError < StandardError; end
-    class MissingStoreError < StandardError; end
-
     autoload :Cache,                'rack/attack/cache'
     autoload :Check,                'rack/attack/check'
     autoload :Throttle,             'rack/attack/throttle'
diff --git a/lib/rack/attack/cache.rb b/lib/rack/attack/cache.rb
index cfa2efa4..5f84950e 100644
--- a/lib/rack/attack/cache.rb
+++ b/lib/rack/attack/cache.rb
@@ -13,7 +13,14 @@ def initialize
 
       attr_reader :store
       def store=(store)
-        @store = StoreProxy.build(store)
+        adapter = StoreAdapter.lookup(store)
+        if adapter
+          @store = adapter.new(store)
+        elsif store?(store)
+          @store = store
+        else
+          raise ArgumentError
+        end
       end
 
       def count(unprefixed_key, period)
@@ -22,9 +29,6 @@ def count(unprefixed_key, period)
       end
 
       def read(unprefixed_key)
-        enforce_store_presence!
-        enforce_store_method_presence!(:read)
-
         store.read("#{prefix}:#{unprefixed_key}")
       end
 
@@ -51,33 +55,19 @@ def key_and_expiry(unprefixed_key, period)
       end
 
       def do_count(key, expires_in)
-        enforce_store_presence!
-        enforce_store_method_presence!(:increment)
-
         result = store.increment(key, 1, expires_in: expires_in)
 
         # NB: Some stores return nil when incrementing uninitialized values
         if result.nil?
-          enforce_store_method_presence!(:write)
-
           store.write(key, 1, expires_in: expires_in)
         end
         result || 1
       end
 
-      def enforce_store_presence!
-        if store.nil?
-          raise Rack::Attack::MissingStoreError
-        end
-      end
+      STORE_METHODS = [:read, :write, :increment, :delete].freeze
 
-      def enforce_store_method_presence!(method_name)
-        if !store.respond_to?(method_name)
-          raise(
-            Rack::Attack::MisconfiguredStoreError,
-            "Configured store #{store.class.name} doesn't respond to ##{method_name} method"
-          )
-        end
+      def store?(object)
+        STORE_METHODS.all? { |meth| object.respond_to?(meth) }
       end
     end
   end
diff --git a/lib/rack/attack/store_proxy/active_support_redis_store_proxy.rb b/lib/rack/attack/store_adapters/active_support_redis_store_adapter.rb
similarity index 52%
rename from lib/rack/attack/store_proxy/active_support_redis_store_proxy.rb
rename to lib/rack/attack/store_adapters/active_support_redis_store_adapter.rb
index 68f0326f..7053c27e 100644
--- a/lib/rack/attack/store_proxy/active_support_redis_store_proxy.rb
+++ b/lib/rack/attack/store_adapters/active_support_redis_store_adapter.rb
@@ -1,37 +1,39 @@
 # frozen_string_literal: true
 
-require 'delegate'
-
 module Rack
   class Attack
-    module StoreProxy
-      class ActiveSupportRedisStoreProxy < SimpleDelegator
+    module StoreAdapters
+      class ActiveSupportRedisStoreAdapter < StoreAdapter
         def self.handle?(store)
           defined?(::Redis) &&
             defined?(::ActiveSupport::Cache::RedisStore) &&
             store.is_a?(::ActiveSupport::Cache::RedisStore)
         end
 
-        def increment(name, amount = 1, options = {})
+        def increment(key, amount = 1, options = {})
           # #increment ignores options[:expires_in].
           #
           # So in order to workaround this we use #write (which sets expiration) to initialize
           # the counter. After that we continue using the original #increment.
-          if options[:expires_in] && !read(name)
-            write(name, amount, options)
+          if options[:expires_in] && !read(key)
+            write(key, amount, options)
 
             amount
           else
-            super
+            store.increment(key, amount, options)
           end
         end
 
-        def read(name, options = {})
-          super(name, options.merge!(raw: true))
+        def read(key, options = {})
+          store.read(key, options.merge!(raw: true))
+        end
+
+        def write(key, value, options = {})
+          store.write(key, value, options.merge!(raw: true))
         end
 
-        def write(name, value, options = {})
-          super(name, value, options.merge!(raw: true))
+        def delete(key, options = {})
+          store.delete(key, options)
         end
       end
     end
diff --git a/lib/rack/attack/store_proxy/dalli_proxy.rb b/lib/rack/attack/store_adapters/dalli_adapter.rb
similarity index 77%
rename from lib/rack/attack/store_proxy/dalli_proxy.rb
rename to lib/rack/attack/store_adapters/dalli_adapter.rb
index 360e2198..b2683eee 100644
--- a/lib/rack/attack/store_proxy/dalli_proxy.rb
+++ b/lib/rack/attack/store_adapters/dalli_adapter.rb
@@ -1,11 +1,9 @@
 # frozen_string_literal: true
 
-require 'delegate'
-
 module Rack
   class Attack
-    module StoreProxy
-      class DalliProxy < SimpleDelegator
+    module StoreAdapters
+      class DalliAdapter < StoreAdapter
         def self.handle?(store)
           return false unless defined?(::Dalli)
 
@@ -18,14 +16,14 @@ def self.handle?(store)
           end
         end
 
-        def initialize(client)
-          super(client)
+        def initialize(store)
+          super
           stub_with_if_missing
         end
 
         def read(key)
           rescuing do
-            with do |client|
+            store.with do |client|
               client.get(key)
             end
           end
@@ -33,7 +31,7 @@ def read(key)
 
         def write(key, value, options = {})
           rescuing do
-            with do |client|
+            store.with do |client|
               client.set(key, value, options.fetch(:expires_in, 0), raw: true)
             end
           end
@@ -41,7 +39,7 @@ def write(key, value, options = {})
 
         def increment(key, amount, options = {})
           rescuing do
-            with do |client|
+            store.with do |client|
               client.incr(key, amount, options.fetch(:expires_in, 0), amount)
             end
           end
@@ -49,7 +47,7 @@ def increment(key, amount, options = {})
 
         def delete(key)
           rescuing do
-            with do |client|
+            store.with do |client|
               client.delete(key)
             end
           end
@@ -58,11 +56,9 @@ def delete(key)
         private
 
         def stub_with_if_missing
-          unless __getobj__.respond_to?(:with)
-            class << self
-              def with
-                yield __getobj__
-              end
+          unless store.respond_to?(:with)
+            def store.with
+              yield store
             end
           end
         end
diff --git a/lib/rack/attack/store_adapters/mem_cache_store_adapter.rb b/lib/rack/attack/store_adapters/mem_cache_store_adapter.rb
new file mode 100644
index 00000000..aa7f0127
--- /dev/null
+++ b/lib/rack/attack/store_adapters/mem_cache_store_adapter.rb
@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Rack
+  class Attack
+    module StoreAdapters
+      class MemCacheStoreAdapter < StoreAdapter
+        def self.handle?(store)
+          defined?(::Dalli) &&
+            defined?(::ActiveSupport::Cache::MemCacheStore) &&
+            store.is_a?(::ActiveSupport::Cache::MemCacheStore)
+        end
+
+        def read(key, options = {})
+          store.read(key, options)
+        end
+
+        def write(key, value, options = {})
+          store.write(key, value, options.merge!(raw: true))
+        end
+
+        def increment(*args)
+          store.increment(*args)
+        end
+
+        def delete(*args)
+          store.delete(*args)
+        end
+      end
+    end
+  end
+end
diff --git a/lib/rack/attack/store_proxy/redis_proxy.rb b/lib/rack/attack/store_adapters/redis_adapter.rb
similarity index 65%
rename from lib/rack/attack/store_proxy/redis_proxy.rb
rename to lib/rack/attack/store_adapters/redis_adapter.rb
index 9fe2bc8c..6440f3b2 100644
--- a/lib/rack/attack/store_proxy/redis_proxy.rb
+++ b/lib/rack/attack/store_adapters/redis_adapter.rb
@@ -1,17 +1,15 @@
 # frozen_string_literal: true
 
-require 'delegate'
-
 module Rack
   class Attack
-    module StoreProxy
-      class RedisProxy < SimpleDelegator
-        def initialize(*args)
+    module StoreAdapters
+      class RedisProxy < StoreAdapter
+        def initialize(store)
           if Gem::Version.new(Redis::VERSION) < Gem::Version.new("3")
             warn 'RackAttack requires Redis gem >= 3.0.0.'
           end
 
-          super(*args)
+          super
         end
 
         def self.handle?(store)
@@ -19,14 +17,14 @@ def self.handle?(store)
         end
 
         def read(key)
-          rescuing { get(key) }
+          rescuing { store.get(key) }
         end
 
         def write(key, value, options = {})
           if (expires_in = options[:expires_in])
-            rescuing { setex(key, expires_in, value) }
+            rescuing { store.setex(key, expires_in, value) }
           else
-            rescuing { set(key, value) }
+            rescuing { store.set(key, value) }
           end
         end
 
@@ -34,9 +32,9 @@ def increment(key, amount, options = {})
           count = nil
 
           rescuing do
-            pipelined do
-              count = incrby(key, amount)
-              expire(key, options[:expires_in]) if options[:expires_in]
+            store.pipelined do
+              count = store.incrby(key, amount)
+              store.expire(key, options[:expires_in]) if options[:expires_in]
             end
           end
 
@@ -44,7 +42,7 @@ def increment(key, amount, options = {})
         end
 
         def delete(key, _options = {})
-          rescuing { del(key) }
+          rescuing { store.del(key) }
         end
 
         private
diff --git a/lib/rack/attack/store_proxy/redis_cache_store_proxy.rb b/lib/rack/attack/store_adapters/redis_cache_store_adapter.rb
similarity index 59%
rename from lib/rack/attack/store_proxy/redis_cache_store_proxy.rb
rename to lib/rack/attack/store_adapters/redis_cache_store_adapter.rb
index f4081bee..260b9925 100644
--- a/lib/rack/attack/store_proxy/redis_cache_store_proxy.rb
+++ b/lib/rack/attack/store_adapters/redis_cache_store_adapter.rb
@@ -1,41 +1,43 @@
 # frozen_string_literal: true
 
-require 'delegate'
-
 module Rack
   class Attack
-    module StoreProxy
-      class RedisCacheStoreProxy < SimpleDelegator
+    module StoreAdapters
+      class RedisCacheStoreAdapter < StoreAdapter
         def self.handle?(store)
           store.class.name == "ActiveSupport::Cache::RedisCacheStore"
         end
 
-        def increment(name, amount = 1, options = {})
+        def increment(key, amount = 1, options = {})
           # RedisCacheStore#increment ignores options[:expires_in].
           #
           # So in order to workaround this we use RedisCacheStore#write (which sets expiration) to initialize
           # the counter. After that we continue using the original RedisCacheStore#increment.
           rescuing do
-            if options[:expires_in] && !read(name)
-              write(name, amount, options)
+            if options[:expires_in] && !read(key)
+              write(key, amount, options)
 
               amount
             else
-              super
+              store.increment(key, amount, options)
             end
           end
         end
 
-        def read(*_args)
-          rescuing { super }
+        def read(*args)
+          rescuing { store.read(*args) }
         end
 
-        def write(name, value, options = {})
+        def write(key, value, options = {})
           rescuing do
-            super(name, value, options.merge!(raw: true))
+            store.write(key, value, options.merge!(raw: true))
           end
         end
 
+        def delete(*args)
+          rescuing { store.delete(*args) }
+        end
+
         private
 
         def rescuing
diff --git a/lib/rack/attack/store_proxy/redis_store_proxy.rb b/lib/rack/attack/store_adapters/redis_store_adapter.rb
similarity index 59%
rename from lib/rack/attack/store_proxy/redis_store_proxy.rb
rename to lib/rack/attack/store_adapters/redis_store_adapter.rb
index 6be54128..bfe1ebab 100644
--- a/lib/rack/attack/store_proxy/redis_store_proxy.rb
+++ b/lib/rack/attack/store_adapters/redis_store_adapter.rb
@@ -1,24 +1,22 @@
 # frozen_string_literal: true
 
-require 'delegate'
-
 module Rack
   class Attack
-    module StoreProxy
-      class RedisStoreProxy < RedisProxy
+    module StoreAdapters
+      class RedisStoreAdapter < RedisAdapter
         def self.handle?(store)
           defined?(::Redis::Store) && store.is_a?(::Redis::Store)
         end
 
         def read(key)
-          rescuing { get(key, raw: true) }
+          rescuing { store.get(key, raw: true) }
         end
 
         def write(key, value, options = {})
           if (expires_in = options[:expires_in])
-            rescuing { setex(key, expires_in, value, raw: true) }
+            rescuing { store.setex(key, expires_in, value, raw: true) }
           else
-            rescuing { set(key, value, raw: true) }
+            rescuing { store.set(key, value, raw: true) }
           end
         end
       end
diff --git a/lib/rack/attack/store_adapters/store_adapter.rb b/lib/rack/attack/store_adapters/store_adapter.rb
new file mode 100644
index 00000000..52fff6be
--- /dev/null
+++ b/lib/rack/attack/store_adapters/store_adapter.rb
@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module Rack
+  class Attack
+    module StoreAdapters
+      class StoreAdapter
+        @adapters = []
+
+        class << self
+          def inherited(klass)
+            @adapters << klass
+          end
+
+          def lookup(store)
+            @adapters.find { |adapter| adapter.handle?(store) }
+          end
+
+          def handle?(store)
+            raise NotImplementedError
+          end
+        end
+
+        attr_reader :store
+
+        def initialize(store)
+          @store = store
+        end
+
+        def read(key)
+          raise NotImplementedError
+        end
+
+        def write(key, value, options = {})
+          raise NotImplementedError
+        end
+
+        def increment(key, amount, options = {})
+          raise NotImplementedError
+        end
+
+        def delete(key, options = {})
+          raise NotImplementedError
+        end
+      end
+    end
+  end
+end
diff --git a/lib/rack/attack/store_proxy.rb b/lib/rack/attack/store_proxy.rb
deleted file mode 100644
index 55c63e21..00000000
--- a/lib/rack/attack/store_proxy.rb
+++ /dev/null
@@ -1,21 +0,0 @@
-# frozen_string_literal: true
-
-module Rack
-  class Attack
-    module StoreProxy
-      PROXIES = [
-        DalliProxy,
-        MemCacheStoreProxy,
-        RedisStoreProxy,
-        RedisProxy,
-        RedisCacheStoreProxy,
-        ActiveSupportRedisStoreProxy
-      ].freeze
-
-      def self.build(store)
-        klass = PROXIES.find { |proxy| proxy.handle?(store) }
-        klass ? klass.new(store) : store
-      end
-    end
-  end
-end
diff --git a/lib/rack/attack/store_proxy/mem_cache_store_proxy.rb b/lib/rack/attack/store_proxy/mem_cache_store_proxy.rb
deleted file mode 100644
index f8c036c9..00000000
--- a/lib/rack/attack/store_proxy/mem_cache_store_proxy.rb
+++ /dev/null
@@ -1,21 +0,0 @@
-# frozen_string_literal: true
-
-require 'delegate'
-
-module Rack
-  class Attack
-    module StoreProxy
-      class MemCacheStoreProxy < SimpleDelegator
-        def self.handle?(store)
-          defined?(::Dalli) &&
-            defined?(::ActiveSupport::Cache::MemCacheStore) &&
-            store.is_a?(::ActiveSupport::Cache::MemCacheStore)
-        end
-
-        def write(name, value, options = {})
-          super(name, value, options.merge!(raw: true))
-        end
-      end
-    end
-  end
-end