diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb index d8308a5..14a4ed9 100644 --- a/lib/rack/attack.rb +++ b/lib/rack/attack.rb @@ -2,18 +2,19 @@ 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 :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' + 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 :MemCacheProxy, 'rack/attack/store_proxy/mem_cache_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 f28ea79..b408c5a 100644 --- a/lib/rack/attack/store_proxy.rb +++ b/lib/rack/attack/store_proxy.rb @@ -1,7 +1,7 @@ module Rack class Attack module StoreProxy - PROXIES = [DalliProxy, RedisStoreProxy] + PROXIES = [DalliProxy, MemCacheProxy, RedisStoreProxy] def self.build(store) # RedisStore#increment needs different behavior, so detect that @@ -14,7 +14,7 @@ module Rack # We also want to use the underlying Dalli client instead of ::ActiveSupport::Cache::MemCacheStore, # but not the Memcache client if using Rails 3.x client = store.instance_variable_get(:@data) - if client.is_a?(Redis::Store) || client.is_a?(Dalli::Client) + if client.is_a?(Redis::Store) || client.is_a?(Dalli::Client) || client.is_a?(MemCache) store = store.instance_variable_get(:@data) end end diff --git a/lib/rack/attack/store_proxy/mem_cache_proxy.rb b/lib/rack/attack/store_proxy/mem_cache_proxy.rb new file mode 100644 index 0000000..098e048 --- /dev/null +++ b/lib/rack/attack/store_proxy/mem_cache_proxy.rb @@ -0,0 +1,51 @@ +module Rack + class Attack + module StoreProxy + class MemCacheProxy < SimpleDelegator + def self.handle?(store) + defined?(::MemCache) && store.is_a?(::MemCache) + end + + def initialize(store) + super(store) + stub_with_if_missing + end + + def read(key) + # Second argument: reading raw value + get(key, true) + rescue MemCache::MemCacheError + end + + def write(key, value, options={}) + # Third argument: writing raw value + set(key, value, options.fetch(:expires_in, 0), true) + rescue MemCache::MemCacheError + end + + def increment(key, amount, options={}) + incr(key, amount) + rescue MemCache::MemCacheError + end + + def delete(key, options={}) + with do |client| + client.delete(key) + end + rescue MemCache::MemCacheError + 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