From 1e6d3d19da2b62bfe7f8fd7827dcad3ee3fe9923 Mon Sep 17 00:00:00 2001
From: Eric Wong <normalperson@yhbt.net>
Date: Mon, 19 Jul 2010 10:09:56 +0000
Subject: rev + em: more easily allow Content-Length in pipe responses

If a response proxying a pipe (or socket) includes a
Content-Length, do not attempt to outsmart the application
and just use the given Content-Length.

This helps avoid exposing applications to weird internals such
as env["rainbows.autochunk"] and X-Rainbows-* response headers.
---
 lib/rainbows/dev_fd_response.rb | 11 ++++++-----
 lib/rainbows/ev_core.rb         | 14 ++++++++++++++
 lib/rainbows/event_machine.rb   | 10 +---------
 lib/rainbows/rev/client.rb      | 10 +---------
 t/fast-pipe-response.ru         |  2 --
 5 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/lib/rainbows/dev_fd_response.rb b/lib/rainbows/dev_fd_response.rb
index c04d64f..451cad7 100644
--- a/lib/rainbows/dev_fd_response.rb
+++ b/lib/rainbows/dev_fd_response.rb
@@ -40,11 +40,12 @@ class Rainbows::DevFdResponse < Struct.new(:app)
       headers['Content-Length'] ||= st.size.to_s
       headers.delete('Transfer-Encoding')
     elsif st.pipe? || st.socket? # epoll-able things
-      if env['rainbows.autochunk']
-        headers['Transfer-Encoding'] = 'chunked'
-        headers.delete('Content-Length')
-      else
-        headers['X-Rainbows-Autochunk'] = 'no'
+      unless headers['Content-Length']
+        if env['rainbows.autochunk']
+          headers['Transfer-Encoding'] = 'chunked'
+        else
+          headers['X-Rainbows-Autochunk'] = 'no'
+        end
       end
 
       # we need to make sure our pipe output is Fiber-compatible
diff --git a/lib/rainbows/ev_core.rb b/lib/rainbows/ev_core.rb
index 5ca693b..dbcdeba 100644
--- a/lib/rainbows/ev_core.rb
+++ b/lib/rainbows/ev_core.rb
@@ -6,6 +6,7 @@ module Rainbows
   module EvCore
     include Unicorn
     include Rainbows::Const
+    include Rainbows::Response
     G = Rainbows::G
     NULL_IO = Unicorn::HttpRequest::NULL_IO
 
@@ -33,6 +34,19 @@ module Rainbows
         quit
     end
 
+    # returns whether to enable response chunking for autochunk models
+    def stream_response_headers(status, headers)
+      if headers['Content-Length']
+        rv = false
+      else
+        rv = !!(headers['Transfer-Encoding'] =~ %r{\Achunked\z}i)
+        rv = false if headers.delete('X-Rainbows-Autochunk') == 'no'
+      end
+      headers[CONNECTION] = CLOSE # TODO: allow keep-alive
+      write(response_header(status, headers))
+      rv
+    end
+
     # TeeInput doesn't map too well to this right now...
     def on_read(data)
       case @state
diff --git a/lib/rainbows/event_machine.rb b/lib/rainbows/event_machine.rb
index 625357e..86cb4eb 100644
--- a/lib/rainbows/event_machine.rb
+++ b/lib/rainbows/event_machine.rb
@@ -50,7 +50,6 @@ module Rainbows
 
     class Client < EM::Connection # :nodoc: all
       include Rainbows::EvCore
-      include Rainbows::Response
       G = Rainbows::G
 
       def initialize(io)
@@ -98,14 +97,7 @@ module Rainbows
 
       # used for streaming sockets and pipes
       def stream_response(status, headers, io)
-        if headers
-          do_chunk = !!(headers['Transfer-Encoding'] =~ %r{\Achunked\z}i)
-          do_chunk = false if headers.delete('X-Rainbows-Autochunk') == 'no'
-          headers[CONNECTION] = CLOSE # TODO: allow keep-alive
-          write(response_header(status, headers))
-        else
-          do_chunk = false
-        end
+        do_chunk = stream_response_headers(status, headers) if headers
         mod = do_chunk ? ResponseChunkPipe : ResponsePipe
         EM.watch(io, mod, self).notify_readable = true
       end
diff --git a/lib/rainbows/rev/client.rb b/lib/rainbows/rev/client.rb
index bc3785f..f067d1b 100644
--- a/lib/rainbows/rev/client.rb
+++ b/lib/rainbows/rev/client.rb
@@ -7,7 +7,6 @@ module Rainbows
     class Client < ::Rev::IO
       include Rainbows::ByteSlice
       include Rainbows::EvCore
-      include Rainbows::Response
       G = Rainbows::G
 
       def initialize(io)
@@ -59,14 +58,7 @@ module Rainbows
 
       # used for streaming sockets and pipes
       def stream_response(status, headers, io, body)
-        if headers
-          do_chunk = !!(headers['Transfer-Encoding'] =~ %r{\Achunked\z}i)
-          do_chunk = false if headers.delete('X-Rainbows-Autochunk') == 'no'
-          headers[CONNECTION] = CLOSE # TODO: allow keep-alive
-          write(response_header(status, headers))
-        else
-          do_chunk = false
-        end
+        do_chunk = stream_response_headers(status, headers) if headers
         # we only want to attach to the Rev::Loop belonging to the
         # main thread in Ruby 1.9
         io = DeferredResponse.new(io, self, do_chunk, body)
diff --git a/t/fast-pipe-response.ru b/t/fast-pipe-response.ru
index 01f4d59..77c8d71 100644
--- a/t/fast-pipe-response.ru
+++ b/t/fast-pipe-response.ru
@@ -1,10 +1,8 @@
 # must be run without Rack::Lint since that clobbers to_path
 use Rainbows::DevFdResponse
 run(lambda { |env|
-  env['rainbows.autochunk'] = false
   [ 200,
     {
-      'X-Rainbows-Autochunk' => 'no',
       'Content-Length' => ::File.stat('random_blob').size.to_s,
       'Content-Type' => 'application/octet-stream',
     },
-- 
cgit v1.2.3-24-ge0c7