rate_limit_headers.rb 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970
  1. # frozen_string_literal: true
  2. module Api::RateLimitHeaders
  3. extend ActiveSupport::Concern
  4. class_methods do
  5. def override_rate_limit_headers(method_name, options = {})
  6. around_action(only: method_name, if: :current_account) do |_controller, block|
  7. block.call
  8. ensure
  9. rate_limiter = RateLimiter.new(current_account, options)
  10. rate_limit_headers = rate_limiter.to_headers
  11. response.headers.merge!(rate_limit_headers) unless response.headers['X-RateLimit-Remaining'].present? && rate_limit_headers['X-RateLimit-Remaining'].to_i > response.headers['X-RateLimit-Remaining'].to_i
  12. end
  13. end
  14. end
  15. included do
  16. before_action :set_rate_limit_headers, if: :rate_limited_request?
  17. end
  18. private
  19. def set_rate_limit_headers
  20. apply_header_limit
  21. apply_header_remaining
  22. apply_header_reset
  23. end
  24. def rate_limited_request?
  25. !request.env['rack.attack.throttle_data'].nil?
  26. end
  27. def apply_header_limit
  28. response.headers['X-RateLimit-Limit'] = rate_limit_limit
  29. end
  30. def rate_limit_limit
  31. api_throttle_data[:limit].to_s
  32. end
  33. def apply_header_remaining
  34. response.headers['X-RateLimit-Remaining'] = rate_limit_remaining
  35. end
  36. def rate_limit_remaining
  37. (api_throttle_data[:limit] - api_throttle_data[:count]).to_s
  38. end
  39. def apply_header_reset
  40. response.headers['X-RateLimit-Reset'] = rate_limit_reset
  41. end
  42. def rate_limit_reset
  43. (request_time + reset_period_offset).iso8601(6)
  44. end
  45. def api_throttle_data
  46. most_limited_type, = request.env['rack.attack.throttle_data'].min_by { |_key, value| value[:limit] - value[:count] }
  47. request.env['rack.attack.throttle_data'][most_limited_type]
  48. end
  49. def request_time
  50. @request_time ||= Time.now.utc
  51. end
  52. def reset_period_offset
  53. api_throttle_data[:period] - (request_time.to_i % api_throttle_data[:period])
  54. end
  55. end