Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feature: support for ActiveSupport::MemCacheStore #153

Merged
merged 9 commits into from
Feb 8, 2016
1 change: 1 addition & 0 deletions lib/rack/attack.rb
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ class Rack::Attack
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'
Expand Down
18 changes: 12 additions & 6 deletions lib/rack/attack/store_proxy.rb
Original file line number Diff line number Diff line change
@@ -1,19 +1,25 @@
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
# (method has an arity of 2; must call #expire separately
if defined?(::ActiveSupport::Cache::RedisStore) && store.is_a?(::ActiveSupport::Cache::RedisStore)
if (defined?(::ActiveSupport::Cache::RedisStore) && store.is_a?(::ActiveSupport::Cache::RedisStore)) ||
(defined?(::ActiveSupport::Cache::MemCacheStore) && store.is_a?(::ActiveSupport::Cache::MemCacheStore))

# ActiveSupport::Cache::RedisStore doesn't expose any way to set an expiry,
# so use the raw Redis::Store instead
store = store.instance_variable_get(:@data)
# so use the raw Redis::Store instead.
# We also want to use the underlying Dalli client instead of ::ActiveSupport::Cache::MemCacheStore,
# and the MemCache client if using Rails 3.x
client = store.instance_variable_get(:@data)
if (defined?(::Redis::Store) && client.is_a?(Redis::Store)) ||
(defined?(Dalli::Client) && client.is_a?(Dalli::Client)) || (defined?(MemCache) && client.is_a?(MemCache))
store = store.instance_variable_get(:@data)
end
end

klass = PROXIES.find { |proxy| proxy.handle?(store) }

klass ? klass.new(store) : store
end

Expand Down
51 changes: 51 additions & 0 deletions lib/rack/attack/store_proxy/mem_cache_proxy.rb
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions rack-attack.gemspec
Original file line number Diff line number Diff line change
Expand Up @@ -32,4 +32,5 @@ Gem::Specification.new do |s|
s.add_development_dependency 'redis-activesupport'
s.add_development_dependency 'dalli'
s.add_development_dependency 'connection_pool'
s.add_development_dependency 'memcache-client'
end
3 changes: 3 additions & 0 deletions spec/integration/rack_attack_cache_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,14 @@ def sleep_until_expired
end

require 'active_support/cache/dalli_store'
require 'active_support/cache/mem_cache_store'
require 'active_support/cache/redis_store'
require 'connection_pool'
cache_stores = [
ActiveSupport::Cache::MemoryStore.new,
ActiveSupport::Cache::DalliStore.new("127.0.0.1"),
ActiveSupport::Cache::RedisStore.new("127.0.0.1"),
ActiveSupport::Cache::MemCacheStore.new("127.0.0.1"),
Dalli::Client.new,
ConnectionPool.new { Dalli::Client.new },
Redis::Store.new
Expand Down Expand Up @@ -54,6 +56,7 @@ def sleep_until_expired
@cache.send(:do_count, @key, @expires_in).must_equal 2
end
end

describe "do_count after expires_in" do
it "must be 1" do
@cache.send(:do_count, @key, @expires_in)
Expand Down