(PUP-11853) Wait for request completion before closing

Before this commit, puppet would attempt to `finish` a
connection in the block passed to an `http.request` function. This
will not work because `finish` sets the `@socket` to nil, and the
net/http request assumes that the `@socket` is set for the duration
of that function. With this change, we mark some state during the
request function that indicates the connection should be closed
after the completion of the `http.request`.

Author: tvpartytonight

lib/puppet/http/client.rb

@@ -367,6 +367,7 @@ def execute_streaming(request, options: {}, &block)
         apply_auth(request, basic_auth) if redirects.zero?
 
         # don't call return within the `request` block
+        close_and_sleep = nil
         http.request(request) do |nethttp|
           response = Puppet::HTTP::ResponseNetHTTP.new(request.uri, nethttp)
           begin
@@ -380,12 +381,14 @@ def execute_streaming(request, options: {}, &block)
               interval = @retry_after_handler.retry_after_interval(request, response, retries)
               retries += 1
               if interval
-                if http.started?
-                  Puppet.debug("Closing connection for #{Puppet::HTTP::Site.from_uri(request.uri)}")
-                  http.finish
+                close_and_sleep = proc do
+                  if http.started?
+                    Puppet.debug("Closing connection for #{Puppet::HTTP::Site.from_uri(request.uri)}")
+                    http.finish
+                  end
+                  Puppet.warning(_("Sleeping for %{interval} seconds before retrying the request") % { interval: interval })
+                  ::Kernel.sleep(interval)
                 end
-                Puppet.warning(_("Sleeping for %{interval} seconds before retrying the request") % { interval: interval })
-                ::Kernel.sleep(interval)
                 next
               end
             end
@@ -404,6 +407,10 @@ def execute_streaming(request, options: {}, &block)
 
           done = true
         end
+      ensure
+        # If a server responded with a retry, make sure the connection is closed and then
+        # sleep the specified time.
+        close_and_sleep.call if close_and_sleep
       end
     end
 

spec/integration/http/client_spec.rb

@@ -175,6 +175,22 @@
     end
   end
 
+  context 'ensure that retrying does not attempt to read the body after closing the connection' do
+    let(:client) { Puppet::HTTP::Client.new(retry_limit: 1) }
+    it 'raises a retry error instead' do
+      response_proc = -> (req, res) {
+        res['Retry-After'] = 1
+        res.status = 503
+      }
+
+      https_server.start_server(response_proc: response_proc) do |port|
+        uri = URI("https://127.0.0.1:#{port}")
+        kwargs = {headers: {'Content-Type' => 'text/plain'}, options: {ssl_context: root_context}}
+        expect{client.post(uri, '', **kwargs)}.to raise_error(Puppet::HTTP::TooManyRetryAfters)
+      end
+    end
+  end
+
   context 'persistent connections' do
     it "detects when the server has closed the connection and reconnects" do
       Puppet[:http_debug] = true