Skip to content

Commit

Permalink
Move from proxies to well-defined adapters [skip ci]
Browse files Browse the repository at this point in the history
  • Loading branch information
fatkodima committed Oct 12, 2019
1 parent 95347e3 commit c90893f
Show file tree
Hide file tree
Showing 11 changed files with 144 additions and 125 deletions.
3 changes: 0 additions & 3 deletions lib/rack/attack.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down
32 changes: 11 additions & 21 deletions lib/rack/attack/cache.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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

Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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)

Expand All @@ -18,38 +16,38 @@ 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
end

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
end

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
end

def delete(key)
rescuing do
with do |client|
store.with do |client|
client.delete(key)
end
end
Expand All @@ -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
Expand Down
31 changes: 31 additions & 0 deletions lib/rack/attack/store_adapters/mem_cache_store_adapter.rb
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -1,50 +1,48 @@
# 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)
defined?(::Redis) && store.is_a?(::Redis)
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

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

count.value if count
end

def delete(key, _options = {})
rescuing { del(key) }
rescuing { store.del(key) }
end

private
Expand Down
Original file line number Diff line number Diff line change
@@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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
Expand Down
Loading

0 comments on commit c90893f

Please sign in to comment.