From d6a89ac03f8e9628ab2e71d8fb56fdbdaec16569 Mon Sep 17 00:00:00 2001
From: Ryan Tomayko <rtomayko@gmail.com>
Date: Thu, 5 Mar 2009 21:20:17 -0800
Subject: [PATCH 2/4] Add Rack::Chunked (Transfer-Encoding: chunked) middleware

---
 lib/rack.rb               |    1 +
 lib/rack/chunked.rb       |   49 +++++++++++++++++++++++++++++++++++
 test/spec_rack_chunked.rb |   62 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 112 insertions(+), 0 deletions(-)
 create mode 100644 lib/rack/chunked.rb
 create mode 100644 test/spec_rack_chunked.rb

diff --git a/lib/rack.rb b/lib/rack.rb
index 2643630..987744f 100644
--- a/lib/rack.rb
+++ b/lib/rack.rb
@@ -28,6 +28,7 @@ module Rack
 
   autoload :Builder, "rack/builder"
   autoload :Cascade, "rack/cascade"
+  autoload :Chunked, "rack/chunked"
   autoload :CommonLogger, "rack/commonlogger"
   autoload :ConditionalGet, "rack/conditionalget"
   autoload :ContentLength, "rack/content_length"
diff --git a/lib/rack/chunked.rb b/lib/rack/chunked.rb
new file mode 100644
index 0000000..280d89d
--- /dev/null
+++ b/lib/rack/chunked.rb
@@ -0,0 +1,49 @@
+require 'rack/utils'
+
+module Rack
+
+  # Middleware that applies chunked transfer encoding to response bodies
+  # when the response does not include a Content-Length header.
+  class Chunked
+    include Rack::Utils
+
+    def initialize(app)
+      @app = app
+    end
+
+    def call(env)
+      status, headers, body = @app.call(env)
+      headers = HeaderHash.new(headers)
+
+      if env['HTTP_VERSION'] == 'HTTP/1.0' ||
+         STATUS_WITH_NO_ENTITY_BODY.include?(status) ||
+         headers['Content-Length'] ||
+         headers['Transfer-Encoding']
+        [status, headers.to_hash, body]
+      else
+        dup.chunk(status, headers, body)
+      end
+    end
+
+    def chunk(status, headers, body)
+      @body = body
+      headers.delete('Content-Length')
+      headers['Transfer-Encoding'] = 'chunked'
+      [status, headers.to_hash, self]
+    end
+
+    def each
+      term = "\r\n"
+      @body.each do |chunk|
+        size = bytesize(chunk)
+        next if size == 0
+        yield [size.to_s(16), term, chunk, term].join
+      end
+      yield ["0", term, "", term].join
+    end
+
+    def close
+      @body.close if @body.respond_to?(:close)
+    end
+  end
+end
diff --git a/test/spec_rack_chunked.rb b/test/spec_rack_chunked.rb
new file mode 100644
index 0000000..39eea48
--- /dev/null
+++ b/test/spec_rack_chunked.rb
@@ -0,0 +1,62 @@
+require 'rack/mock'
+require 'rack/chunked'
+require 'rack/utils'
+
+context "Rack::Chunked" do
+
+  before do
+    @env = Rack::MockRequest.
+      env_for('/', 'HTTP_VERSION' => '1.1', 'REQUEST_METHOD' => 'GET')
+  end
+
+  specify 'chunks responses with no Content-Length' do
+    app = lambda { |env| [200, {}, ['Hello', ' ', 'World!']] }
+    response = Rack::MockResponse.new(*Rack::Chunked.new(app).call(@env))
+    response.headers.should.not.include 'Content-Length'
+    response.headers['Transfer-Encoding'].should.equal 'chunked'
+    response.body.should.equal "5\r\nHello\r\n1\r\n \r\n6\r\nWorld!\r\n0\r\n\r\n"
+  end
+
+  specify 'chunks empty bodies properly' do
+    app = lambda { |env| [200, {}, []] }
+    response = Rack::MockResponse.new(*Rack::Chunked.new(app).call(@env))
+    response.headers.should.not.include 'Content-Length'
+    response.headers['Transfer-Encoding'].should.equal 'chunked'
+    response.body.should.equal "0\r\n\r\n"
+  end
+
+  specify 'does not modify response when Content-Length header present' do
+    app = lambda { |env| [200, {'Content-Length'=>'12'}, ['Hello', ' ', 'World!']] }
+    status, headers, body = Rack::Chunked.new(app).call(@env)
+    status.should.equal 200
+    headers.should.not.include 'Transfer-Encoding'
+    headers.should.include 'Content-Length'
+    body.join.should.equal 'Hello World!'
+  end
+
+  specify 'does not modify response when client is HTTP/1.0' do
+    app = lambda { |env| [200, {}, ['Hello', ' ', 'World!']] }
+    @env['HTTP_VERSION'] = 'HTTP/1.0'
+    status, headers, body = Rack::Chunked.new(app).call(@env)
+    status.should.equal 200
+    headers.should.not.include 'Transfer-Encoding'
+    body.join.should.equal 'Hello World!'
+  end
+
+  specify 'does not modify response when Transfer-Encoding header already present' do
+    app = lambda { |env| [200, {'Transfer-Encoding' => 'identity'}, ['Hello', ' ', 'World!']] }
+    status, headers, body = Rack::Chunked.new(app).call(@env)
+    status.should.equal 200
+    headers['Transfer-Encoding'].should.equal 'identity'
+    body.join.should.equal 'Hello World!'
+  end
+
+  [100, 204, 304].each do |status_code|
+    specify "does not modify response when status code is #{status_code}" do
+      app = lambda { |env| [status_code, {}, []] }
+      status, headers, body = Rack::Chunked.new(app).call(@env)
+      status.should.equal status_code
+      headers.should.not.include 'Transfer-Encoding'
+    end
+  end
+end
-- 
1.6.2

