Class: RateLimiting

Inherits:
Object
  • Object
show all
Defined in:
lib/rate_limiting.rb

Instance Method Summary collapse

Constructor Details

#initialize(app, &block) ⇒ RateLimiting

Returns a new instance of RateLimiting.



6
7
8
9
10
11
12
 
# File 'lib/rate_limiting.rb', line 6

def initialize(app, &block)
  @app = app
  @logger =  nil
  @rules = []
  @cache = {}
  block.call(self)
end

Instance Method Details

#allowed?(request) ⇒ Boolean

Returns:

  • (Boolean) 


96
97
98
99
100
101
102
103
 
# File 'lib/rate_limiting.rb', line 96

def allowed?(request)
  if rule = find_matching_rule(request)
    logger.debug "[#{self}] #{request.ip}:#{request.path}: Rate limiting rule matched."
    apply_rule(request, rule)
  else
    true
  end
end

#apply_rule(request, rule) ⇒ Object 



112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
 
# File 'lib/rate_limiting.rb', line 112

def apply_rule(request, rule)
  key = rule.get_key(request)
  if cache_has?(key)
    record = cache_get(key)
    logger.debug "[#{self}] #{request.ip}:#{request.path}: Rate limiting entry: '#{key}' => #{record}"
    if (reset = Time.at(record.split(':')[1].to_i)) > Time.now
      # rule hasn't been reset yet
      times = record.split(':')[0].to_i
      cache_set(key, "#{times + 1}:#{reset.to_i}")
      if (times) < rule.limit
        # within rate limit
        response = get_header(times + 1, reset, rule.limit)
      else
        logger.debug "[#{self}] #{request.ip}:#{request.path}: Rate limited; request rejected."
        return false
      end
    else
      response = get_header(1, rule.get_expiration, rule.limit)
      cache_set(key, "1:#{rule.get_expiration.to_i}")
    end
  else
    response = get_header(1, rule.get_expiration, rule.limit)
    cache_set(key, "1:#{rule.get_expiration.to_i}")
  end
  response
end

#cache ⇒ Object 



43
44
45
46
47
48
 
# File 'lib/rate_limiting.rb', line 43

def cache
  case @cache
    when Proc then @cache.call
    else @cache
  end
end

#cache_get(key) ⇒ Object 



62
63
64
65
66
67
68
69
70
71
 
# File 'lib/rate_limiting.rb', line 62

def cache_get(key)
  case
  when cache.respond_to?(:[])
    return cache[key]
  when cache.respond_to?(:get)
    return cache.get(key) || nil
  when cache.respond_to?(:fetch)
    return cache.fetch(key)
  end
end

#cache_has?(key) ⇒ Boolean

Returns:

  • (Boolean) 


50
51
52
53
54
55
56
57
58
59
60
 
# File 'lib/rate_limiting.rb', line 50

def cache_has?(key)
  case
  when cache.respond_to?(:has_key?)
    cache.has_key?(key)
  when cache.respond_to?(:get)
    cache.get(key) rescue false
  when cache.respond_to?(:exist?)
    cache.exist?(key)
  else false
  end
end

#cache_set(key, value) ⇒ Object 



73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
 
# File 'lib/rate_limiting.rb', line 73

def cache_set(key, value)
  case
  when cache.respond_to?(:[])
    begin
      cache[key] = value
    rescue TypeError => e
      cache[key] = value.to_s
    end
  when cache.respond_to?(:set)
    cache.set(key, value)
  when cache.respond_to?(:write)
    begin
      cache.write(key, value)
    rescue TypeError => e
      cache.write(key, value.to_s)
    end
  end
end

#call(env) ⇒ Object 



14
15
16
17
18
 
# File 'lib/rate_limiting.rb', line 14

def call(env)
  request = Rack::Request.new(env)
  @logger = env['rack.logger']
  (limit_header = allowed?(request)) ? respond(env, limit_header) : rate_limit_exceeded(env['HTTP_ACCEPT'])
end

#define_rule(options) ⇒ Object 



35
36
37
 
# File 'lib/rate_limiting.rb', line 35

def define_rule(options)
  @rules << Rule.new(options)
end

#find_matching_rule(request) ⇒ Object 



105
106
107
108
109
110
 
# File 'lib/rate_limiting.rb', line 105

def find_matching_rule(request)
  @rules.each do |rule|
    return rule if request.path =~ rule.match
  end
  nil
end

#get_header(times, reset, limit) ⇒ Object 



139
140
141
 
# File 'lib/rate_limiting.rb', line 139

def get_header(times, reset, limit)
  {'x-RateLimit-Limit' => limit.to_s, 'x-RateLimit-Remaining' => (limit - times).to_s, 'x-RateLimit-Reset' => reset.strftime("%d%m%y%H%M%S") }
end

#logger ⇒ Object 



92
93
94
 
# File 'lib/rate_limiting.rb', line 92

def logger
  @logger || Rack::NullLogger.new(nil)
end

#rate_limit_exceeded(accept) ⇒ Object 



25
26
27
28
29
30
31
32
33
 
# File 'lib/rate_limiting.rb', line 25

def rate_limit_exceeded(accept)
  case accept.gsub(/;.*/, "").split(',')[0]
  when "text/xml"         then message, type  = xml_error("403", "Rate Limit Exceeded"), "text/xml"
  when "application/json" then  message, type  = ["Rate Limit Exceeded"].to_json, "application/json"
  else
    message, type  = ["Rate Limit Exceeded"], "text/html"
  end
  [403, {"Content-Type" => type}, message]
end

#respond(env, limit_header) ⇒ Object 



20
21
22
23
 
# File 'lib/rate_limiting.rb', line 20

def respond(env, limit_header)
  status, header, response = @app.call(env)
  (limit_header.class == Hash) ? [status, header.merge(limit_header), response] : [status, header, response]
end

#set_cache(cache) ⇒ Object 



39
40
41
 
# File 'lib/rate_limiting.rb', line 39

def set_cache(cache)
  @cache = cache
end

#xml_error(code, message) ⇒ Object 



143
144
145
 
# File 'lib/rate_limiting.rb', line 143

def xml_error(code, message)
  "<?xml version=\"1.0\"?>\n<error>\n  <code>#{code}</code>\n  <message>#{message}</message>\n</error>"
end