Permalink
Browse files

correctly dechunk data from response and handle flushing

this is due rails `render stream: true` support (closes #117)

java servlet containers all seem to auto-chunk thus they expect non-chunked data otherwise it might end up double chunked !

inspired by code written by Tim Olsen (@tolsen)
  • Loading branch information...
1 parent 78d85a1 commit c06add59e1ee80446023a82184910f61199cc32f @kares kares committed Aug 30, 2012
Showing with 167 additions and 45 deletions.
  1. +55 −19 src/main/ruby/jruby/rack/response.rb
  2. +112 −26 src/spec/ruby/jruby/rack/response_spec.rb
@@ -9,10 +9,12 @@ module JRuby
module Rack
class Response
include org.jruby.rack.RackResponse
- java_import java.nio.channels.Channels
+ java_import 'java.nio.ByteBuffer'
+ java_import 'java.nio.channels.Channels'
- def initialize(arr)
- @status, @headers, @body = *arr
+ # Expects a Rack response: [status, headers, body].
+ def initialize(array)
+ @status, @headers, @body = *array
end
def getStatus
@@ -23,10 +25,6 @@ def getHeaders
@headers
end
- def chunked?
- (@headers && @headers['Transfer-Encoding'] == "chunked")
- end
-
def getBody
body = ""
@body.each { |part| body << part }
@@ -42,7 +40,7 @@ def respond(response)
write_body(response)
end
end
-
+
def write_status(response)
response.setStatus(@status.to_i)
end
@@ -54,9 +52,12 @@ def write_headers(response)
response.setContentType(v.to_s)
when /^Content-Length$/i
length = v.to_i
- # setContentLength accepts only int, addHeader must be used for large files (>2GB)
- response.setContentLength(length) unless chunked? || length >= 2_147_483_648
+ # setContentLength(int) ... addHeader must be used for large files (>2GB)
+ response.setContentLength(length) if ! chunked? && length < 2_147_483_648
else
+ # servlet container auto handle chunking when response is flushed
+ # (and Content-Length headers has not been set) :
+ next if ( k == 'Transfer-Encoding' && v == 'chunked' )
# NOTE: effectively the same as `v.split("\n").each` which is what
# rack handler does to guard against response splitting attacks !
# https://github.com/jruby/jruby-rack/issues/81
@@ -79,29 +80,46 @@ def write_headers(response)
end
def write_body(response)
- outputstream = response.getOutputStream
+ output_stream = response.getOutputStream
begin
if @body.respond_to?(:call) && ! @body.respond_to?(:each)
- @body.call(outputstream)
+ @body.call(output_stream)
elsif @body.respond_to?(:to_channel) &&
! object_polluted_with_anyio?(@body, :to_channel)
@body = @body.to_channel # so that we close the channel
- transfer_channel(@body, outputstream)
+ transfer_channel(@body, output_stream)
elsif @body.respond_to?(:to_inputstream) &&
! object_polluted_with_anyio?(@body, :to_inputstream)
@body = @body.to_inputstream # so that we close the stream
- transfer_channel(Channels.newChannel(@body), outputstream)
+ transfer_channel(Channels.newChannel(@body), output_stream)
elsif @body.respond_to?(:body_parts) && @body.body_parts.respond_to?(:to_channel) &&
! object_polluted_with_anyio?(@body.body_parts, :to_channel)
# ActionDispatch::Response "raw" body access in case it's a File
@body = @body.body_parts.to_channel # so that we close the channel
- transfer_channel(@body, outputstream)
+ transfer_channel(@body, output_stream)
else
# 1.8 has a String#each method but 1.9 does not :
method = @body.respond_to?(:each_line) ? :each_line : :each
- @body.send(method) do |line|
- outputstream.write(line.to_java_bytes)
- outputstream.flush if chunked?
+ if dechunk?
+ # NOTE: due Rails 3.2 stream-ed rendering http://git.io/ooCOtA#L223
+ term = "\r\n"; tail = "0#{term}#{term}".freeze
+ chunk = /([0-9]|[a-f]+)#{Regexp.escape(term)}(.*)#{Regexp.escape(term)}/mo
+ @body.send(method) do |line|
+ if line == tail
+ # NOOP
+ elsif line =~ chunk # (size.to_s(16)) term (chunk) term
+ # NOTE: not checking whether we have the correct size ?!
+ output_stream.write $2.to_java_bytes
+ else # seems it's not a chunk ... thus let it flow :
+ output_stream.write line.to_java_bytes
+ end
+ output_stream.flush
+ end
+ else
+ @body.send(method) do |line|
+ output_stream.write(line.to_java_bytes)
+ output_stream.flush if flush?
+ end
end
end
rescue LocalJumpError => e
@@ -116,6 +134,24 @@ def write_body(response)
end
end
+ protected
+
+ # returns true if a chunked encoding is detected
+ def chunked?
+ return @chunked unless @chunked.nil?
+ @chunked = !! ( @headers && @headers['Transfer-Encoding'] == 'chunked' )
+ end
+
+ # returns true if output (body) should be flushed after each written line
+ def flush?
+ chunked? || ! ( @headers && @headers['Content-Length'] )
+ end
+
+ # this should be true whenever the response should be de-chunked
+ def dechunk?
+ chunked?
+ end
+
private
BUFFER_SIZE = 16 * 1024
@@ -125,7 +161,7 @@ def transfer_channel(channel, outputstream)
if channel.respond_to?(:transfer_to)
channel.transfer_to(0, channel.size, outputchannel)
else
- buffer = java.nio.ByteBuffer.allocate(BUFFER_SIZE)
+ buffer = ByteBuffer.allocate(BUFFER_SIZE)
while channel.read(buffer) != -1
buffer.flip
outputchannel.write(buffer)
@@ -9,6 +9,7 @@
require 'jruby/rack/response'
describe JRuby::Rack::Response do
+
before :each do
@status, @headers, @body = mock("status"), mock("headers"), mock("body")
@headers.stub!(:[]).and_return nil
@@ -82,14 +83,6 @@ class << str; undef_method :each; end if str.respond_to?(:each)
@response.write_headers(@servlet_response)
end
- it "should detect a chunked response when the Transfer-Encoding header is set" do
- @headers = { "Transfer-Encoding" => "chunked" }
- @response = JRuby::Rack::Response.new([@status, @headers, @body])
- @servlet_response.should_receive(:addHeader).with("Transfer-Encoding", "chunked")
- @response.write_headers(@servlet_response)
- @response.chunked?.should eql(true)
- end
-
it "should write the status first, followed by the headers, and the body last" do
@servlet_response.should_receive(:committed?).and_return false
@response.should_receive(:write_status).ordered
@@ -112,39 +105,132 @@ class << str; undef_method :each; end if str.respond_to?(:each)
@response.getBody.should == "hello"
end
+ it "detects a chunked response when the Transfer-Encoding header is set" do
+ @headers = { "Transfer-Encoding" => "chunked" }
+ @response = JRuby::Rack::Response.new([@status, @headers, @body])
+ # NOTE: servlet container auto handle chunking when flushed no need to set :
+ @servlet_response.should_not_receive(:addHeader).with("Transfer-Encoding", "chunked")
+ @response.write_headers(@servlet_response)
+ @response.send(:chunked?).should be true
+ end
+
describe "#write_body" do
+
let(:stream) do
StubOutputStream.new.tap do |stream|
@servlet_response.stub!(:getOutputStream).and_return stream
end
end
-
- it "does not flush after write if Transfer-Encoding header is not set" do
- @body.should_receive(:each).
- and_yield("hello").
- and_yield("there")
- @servlet_response.should_not_receive(:addHeader).with("Transfer-Encoding", "chunked")
- @response.chunked?.should eql(false)
- stream.should_receive(:write).exactly(2).times
- stream.should_not_receive(:flush)
-
- @response.write_body(@servlet_response)
- end
-
+
it "writes the body to the stream and flushes when the response is chunked" do
@headers = { "Transfer-Encoding" => "chunked" }
@response = JRuby::Rack::Response.new([@status, @headers, @body])
- @servlet_response.should_receive(:addHeader).with("Transfer-Encoding", "chunked")
+ # NOTE: servlet container auto handle chunking when flushed no need to set :
+ @servlet_response.should_not_receive(:addHeader).with("Transfer-Encoding", "chunked")
@response.write_headers(@servlet_response)
- @response.chunked?.should eql(true)
- @body.should_receive(:each).ordered.
- and_yield("hello").
- and_yield("there")
+ @response.send(:chunked?).should == true
+ @body.should_receive(:each).ordered.and_yield("hello").and_yield("there")
stream.should_receive(:write).exactly(2).times
stream.should_receive(:flush).exactly(2).times
@response.write_body(@servlet_response)
end
+ it "dechunks the body when a chunked response is detected",
+ :lib => [ :rails31, :rails32, :rails40 ] do
+ require 'rack/chunked'
+
+ headers = {
+ "Cache-Control" => 'no-cache',
+ "Transfer-Encoding" => 'chunked'
+ }
+ body = [
+ "1".freeze,
+ "\nsecond chunk",
+ "a multi\nline chunk \n42",
+ "yet-another-chunk\n",
+ "terminated chunk\r\n",
+ "\r\nthe very\r\n last\r\n\r\n chunk",
+ ]
+ body = Rack::Chunked::Body.new body
+ response = JRuby::Rack::Response.new([ 200, headers, body ])
+ @servlet_response.stub!(:getOutputStream).and_return stream = mock("stream")
+ @servlet_response.stub!(:addHeader)
+ response.write_headers(@servlet_response)
+
+ times = 0
+ stream.should_receive(:write).exactly(6).times.with do |bytes|
+ str = String.from_java_bytes(bytes)
+ case times += 1
+ when 1 then str.should == "1"
+ when 2 then str.should == "\nsecond chunk"
+ when 3 then str.should == "a multi\nline chunk \n42"
+ when 4 then str.should == "yet-another-chunk\n"
+ when 5 then str.should == "terminated chunk\r\n"
+ when 6 then str.should == "\r\nthe very\r\n last\r\n\r\n chunk"
+ else
+ fail("unexpected :write received with #{str.inspect}")
+ end
+ end
+ stream.should_receive(:flush).exactly(6+1).times # +1 for tail chunk
+
+ response.write_body(@servlet_response)
+ end
+
+ it "handles dechunking gracefully when body is not chunked" do
+ headers = {
+ "Cache-Control" => 'no-cache',
+ "Transfer-Encoding" => 'chunked'
+ }
+ body = [
+ "1".freeze,
+ "a multi\nline chunk \n42",
+ "\r\nthe very\r\n last\r\n\r\n chunk",
+ ]
+ response = JRuby::Rack::Response.new([ 200, headers, body ])
+ @servlet_response.stub!(:getOutputStream).and_return stream = mock("stream")
+ @servlet_response.stub!(:addHeader)
+ response.write_headers(@servlet_response)
+
+ times = 0
+ stream.should_receive(:write).exactly(3).times.with do |bytes|
+ str = String.from_java_bytes(bytes)
+ case times += 1
+ when 1 then str.should == "1"
+ when 2 then str.should == "a multi\nline chunk \n42"
+ when 3 then str.should == "\r\nthe very\r\n last\r\n\r\n chunk"
+ else
+ fail("unexpected :write received with #{str.inspect}")
+ end
+ end
+ stream.should_receive(:flush).exactly(3).times
+
+ response.write_body(@servlet_response)
+ end
+
+ it "flushed the body when no Content-Length set" do
+ @response = JRuby::Rack::Response.new([ 200, {}, @body ])
+ @servlet_response.stub!(:addHeader)
+ @body.should_receive(:each).ordered.and_yield("hello").and_yield("there")
+ @response.write_headers(@servlet_response)
+ stream.should_receive(:write).once.ordered
+ stream.should_receive(:flush).once.ordered
+ stream.should_receive(:write).once.ordered
+ stream.should_receive(:flush).once.ordered
+ @response.write_body(@servlet_response)
+ end
+
+ it "does not flush the body when Content-Length set" do
+ @headers = { "Content-Length" => 10 }
+ @response = JRuby::Rack::Response.new([ 200, @headers, @body ])
+ @servlet_response.stub!(:addHeader)
+ @servlet_response.stub!(:setContentLength)
+ @body.should_receive(:each).ordered.and_yield("hello").and_yield("there")
+ @response.write_headers(@servlet_response)
+ stream.should_receive(:write).twice
+ stream.should_receive(:flush).never
+ @response.write_body(@servlet_response)
+ end
+
it "writes the body to the servlet response" do
@body.should_receive(:each).
and_yield("hello").

0 comments on commit c06add5

Please sign in to comment.