rate_limiter.rb 1.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. # frozen_string_literal: true
  2. class RateLimiter
  3. include Redisable
  4. FAMILIES = {
  5. follows: {
  6. limit: 400,
  7. period: 24.hours.freeze,
  8. }.freeze,
  9. statuses: {
  10. limit: 300,
  11. period: 3.hours.freeze,
  12. }.freeze,
  13. reports: {
  14. limit: 400,
  15. period: 24.hours.freeze,
  16. }.freeze,
  17. }.freeze
  18. def initialize(by, options = {})
  19. @by = by
  20. @family = options[:family]
  21. @limit = FAMILIES[@family][:limit]
  22. @period = FAMILIES[@family][:period].to_i
  23. end
  24. def record!
  25. count = redis.get(key)
  26. if count.nil?
  27. redis.set(key, 0)
  28. redis.expire(key, (@period - (last_epoch_time % @period) + 1).to_i)
  29. end
  30. raise Mastodon::RateLimitExceededError if count.present? && count.to_i >= @limit
  31. redis.incr(key)
  32. end
  33. def rollback!
  34. redis.decr(key)
  35. end
  36. def to_headers(now = Time.now.utc)
  37. {
  38. 'X-RateLimit-Limit' => @limit.to_s,
  39. 'X-RateLimit-Remaining' => (@limit - (redis.get(key) || 0).to_i).to_s,
  40. 'X-RateLimit-Reset' => (now + (@period - now.to_i % @period)).iso8601(6),
  41. }
  42. end
  43. private
  44. def key
  45. @key ||= "rate_limit:#{@by.id}:#{@family}:#{(last_epoch_time / @period).to_i}"
  46. end
  47. def last_epoch_time
  48. @last_epoch_time ||= Time.now.to_i
  49. end
  50. end