From 5d72c6e5f92190bef89b67a5629fa6051440a72d Mon Sep 17 00:00:00 2001 From: hakanensari Date: Tue, 15 Apr 2014 16:18:34 +0100 Subject: [PATCH] Move individual proxy classes to separate files --- lib/rack/attack.rb | 22 ++-- lib/rack/attack/store_proxy.rb | 103 +----------------- lib/rack/attack/store_proxy/dalli_proxy.rb | 65 +++++++++++ .../attack/store_proxy/redis_store_proxy.rb | 42 +++++++ spec/integration/rack_attack_cache_spec.rb | 8 -- spec/rack_attack_dalli_proxy_spec.rb | 10 ++ 6 files changed, 135 insertions(+), 115 deletions(-) create mode 100644 lib/rack/attack/store_proxy/dalli_proxy.rb create mode 100644 lib/rack/attack/store_proxy/redis_store_proxy.rb create mode 100644 spec/rack_attack_dalli_proxy_spec.rb diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb index 24a74b5..4bfe33c 100644 --- a/lib/rack/attack.rb +++ b/lib/rack/attack.rb @@ -2,16 +2,18 @@ require 'rack' require 'forwardable' class Rack::Attack - autoload :Cache, 'rack/attack/cache' - autoload :Check, 'rack/attack/check' - autoload :Throttle, 'rack/attack/throttle' - autoload :Whitelist, 'rack/attack/whitelist' - autoload :Blacklist, 'rack/attack/blacklist' - autoload :Track, 'rack/attack/track' - autoload :StoreProxy,'rack/attack/store_proxy' - autoload :Fail2Ban, 'rack/attack/fail2ban' - autoload :Allow2Ban, 'rack/attack/allow2ban' - autoload :Request, 'rack/attack/request' + autoload :Cache, 'rack/attack/cache' + autoload :Check, 'rack/attack/check' + autoload :Throttle, 'rack/attack/throttle' + autoload :Whitelist, 'rack/attack/whitelist' + autoload :Blacklist, 'rack/attack/blacklist' + autoload :Track, 'rack/attack/track' + autoload :StoreProxy, 'rack/attack/store_proxy' + autoload :DalliProxy, 'rack/attack/store_proxy/dalli_proxy' + autoload :RedisStoreProxy, 'rack/attack/store_proxy/redis_store_proxy' + autoload :Fail2Ban, 'rack/attack/fail2ban' + autoload :Allow2Ban, 'rack/attack/allow2ban' + autoload :Request, 'rack/attack/request' class << self diff --git a/lib/rack/attack/store_proxy.rb b/lib/rack/attack/store_proxy.rb index 3c4f60d..5c476ff 100644 --- a/lib/rack/attack/store_proxy.rb +++ b/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) - RedisStoreProxy.new(store) - elsif defined?(::Dalli) && store.is_a?(::Dalli::Client) - 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 + klass = PROXIES.find { |proxy| proxy.handle?(store) } + + klass ? klass.new(store) : store 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 diff --git a/lib/rack/attack/store_proxy/dalli_proxy.rb b/lib/rack/attack/store_proxy/dalli_proxy.rb new file mode 100644 index 0000000..703f2d6 --- /dev/null +++ b/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 diff --git a/lib/rack/attack/store_proxy/redis_store_proxy.rb b/lib/rack/attack/store_proxy/redis_store_proxy.rb new file mode 100644 index 0000000..5677898 --- /dev/null +++ b/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 diff --git a/spec/integration/rack_attack_cache_spec.rb b/spec/integration/rack_attack_cache_spec.rb index df92c41..14bec01 100644 --- a/spec/integration/rack_attack_cache_spec.rb +++ b/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 diff --git a/spec/rack_attack_dalli_proxy_spec.rb b/spec/rack_attack_dalli_proxy_spec.rb new file mode 100644 index 0000000..50a2b9c --- /dev/null +++ b/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