Last active
April 9, 2021 20:05
Simple action rate limit module for Ruby on Rails, intended to be included in ApplicationController
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# frozen_string_literal: true | |
module ActionRateLimit | |
LIMIT_BY_OPTIONS = { | |
user: proc { request.cookies['user_id'] }, | |
ip: proc { request.remote_ip }, | |
}.freeze | |
def self.included(klass) | |
klass.extend ClassMethods | |
end | |
module ClassMethods | |
# Creates rate limit rule on the action | |
# action - [symbol] the action in the controller to be limited | |
# limit - [int] max number of requests that can be accepted within period of time | |
# period - [int] number of seconds, can be defined as 1.hour. Can be between 10 seconds and 1 day | |
# limit_by - proc or one of {:user, :ip} - defines what value to group requests by. For login-only requests | |
# makes most sense to use `:user` and put this rule after requiring login | |
# Note: if `limit_by` will evaluate to `nil`, the rate limit check will be skipped | |
def rate_limit(action, limit:, period:, limit_by:) | |
limit_by = LIMIT_BY_OPTIONS[limit_by] || limit_by | |
raise ArgumentError, "id need to be proc or one of: #{LIMIT_BY_OPTIONS.keys.inspect}" unless limit_by.is_a?(Proc) | |
raise ArgumentError, "invalid period" unless (10..1.day).cover?(period) | |
raise ArgumentError, "invalid limit" if limit < 1 | |
before_action only: action do | |
limit_by_value = instance_eval(&limit_by) | |
rate_limit_check!(action, limit, period, limit_by_value) if limit_by_value | |
end | |
end | |
end | |
private | |
def rate_limit_check!(action, limit, period, limit_by_value) | |
now_ts = Time.now.to_f | |
key = "rate_limit:#{self.class.name}##{action}:#{limit_by_value}" | |
redis.zremrangebyscore(key, 0, now_ts - period) | |
return render status: :too_many_requests, plain: "rate limit reached" if redis.zcard(key) >= limit | |
redis.multi do |r| | |
r.zadd(key, now_ts, now_ts) | |
r.expire(key, period) | |
end | |
end | |
def redis | |
# implement with your own reference to redis client | |
end | |
end |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment