diff --git a/app/lib/shunter.rb b/app/lib/shunter.rb new file mode 100644 index 000000000..045f7fdd0 --- /dev/null +++ b/app/lib/shunter.rb @@ -0,0 +1,27 @@ +module Shunter + module_function + + class ThrottleError < StandardError; end + + BASE_LOGGER = ::Logger.new($stdout) + ONE_MINUTE = 60 + ONE_HUNDRED_REQUESTS = 100 + + BASE_CONNECTION = ENV['shunter_redis_connection'] || { host: 'redis', port: 6379 } + + def default_timespan + ENV['shunter_default_timespan'] || ONE_MINUTE + end + + def default_threshold + ENV['shunter_default_threshold'] || ONE_HUNDRED_REQUESTS + end + + def default_adapter + ENV['shunter_default_adapter'] || 'Shunter::Adapters::Redis' + end + + def feature_enabled? + ActiveModel::Type::Boolean.new.cast(ENV['shunter_enabled'] || 'false') + end +end diff --git a/app/lib/shunter/adapters/memory.rb b/app/lib/shunter/adapters/memory.rb new file mode 100644 index 000000000..eb0b25b27 --- /dev/null +++ b/app/lib/shunter/adapters/memory.rb @@ -0,0 +1,31 @@ +# frozen_string_literal: true + +module Shunter + module Adapters + class Memory + attr_reader :store + + def initialize(_options = {}) + @@store ||= {} + end + + def find_counter(key) + @@store[key] + end + + def write_counter(key) + @@store[key] = 1 + end + + def increment_counter(key) + @@store[key] += 1 + end + + def clear! + @@store = {} + end + + def expire_counter(_key, _timespan); end + end + end +end diff --git a/app/lib/shunter/adapters/redis.rb b/app/lib/shunter/adapters/redis.rb new file mode 100644 index 000000000..adf27b359 --- /dev/null +++ b/app/lib/shunter/adapters/redis.rb @@ -0,0 +1,29 @@ +# frozen_string_literal: true + +module Shunter + module Adapters + class Redis + attr_reader :redis + + def initialize(options) + @redis = ::Redis.new(options) + end + + def find_counter(key) + @redis.get(key) + end + + def write_counter(key) + @redis.set(key, 1) + end + + def increment_counter(key) + @redis.incr(key) + end + + def expire_counter(key, timespan) + @redis.expire(key, timespan) + end + end + end +end diff --git a/app/lib/shunter/base.rb b/app/lib/shunter/base.rb new file mode 100644 index 000000000..f3f6867f0 --- /dev/null +++ b/app/lib/shunter/base.rb @@ -0,0 +1,63 @@ +# frozen_string_literal: true + +module Shunter + class Base + attr_accessor :user_id, :adapter + + def initialize(options = {}) + @user_id = options[:user_id] + adapter_klass = Shunter.default_adapter.constantize + @adapter = adapter_klass.new(options[:conn_options]) + end + + def user_key + "counting_#{@user_id}" + end + + def blocked_user_key + "blocked_#{@user_id}" + end + + def throttle + return false if blocked? + + valid_counter? + end + + def blocked? + adapter.find_counter(blocked_user_key).present? + end + + def valid_counter? + if adapter.find_counter(user_key) + number_of_requests = adapter.increment_counter(user_key) + if number_of_requests > allowed_requests.to_i + init_counter(blocked_user_key) + return false + end + else + init_counter(user_key) + end + true + end + + private + + def init_counter(key) + adapter.write_counter(key) + adapter.expire_counter(key, timespan) + end + + def allowed_requests + Shunter.default_threshold + end + + def timespan + Shunter.default_timespan + end + + def logger + Shunter::BASE_LOGGER + end + end +end diff --git a/app/lib/shunter/integration/throttle.rb b/app/lib/shunter/integration/throttle.rb new file mode 100644 index 000000000..0b4517112 --- /dev/null +++ b/app/lib/shunter/integration/throttle.rb @@ -0,0 +1,44 @@ +# frozen_string_literal: true + +require 'active_support/concern' + +module Shunter + module Integration + module Throttle + extend ActiveSupport::Concern + + included do |base| + actions = base.const_defined?('THROTTLED_ACTIONS') && base.const_get('THROTTLED_ACTIONS') + return if actions.blank? + + around_action :throttle, only: actions + + def throttle + unless throttled_user.present? && Shunter.feature_enabled? + yield if block_given? + return + end + + user_id = throttled_user.id + + shunter = Shunter::Base.new(conn_options: connection_options, user_id: user_id) + if shunter.throttle + logger.info "Request from #{throttled_user.class}/#{throttled_user.id} is coming through throttling" + yield if block_given? + else + logger.info "Too many requests from #{throttled_user.class}/#{throttled_user.id}." + raise Shunter::ThrottleError + end + end + end + + def connection_options + Shunter::BASE_CONNECTION + end + + def logger + Shunter::BASE_LOGGER + end + end + end +end diff --git a/config/environments/test.rb b/config/environments/test.rb index 028c61b47..6e680b9c0 100644 --- a/config/environments/test.rb +++ b/config/environments/test.rb @@ -41,4 +41,6 @@ Rails.application.configure do # If set to :null_store, Setting.x returns nil after first spec runs (database is emptied) config.cache_store = :memory_store + + config.log_level = :fatal end diff --git a/test/integration/epp/domain/info/base_test.rb b/test/integration/epp/domain/info/base_test.rb index 99de33f29..25c65d476 100644 --- a/test/integration/epp/domain/info/base_test.rb +++ b/test/integration/epp/domain/info/base_test.rb @@ -193,7 +193,7 @@ class EppDomainInfoBaseTest < EppTestCase - + #{domain.name} diff --git a/test/lib/shunter/base_test.rb b/test/lib/shunter/base_test.rb new file mode 100644 index 000000000..c251724ed --- /dev/null +++ b/test/lib/shunter/base_test.rb @@ -0,0 +1,33 @@ +# frozen_string_literal: true + +require "test_helper" +require "action_controller" +require "pry" + +class BaseTest < Minitest::Test + ENV["shunter_enabled"] = 'true' + + def test_throttling_works_on_inclusion + ENV["shunter_default_adapter"] = "Shunter::Adapters::Memory" + ENV["shunter_default_threshold"] = "100" + adapter = ENV["shunter_default_adapter"].constantize.new + adapter.clear! + + TestKlass.new.throttle do + TestKlass.new.test + end + end + + class TestKlass < ::ActionController::Base + THROTTLED_ACTIONS = %i[test].freeze + include Shunter::Integration::Throttle + + def test + "test" + end + + def throttled_user + @throttled_user ||= OpenStruct.new(id: 1) + end + end +end diff --git a/test/test_helper.rb b/test/test_helper.rb index 0cd407f84..cce9afa97 100644 --- a/test/test_helper.rb +++ b/test/test_helper.rb @@ -18,7 +18,6 @@ require 'capybara/minitest' require 'webmock/minitest' require 'support/assertions/epp_assertions' require 'sidekiq/testing' -require 'spy/integration' Sidekiq::Testing.fake!