rack-attack の throttle の実装がどうなっているかコード読んでみた

こんにちは。中華スープに入っている揚げ卵が好きです、きたけーです。

先日、ブログで紹介したrack-attack ( rack-attack middleware を使って IPアドレスでアクセスを制限したり、同一IPアドレスからの大量のリクエストを防ぐ術 - きたけーTechブログ) ですが、実装どうなってんのかなぁと思ってコードを読んでみました(特にthrottleを中心に)。

設定

大本はここ rack-attack/attack.rb at 25739e24dbc820898cf45d120b1e794e8f6f4df1 · kickstarter/rack-attack · GitHub

前回、ブログで紹介した以下の様な記述では、

Rack::Attack.throttle('reports limit', limit: 1, period: 60.seconds) do |req|
  req.ip if req.post? && req.path.match(%r{^/reports})
end

以下のように Throttle のオブジェクトを生成してハッシュに格納しているのが分かる (https://github.com/kickstarter/rack-attack/blob/25739e24dbc820898cf45d120b1e794e8f6f4df1/lib/rack/attack.rb#L30)

def throttle(name, options, &block)
  self.throttles[name] = Throttle.new(name, options, block)
end

リクエストがきたら

リクエストがきたら、throttlesハッシュに格納されているオブジェクトひとつひとつに対して[]メソッドをリクエストオブジェクトを引数にして呼び出す(https://github.com/kickstarter/rack-attack/blob/25739e24dbc820898cf45d120b1e794e8f6f4df1/lib/rack/attack.rb#L55)。

def throttled?(req)
  throttles.any? do |name, throttle|
    throttle[req]
  end
end

ここでもし、ひとつでも真を返すものがあれば、ステータスコード「429 Too Many Requests」のレスポンスを返す(https://github.com/kickstarter/rack-attack/blob/25739e24dbc820898cf45d120b1e794e8f6f4df1/lib/rack/attack.rb#L84)。

  @throttled_response   = lambda {|env|
    retry_after = (env['rack.attack.match_data'] || {})[:period]
    [429, {'Content-Type' => 'text/plain', 'Retry-After' => retry_after.to_s}, ["Retry later\n"]]
  }

Rack::Attack::Throttle#[]に潜る

実装はこんなかんじ(https://github.com/kickstarter/rack-attack/blob/25739e24dbc820898cf45d120b1e794e8f6f4df1/lib/rack/attack/throttle.rb#L20)

      def [](req)
        discriminator = block[req]
        return false unless discriminator

        current_period = period.respond_to?(:call) ? period.call(req) : period
        current_limit  = limit.respond_to?(:call) ? limit.call(req) : limit
        key            = "#{name}:#{discriminator}"
        count          = cache.count(key, current_period)

        data = {
          :count => count,
          :period => current_period,
          :limit => current_limit
        }
        (req.env['rack.attack.throttle_data'] ||= {})[name] = data

        (count > current_limit).tap do |throttled|
          if throttled
            req.env['rack.attack.matched']             = name
            req.env['rack.attack.match_discriminator'] = discriminator
            req.env['rack.attack.match_type']          = type
            req.env['rack.attack.match_data']          = data
            Rack::Attack.instrument(req)
          end
        end
      end

discriminatorには、記事の最初で記述した設定のブロックをProcオブジェクトとして評価している。対象のリクエストであればIPアドレスがくる。対象のリクエストでなければ、nilがくるのでメソッドとしてはfalseを返す。

その後は、periodとlimitがProcオブジェクトらしかったら評価して(ダックタイピングだ！)、らしくなかったら、そのまま返す。periodとlimitはProcオブジェクトが指定できるんですね。(勉強になった)

Rack::Attack::Cache

cacheは Rack::Attack::Cacheというクラスのインスタンスで、こいつはデフォルトだとRails Cacheをラップして、リクエストの回数をカウントしたりしている。ここが肝っぽい(https://github.com/kickstarter/rack-attack/blob/25739e24dbc820898cf45d120b1e794e8f6f4df1/lib/rack/attack/cache.rb#L17)

def count(unprefixed_key, period)
  epoch_time = Time.now.to_i
  # Add 1 to expires_in to avoid timing error: http://git.io/i1PHXA
  expires_in = period - (epoch_time % period) + 1
  key = "#{prefix}:#{(epoch_time/period).to_i}:#{unprefixed_key}"
  do_count(key, expires_in)
end

epoch時間と設定した時間間隔の商をつかってキャッシュキーを生成したり、余りから有効期限を設定している。

Rack::Attack::Throttle#[]に戻る

最後はenvに色々な値を設定したりなど。tapつかう必要あるのかな、と思ったんですけど、メソッドとしてリミットを超えたかどうか真偽値を返す必要があるからでした。

まとめ

何気なく読み始めたんですけど、素直に実装されていて、かつ、コードもきれいなので読みやすかったです。 1時間くらいの社内勉強会で、良質なrack middlewareのコードリーディング会をやるといいかもしれないですね。

最近、お仕事でRubyに触れる機会が減ってきているので、積極的にコード読んだり、Railsのコード書いたりしよう。