Move streaming requests to generic HTTP class

4 files changed, 112 insertions, 106 deletions

diff --git a/lib/chef/http.rb b/lib/chef/http.rb
index 94d9d4d145..77faa686a0 100644
--- a/lib/chef/http.rb
+++ b/lib/chef/http.rb
@@ -35,6 +35,29 @@ class Chef
   # Basic HTTP client, with support for adding features via middleware
   class HTTP
 
+    # Class for applying middleware behaviors to streaming
+    # responses. Collects stream handlers (if any) from each
+    # middleware. When #handle_chunk is called, the chunk gets
+    # passed to all handlers in turn for processing.
+    class StreamHandler
+      def initialize(middlewares, response)
+        middlewares = middlewares.flatten
+        @stream_handlers = []
+        middlewares.each do |middleware|
+          stream_handler = middleware.stream_response_handler(response)
+          @stream_handlers << stream_handler unless stream_handler.nil?
+        end
+      end
+
+      def handle_chunk(next_chunk)
+        @stream_handlers.inject(next_chunk) do |chunk, handler|
+          handler.handle_chunk(chunk)
+        end
+      end
+
+    end
+
+
     def self.middlewares
       @middlewares ||= []
     end
@@ -106,14 +129,8 @@ class Chef
       api_request(:DELETE, create_url(path), headers)
     end
 
-    def create_url(path)
-      if path =~ /^(http|https):\/\//
-        URI.parse(path)
-      else
-        URI.parse("#{@url}/#{path}")
-      end
-    end
-
+    # Makes an HTTP request to +url+ with the given +method+, +headers+, and
+    # +data+ (if applicable).
     def request(method, url, headers={}, data=false)
 
       method, url, headers, data = apply_request_middleware(method, url, headers, data)
@@ -131,6 +148,57 @@ class Chef
       raise
     end
 
+    # Makes a streaming download request, streaming the response body to a
+    # tempfile. If a block is given, the tempfile is passed to the block and
+    # the tempfile will automatically be unlinked after the block is executed.
+    #
+    # If no block is given, the tempfile is returned, which means it's up to
+    # you to unlink the tempfile when you're done with it.
+    def streaming_request(url, headers, &block)
+      response, rest_request, return_value = nil, nil, nil
+      tempfile = nil
+
+      method = :GET
+      method, url, headers, data = apply_request_middleware(method, url, headers, data)
+
+      response, rest_request, return_value = send_http_request(method, url, headers, data) do |http_response|
+        if http_response.kind_of?(Net::HTTPSuccess)
+          tempfile = stream_to_tempfile(url, http_response)
+          if block_given?
+            begin
+              yield tempfile
+            ensure
+              tempfile && tempfile.close!
+            end
+          end
+        end
+      end
+      unless response.kind_of?(Net::HTTPSuccess) or response.kind_of?(Net::HTTPRedirection)
+        response.error!
+      end
+      tempfile
+    rescue Exception => e
+      log_failed_request(response, return_value) unless response.nil?
+      if e.respond_to?(:chef_rest_request=)
+        e.chef_rest_request = rest_request
+      end
+      raise
+    end
+
+    def http_client
+      BasicClient.new(create_url(url))
+    end
+
+    protected
+
+    def create_url(path)
+      if path =~ /^(http|https):\/\//
+        URI.parse(path)
+      else
+        URI.parse("#{@url}/#{path}")
+      end
+    end
+
     def apply_request_middleware(method, url, headers, data)
       middlewares.inject([method, url, headers, data]) do |req_data, middleware|
         middleware.handle_request(*req_data)
@@ -154,10 +222,6 @@ class Chef
       response.kind_of?(Net::HTTPSuccess) || response.kind_of?(Net::HTTPRedirection)
     end
 
-    def http_client
-      BasicClient.new(create_url(url))
-    end
-
     # Runs a synchronous HTTP request, with no middleware applied (use #request
     # to have the middleware applied). The entire response will be loaded into memory.
     def send_http_request(method, url, headers, body, &response_handler)
@@ -197,6 +261,7 @@ class Chef
       end
     end
 
+
     # Wraps an HTTP request with retry logic.
     # === Arguments
     # url:: URL of the request, used for error messages
@@ -274,6 +339,28 @@ class Chef
       headers
     end
 
+    def stream_to_tempfile(url, response)
+      tf = Tempfile.open("chef-rest")
+      if Chef::Platform.windows?
+        tf.binmode # required for binary files on Windows platforms
+      end
+      Chef::Log.debug("Streaming download from #{url.to_s} to tempfile #{tf.path}")
+      # Stolen from http://www.ruby-forum.com/topic/166423
+      # Kudos to _why!
+
+      stream_handler = StreamHandler.new(middlewares, response)
+
+      response.read_body do |chunk|
+        tf.write(stream_handler.handle_chunk(chunk))
+      end
+      tf.close
+      tf
+    rescue Exception
+      tf.close!
+      raise
+    end
+
+
     public
 
     ############################################################################
diff --git a/lib/chef/http/json_to_model_inflater.rb b/lib/chef/http/json_to_model_inflater.rb
index 3e941ad105..0b46054510 100644
--- a/lib/chef/http/json_to_model_inflater.rb
+++ b/lib/chef/http/json_to_model_inflater.rb
@@ -29,7 +29,10 @@ class Chef
       end
 
       def handle_request(method, url, headers={}, data=false)
-        headers['Accept']       = "application/json"
+        # Ideally this should always set Accept to application/json, but
+        # Chef::REST is sometimes used to make non-JSON requests, so it sets
+        # Accept to the desired value before middlewares get called.
+        headers['Accept']       ||= "application/json"
         headers["Content-Type"] = 'application/json' if data
         json_body = data ? Chef::JSONCompat.to_json(data) : nil
         # Force encoding to binary to fix SSL related EOFErrors
diff --git a/lib/chef/rest.rb b/lib/chef/rest.rb
index 033153f739..b264f96336 100644
--- a/lib/chef/rest.rb
+++ b/lib/chef/rest.rb
@@ -140,47 +140,15 @@ class Chef
       end
     end
 
-    alias :retriable_rest_request :retriable_http_request
-
-    # Makes a streaming download request. <b>Doesn't speak JSON.</b>
-    # Streams the response body to a tempfile. If a block is given, it's
-    # passed to Tempfile.open(), which means that the tempfile will automatically
-    # be unlinked after the block is executed.
-    #
-    # If no block is given, the tempfile is returned, which means it's up to
-    # you to unlink the tempfile when you're done with it.
+    # Customized streaming behavior; sets the accepted content type to "*/*"
+    # if not otherwise specified for compatibility purposes
     def streaming_request(url, headers, &block)
-      rest_request = nil # ensure this is defined, it's referenced in rescue clause
-
-      # Manually apply middleware to avoid specifying an Accept
-      # application/json content type preference:
-      method, url, headers, data = [@decompressor, @authenticator].inject([:GET, url, headers, nil]) do |req_data, middleware|
-        middleware.handle_request(*req_data)
-      end
-      response, rest_request, return_value = send_http_request(method, url, headers, data) do |http_response|
-        @last_response = http_response
-        if http_response.kind_of?(Net::HTTPSuccess)
-          tempfile = stream_to_tempfile(url, http_response)
-          if block_given?
-            begin
-              yield tempfile
-            ensure
-              tempfile && tempfile.close!
-            end
-          end
-          return tempfile
-        end
-      end
-      unless response.kind_of?(Net::HTTPSuccess) or response.kind_of?(Net::HTTPRedirection)
-        response.error!
-      end
-    rescue Exception => e
-      if e.respond_to?(:chef_rest_request=)
-        e.chef_rest_request = rest_request
-      end
-      raise
+      headers["Accept"] ||= "*/*"
+      super
     end
 
+    alias :retriable_rest_request :retriable_http_request
+
     def follow_redirect
       unless @sign_on_redirect
         @authenticator.sign_request = false
@@ -190,64 +158,12 @@ class Chef
       @authenticator.sign_request = true
     end
 
-    def http_client
-      BasicClient.new(create_url(url))
-    end
-
-    class StreamHandler
-      def initialize(middlewares, response)
-        middlewares = middlewares.flatten
-        @stream_handlers = []
-        middlewares.each do |middleware|
-          stream_handler = middleware.stream_response_handler(response)
-          @stream_handlers << stream_handler unless stream_handler.nil?
-        end
-      end
-
-      def handle_chunk(next_chunk)
-        @stream_handlers.inject(next_chunk) do |chunk, handler|
-          handler.handle_chunk(chunk)
-        end
-      end
-
-    end
-
-    private
-
-    def stream_to_tempfile(url, response)
-      tf = Tempfile.open("chef-rest")
-      if Chef::Platform.windows?
-        tf.binmode # required for binary files on Windows platforms
-      end
-      Chef::Log.debug("Streaming download from #{url.to_s} to tempfile #{tf.path}")
-      # Stolen from http://www.ruby-forum.com/topic/166423
-      # Kudos to _why!
-
-      stream_handler = StreamHandler.new(middlewares, response)
-
-      response.read_body do |chunk|
-        tf.write(stream_handler.handle_chunk(chunk))
-      end
-      tf.close
-      tf
-    rescue Exception
-      tf.close!
-      raise
-    end
-
-    public
+    public :create_url
 
     ############################################################################
     # DEPRECATED
     ############################################################################
 
-    # This is only kept around to provide access to cache control data in
-    # lib/chef/provider/remote_file/http.rb
-    # Find a better API.
-    def last_response
-      @last_response
-    end
-
     def decompress_body(body)
       @decompressor.decompress_body(body)
     end
diff --git a/spec/unit/rest_spec.rb b/spec/unit/rest_spec.rb
index c25c8233b1..073ff691d6 100644
--- a/spec/unit/rest_spec.rb
+++ b/spec/unit/rest_spec.rb
@@ -183,7 +183,7 @@ describe Chef::REST do
       end
 
       it "should build a new HTTP GET request without the application/json accept header" do
-        expected_headers = {'X-Chef-Version' => Chef::VERSION, 'Accept-Encoding' => Chef::REST::RESTRequest::ENCODING_GZIP_DEFLATE}
+        expected_headers = {'Accept' => "*/*", 'X-Chef-Version' => Chef::VERSION, 'Accept-Encoding' => Chef::REST::RESTRequest::ENCODING_GZIP_DEFLATE}
         Net::HTTP::Get.should_receive(:new).with("/?foo=bar", expected_headers).and_return(@request_mock)
         @rest.streaming_request(@url, {})
       end
@@ -428,7 +428,7 @@ describe Chef::REST do
       end
 
       it " build a new HTTP GET request without the application/json accept header" do
-        expected_headers = {'X-Chef-Version' => Chef::VERSION, 'Accept-Encoding' => Chef::REST::RESTRequest::ENCODING_GZIP_DEFLATE}
+        expected_headers = {'Accept' => "*/*", 'X-Chef-Version' => Chef::VERSION, 'Accept-Encoding' => Chef::REST::RESTRequest::ENCODING_GZIP_DEFLATE}
         Net::HTTP::Get.should_receive(:new).with("/?foo=bar", expected_headers).and_return(@request_mock)
         @rest.streaming_request(@url, {})
       end