1 files changed, 134 insertions, 0 deletions
diff --git a/web/server/h2o/libh2o/share/h2o/mruby/dos_detector.rb b/web/server/h2o/libh2o/share/h2o/mruby/dos_detector.rb
new file mode 100644
index 00000000..15de2bc7
--- /dev/null
+++ b/web/server/h2o/libh2o/share/h2o/mruby/dos_detector.rb
@@ -0,0 +1,134 @@
+# Copyright (c) 2016 DeNA Co., Ltd., Ichito Nagata
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+#
+require File.expand_path(File.dirname(__FILE__)) + "/lru_cache.rb"
+
+class DoSDetector
+
+  def initialize(config={})
+    config = {
+      :strategy   => CountingStrategy.new,
+      :callback   => self.class.default_callback,
+      :forwarded  => true,
+      :cache_size => 128,
+    }.merge(config)
+
+    @strategy = config[:strategy]
+    @callback = config[:callback]
+    @forwarded = !!(config[:forwarded])
+    @cache = LRUCache.new(config[:cache_size])
+    raise "strategy must not be nil" if @strategy.nil?
+    raise "callback must not be nil" if @callback.nil?
+  end
+
+  def self.default_callback
+    Proc.new do |env, detected, ip|
+      if detected
+        [ 403, { "Content-Type" => "text/plain" }, [ "Forbidden" ] ]
+      else
+        [ 399, {}, [] ]
+      end
+    end
+  end
+
+  def self.fallthrough_callback
+    Proc.new do |env, detected, ip, vars|
+      if detected
+        vars ||= {}
+        env_headers = vars.merge({:ip => ip}).map { |k, v| [ "x-fallthru-set-dos-#{k}", v.to_s ] }.to_h
+        [ 399, env_headers, [] ]
+      else
+        [ 399, {}, [] ]
+      end
+    end
+  end
+
+  def call(env)
+    now = Time.now.to_i
+
+    ip = env['REMOTE_ADDR']
+    if @forwarded && (xff = env['HTTP_X_FORWARDED_FOR'])
+      ip = xff.split(",")[0]
+    end
+
+    unless client = @cache.get(ip)
+      client = { :ip => ip }
+      @cache.set(ip, client)
+    end
+
+    detected, *args = @strategy.detect?(client, now, env)
+    return @callback.call(env, detected, ip, *args)
+  end
+
+  class CountingStrategy
+
+    def initialize(config={})
+      config = {
+        :period     => 10,
+        :threshold  => 100,
+        :ban_period => 300,
+      }.merge(config)
+      @period = config[:period]
+      @threshold = config[:threshold]
+      @ban_period = config[:ban_period]
+      raise "period must be greater than zero" if @period <= 0
+      raise "threshold must be greater than zero" if @threshold <= 0
+      raise "ban_period must not be negative" if @ban_period < 0
+    end
+
+    def detect?(client, now, env)
+      count = countup(client, now)
+
+      banned_until = client[:banned_until] || 0
+      if banned_until >= now
+        detected = true
+      else
+        detected = count >= @threshold 
+        if detected
+          banned_until = now + @ban_period
+          client[:banned_until] = banned_until
+        end
+      end
+
+      return detected, { :count => count, :banned_until => banned_until }
+    end
+
+    private
+
+    def countup(client, now)
+      count = client[:count] || 0
+      period_index = client[:period_index] || 0
+
+      current_period_index = (now / @period).floor
+      if current_period_index > period_index
+        count -= (current_period_index - period_index) * @threshold
+        count = 0 if count < 0
+        client[:period_index] = current_period_index
+      end
+
+      count += 1
+      client[:count] = count
+
+      return count
+    end
+
+  end
+
+end