From 514bb6b4af9012d52770068eac5f76b9bbef3900 Mon Sep 17 00:00:00 2001
From: Aleksey V Zapparov <ixti@member.fsf.org>
Date: Fri, 28 Mar 2014 04:05:48 +0100
Subject: [PATCH] Fix readpartial loop

---
 lib/http/client.rb       | 64 +++++++++++++++++-----------------------
 spec/http/client_spec.rb | 58 ++++++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 37 deletions(-)

diff --git a/lib/http/client.rb b/lib/http/client.rb
index d2bda4a5..9208f027 100644
--- a/lib/http/client.rb
+++ b/lib/http/client.rb
@@ -39,52 +39,22 @@ def perform(req, options)
 
       if options.follow
         res = Redirector.new(options.follow).perform req, res do |request|
-          # TODO: keep-alive
-          @parser.reset
-          finish_response
-
           perform_without_following_redirects request, options
         end
       end
 
-      @body_remaining = Integer(res['Content-Length']) if res['Content-Length']
       res
     end
 
     # Read a chunk of the body
-    def readpartial(size = BUFFER_SIZE) # rubocop:disable CyclomaticComplexity
-      if @parser.finished? || (@body_remaining && @body_remaining.zero?)
-        chunk = @parser.chunk
-
-        if !chunk && @body_remaining && !@body_remaining.zero?
-          fail StateError, "expected #{@body_remaining} more bytes of body"
-        end
-
-        @body_remaining -= chunk.bytesize if chunk
-        return chunk
-      end
-
-      fail StateError, 'not connected' unless @socket
+    def readpartial(size = BUFFER_SIZE)
+      return unless @socket
 
+      read_more size
       chunk = @parser.chunk
-      unless chunk
-        begin
-          @parser << @socket.readpartial(BUFFER_SIZE)
-          chunk = @parser.chunk
-        rescue EOFError
-          chunk = nil
-        end
-
-        # TODO: consult @body_remaining here and raise if appropriate
-        return unless chunk
-      end
-
-      if @body_remaining
-        @body_remaining -= chunk.bytesize
-        @body_remaining = nil if @body_remaining < 1
-      end
 
       finish_response if @parser.finished?
+
       chunk
     end
 
@@ -92,6 +62,12 @@ def readpartial(size = BUFFER_SIZE) # rubocop:disable CyclomaticComplexity
 
     # Perform a single (no follow) HTTP request
     def perform_without_following_redirects(req, options)
+      # finish previous response if client was re-used
+      # TODO: this is pretty wrong, as socket shoud be part of response
+      #       connection, so that re-use of client will not break multiple
+      #       chunked responses
+      finish_response
+
       uri = req.uri
 
       # TODO: keep-alive support
@@ -101,13 +77,17 @@ def perform_without_following_redirects(req, options)
       req.stream @socket
 
       begin
-        @parser << @socket.readpartial(BUFFER_SIZE) until @parser.headers
+        read_more BUFFER_SIZE until @parser.headers
       rescue IOError, Errno::ECONNRESET, Errno::EPIPE => ex
         raise IOError, "problem making HTTP request: #{ex}"
       end
 
       body = Response::Body.new(self)
-      Response.new(@parser.status_code, @parser.http_version, @parser.headers, body, uri)
+      res  = Response.new(@parser.status_code, @parser.http_version, @parser.headers, body, uri)
+
+      finish_response if :head == req.verb
+
+      res
     end
 
     # Initialize TLS connection
@@ -135,10 +115,20 @@ def make_request_body(opts, headers)
 
     # Callback for when we've reached the end of a response
     def finish_response
-      # TODO: keep-alive support
+      @socket.close if @socket && !@socket.closed?
+      @parser.reset
+
       @socket = nil
     end
 
+    # Feeds some more data into parser
+    def read_more(size)
+      @parser << @socket.readpartial(size) unless @parser.finished?
+      return true
+    rescue EOFError
+      return false
+    end
+
     # Moves uri get params into the opts.params hash
     # @return [Array<URI, Hash>]
     def normalize_get_params(uri, opts)
diff --git a/spec/http/client_spec.rb b/spec/http/client_spec.rb
index 7b2da6e4..9d4abcc1 100644
--- a/spec/http/client_spec.rb
+++ b/spec/http/client_spec.rb
@@ -120,4 +120,62 @@ def simple_response(body, status = 200)
       end
     end
   end
+
+  describe '#perform' do
+    let(:client) { described_class.new }
+
+    it 'calls finish_response before actual performance' do
+      TCPSocket.stub(:open) { throw :halt }
+      expect(client).to receive(:finish_response)
+      catch(:halt) { client.head "http://127.0.0.1:#{ExampleService::PORT}/" }
+    end
+
+    it 'calls finish_response once body was fully flushed' do
+      expect(client).to receive(:finish_response).twice.and_call_original
+      client.get("http://127.0.0.1:#{ExampleService::PORT}/").to_s
+    end
+
+    context 'with HEAD request' do
+      it 'does not iterates through body' do
+        expect(client).to_not receive(:readpartial)
+        client.head("http://127.0.0.1:#{ExampleService::PORT}/")
+      end
+
+      it 'finishes response after headers were received' do
+        expect(client).to receive(:finish_response).twice.and_call_original
+        client.head("http://127.0.0.1:#{ExampleService::PORT}/")
+      end
+    end
+
+    context 'when server fully flushes response in one chunk' do
+      before do
+        socket_spy = double
+
+        chunks = [
+          <<-RESPONSE.gsub(/^\s*\| */, '').gsub(/\n/, "\r\n")
+          | HTTP/1.1 200 OK
+          | Content-Type: text/html
+          | Server: WEBrick/1.3.1 (Ruby/1.9.3/2013-11-22)
+          | Date: Mon, 24 Mar 2014 00:32:22 GMT
+          | Content-Length: 15
+          | Connection: Keep-Alive
+          |
+          | <!doctype html>
+          RESPONSE
+        ]
+
+        socket_spy.stub(:close) { nil }
+        socket_spy.stub(:closed?) { true }
+        socket_spy.stub(:readpartial) { chunks.shift }
+        socket_spy.stub(:<<) { nil }
+
+        TCPSocket.stub(:open) { socket_spy }
+      end
+
+      it 'properly reads body' do
+        body = client.get("http://127.0.0.1:#{ExampleService::PORT}/").to_s
+        expect(body).to eq '<!doctype html>'
+      end
+    end
+  end
 end