From 46d79be0ad3de48ef0a677537becb3508ccad31e Mon Sep 17 00:00:00 2001
From: Eric Wong <normalperson@yhbt.net>
Date: Tue, 28 Dec 2010 01:14:43 +0000
Subject: enable the keepalive_requests config option

This will allow servers to limit the number of keepalive
requests that can be made over a single connection to
prevent denial-of-service and also to improve fairness
in load-balancing.
---
 lib/rainbows/configurator.rb          |  9 +++++++
 lib/rainbows/http_server.rb           |  8 +++++-
 t/t0040-keepalive_requests-setting.sh | 51 +++++++++++++++++++++++++++++++++++
 3 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100755 t/t0040-keepalive_requests-setting.sh

diff --git a/lib/rainbows/configurator.rb b/lib/rainbows/configurator.rb
index e69a3fb..3203c5a 100644
--- a/lib/rainbows/configurator.rb
+++ b/lib/rainbows/configurator.rb
@@ -13,6 +13,7 @@ module Rainbows::Configurator
   #     worker_connections 400
   #     keepalive_timeout 0 # zero disables keepalives entirely
   #     client_max_body_size 5*1024*1024 # 5 megabytes
+  #     keepalive_requests 666 # default:100
   #   end
   #
   #   # the rest of the Unicorn configuration
@@ -33,6 +34,14 @@ module Rainbows::Configurator
   # The default +client_max_body_size+ is 1 megabyte (1024 * 1024 bytes),
   # setting this to +nil+ will disable body size checks and allow any
   # size to be specified.
+  #
+  # The default +keepalive_requests+ is 100, meaning a client may
+  # complete 100 keepalive requests after the initial request before
+  # \Rainbows! forces a disconnect.  Lowering this can improve
+  # load-balancing characteristics as it forces HTTP/1.1 clients to
+  # reconnect after the specified number of requests, hopefully to a
+  # less busy host or worker process.  This may also be used to mitigate
+  # denial-of-service attacks that use HTTP pipelining.
   def Rainbows!(&block)
     block_given? or raise ArgumentError, "Rainbows! requires a block"
     Rainbows::HttpServer.setup(block)
diff --git a/lib/rainbows/http_server.rb b/lib/rainbows/http_server.rb
index d02af72..906fa0a 100644
--- a/lib/rainbows/http_server.rb
+++ b/lib/rainbows/http_server.rb
@@ -90,10 +90,16 @@ class Rainbows::HttpServer < Unicorn::HttpServer
 
   def keepalive_timeout(nr)
     (Integer === nr && nr >= 0) or
-      raise ArgumentError, "keepalive must be a non-negative Integer"
+      raise ArgumentError, "keepalive_timeout must be a non-negative Integer"
     G.kato = nr
   end
 
+  def keepalive_requests(nr)
+    Integer === nr or
+      raise ArgumentError, "keepalive_requests must be a non-negative Integer"
+    Unicorn::HttpRequest.keepalive_requests = nr
+  end
+
   def client_max_body_size(nr)
     err = "client_max_body_size must be nil or a non-negative Integer"
     case nr
diff --git a/t/t0040-keepalive_requests-setting.sh b/t/t0040-keepalive_requests-setting.sh
new file mode 100755
index 0000000..aee6cd3
--- /dev/null
+++ b/t/t0040-keepalive_requests-setting.sh
@@ -0,0 +1,51 @@
+#!/bin/sh
+. ./test-lib.sh
+t_plan 6 "keepalive_requests limit tests for $model"
+
+t_begin "setup and start" && {
+	rainbows_setup $model 50 666
+	rtmpfiles curl_out curl_err
+	grep 'keepalive_timeout 666' $unicorn_config
+        ed -s $unicorn_config <<EOF
+,s/listen.*/&, :tcp_nodelay => true/
+w
+EOF
+	grep nodelay $unicorn_config
+	rainbows -E none -D env.ru -c $unicorn_config
+	rainbows_wait_start
+}
+
+t_begin "curl requests hit default keepalive_requests limit" && {
+	curl -sSfv http://$listen/[0-101] > $curl_out 2> $curl_err
+	test 1 -eq $(grep 'Connection: close' $curl_err |wc -l)
+	test 101 -eq $(grep 'Connection: keep-alive' $curl_err |wc -l)
+}
+
+t_begin "reload with smaller keepalive_requests limit" && {
+	ed -s $unicorn_config <<EOF
+,g/Rainbows!/
+a
+  keepalive_requests 5
+.
+w
+EOF
+	kill -HUP $rainbows_pid
+	test x"$(cat $fifo)" = xSTART
+}
+
+t_begin "curl requests hit smaller keepalive_requests limit" && {
+	rm -f $curl_out $curl_err
+	curl -sSfv http://$listen/[1-13] > $curl_out 2> $curl_err
+	test 2 -eq $(grep 'Connection: close' $curl_err |wc -l)
+	test 11 -eq $(grep 'Connection: keep-alive' $curl_err |wc -l)
+}
+
+t_begin "killing succeeds" && {
+	kill $rainbows_pid
+}
+
+t_begin "check stderr" && {
+	check_stderr
+}
+
+t_done
-- 
cgit v1.2.3-24-ge0c7