Pass headers down instead of using a block

Right now we are buffering entire gems in memory when we download them.

  * https://github.com/rubygems/rubygems/blob/66c94f65c8a05484f2d12cb0bccd8513da5c65ef/lib/rubygems/request.rb#L208-L216

I would like to open the possibility of streaming assets rather than
storing the entire asset in memory.  In order to accomplish this, I
would like to allow callers to pass a block to indicate that they "want
to stream".

Before this commit, we would use the block to set headers on
the request object.  This commit passes the headers down to construct
the request object.  Callers know the headers they want to set before
calling the method, so it makes sense to pass them in via parameters
rather than as a block

Author: tenderlove

lib/rubygems/remote_fetcher.rb

@@ -210,9 +210,7 @@ def fetch_file(uri, *_)
 
   def fetch_http(uri, last_modified = nil, head = false, depth = 0)
     fetch_type = head ? Gem::Net::HTTP::Head : Gem::Net::HTTP::Get
-    response   = request uri, fetch_type, last_modified do |req|
-      headers.each {|k,v| req.add_field(k,v) }
-    end
+    response   = request uri, fetch_type, last_modified, headers
 
     case response
     when Gem::Net::HTTPOK, Gem::Net::HTTPNotModified then
@@ -308,15 +306,13 @@ def cache_update_path(uri, path = nil, update = true)
   # a Gem::Net::HTTP response object.  request maintains a table of persistent
   # connections to reduce connect overhead.
 
-  def request(uri, request_class, last_modified = nil)
+  def request(uri, request_class, last_modified = nil, headers = {})
     proxy = proxy_for @proxy, uri
     pool  = pools_for(proxy).pool_for uri
 
-    request = Gem::Request.new uri, request_class, last_modified, pool
+    request = Gem::Request.new uri, request_class, last_modified, pool, headers
 
-    request.fetch do |req|
-      yield req if block_given?
-    end
+    request.fetch
   end
 
   def https?(uri)

lib/rubygems/request.rb

@@ -14,7 +14,7 @@ def self.create_with_proxy(uri, request_class, last_modified, proxy) # :nodoc:
     proxy ||= get_proxy_from_env(uri.scheme)
     pool = ConnectionPools.new proxy_uri(proxy), cert_files
 
-    new(uri, request_class, last_modified, pool.pool_for(uri))
+    new(uri, request_class, last_modified, pool.pool_for(uri), {})
   end
 
   def self.proxy_uri(proxy) # :nodoc:
@@ -26,12 +26,13 @@ def self.proxy_uri(proxy) # :nodoc:
     end
   end
 
-  def initialize(uri, request_class, last_modified, pool)
+  def initialize(uri, request_class, last_modified, pool, initheader)
     @uri = uri
     @request_class = request_class
     @last_modified = last_modified
     @requests = Hash.new(0).compare_by_identity
     @user_agent = user_agent
+    @initheader = initheader
 
     @connection_pool = pool
   end
@@ -139,7 +140,7 @@ def connection_for(uri)
   end
 
   def fetch
-    request = @request_class.new @uri.request_uri
+    request = @request_class.new @uri.request_uri, @initheader
 
     unless @uri.nil? || @uri.user.nil? || @uri.user.empty?
       request.basic_auth Gem::UriFormatter.new(@uri.user).unescape,
@@ -155,8 +156,6 @@ def fetch
       request.add_field "If-Modified-Since", @last_modified.httpdate
     end
 
-    yield request if block_given?
-
     perform_request request
   end
 

lib/rubygems/s3_uri_signer.rb

@@ -177,13 +177,11 @@ def ec2_metadata_credentials_imds_v1
   end
 
   def ec2_metadata_request(url, token:)
-    request = ec2_iam_request(Gem::URI(url), Gem::Net::HTTP::Get)
+    request = ec2_iam_request(Gem::URI(url), Gem::Net::HTTP::Get, token ? {
+      "X-aws-ec2-metadata-token" => token
+    } : {})
 
-    response = request.fetch do |req|
-      if token
-        req.add_field "X-aws-ec2-metadata-token", token
-      end
-    end
+    response = request.fetch
 
     case response
     when Gem::Net::HTTPOK then
@@ -194,11 +192,11 @@ def ec2_metadata_request(url, token:)
   end
 
   def ec2_metadata_token
-    request = ec2_iam_request(Gem::URI(EC2_IAM_TOKEN), Gem::Net::HTTP::Put)
+    request = ec2_iam_request(Gem::URI(EC2_IAM_TOKEN), Gem::Net::HTTP::Put, {
+      "X-aws-ec2-metadata-token-ttl-seconds" => "60"
+    })
 
-    response = request.fetch do |req|
-      req.add_field "X-aws-ec2-metadata-token-ttl-seconds", 60
-    end
+    response = request.fetch
 
     case response
     when Gem::Net::HTTPOK then
@@ -208,9 +206,9 @@ def ec2_metadata_token
     end
   end
 
-  def ec2_iam_request(uri, verb)
+  def ec2_iam_request(uri, verb, headers)
     @request_pool ||= create_request_pool(uri)
-    Gem::Request.new(uri, verb, nil, @request_pool)
+    Gem::Request.new(uri, verb, nil, @request_pool, headers)
   end
 
   def create_request_pool(uri)

test/rubygems/test_gem_remote_fetcher.rb

@@ -516,7 +516,7 @@ def test_fetch_http
     @fetcher = fetcher
     url = "http://gems.example.com/redirect"
 
-    def fetcher.request(uri, request_class, last_modified = nil)
+    def fetcher.request(uri, request_class, last_modified = nil, headers = {})
       url = "http://gems.example.com/redirect"
       if defined? @requested
         res = Gem::Net::HTTPOK.new nil, 200, nil
@@ -541,7 +541,7 @@ def test_fetch_http_redirects
     @fetcher = fetcher
     url = "http://gems.example.com/redirect"
 
-    def fetcher.request(uri, request_class, last_modified = nil)
+    def fetcher.request(uri, request_class, last_modified = nil, headers = {})
       url = "http://gems.example.com/redirect"
       res = Gem::Net::HTTPMovedPermanently.new nil, 301, nil
       res.add_field "Location", url
@@ -560,7 +560,7 @@ def test_fetch_http_redirects_without_location
     @fetcher = fetcher
     url = "http://gems.example.com/redirect"
 
-    def fetcher.request(uri, request_class, last_modified = nil)
+    def fetcher.request(uri, request_class, last_modified = nil, headers = {})
       res = Gem::Net::HTTPMovedPermanently.new nil, 301, nil
       res
     end
@@ -572,18 +572,6 @@ def fetcher.request(uri, request_class, last_modified = nil)
     assert_equal "redirecting but no redirect location was given (#{url})", e.message
   end
 
-  def test_request_block
-    fetcher = Gem::RemoteFetcher.new nil
-    @fetcher = fetcher
-
-    assert_throws :block_called do
-      fetcher.request Gem::URI("http://example"), Gem::Net::HTTP::Get do |req|
-        assert_kind_of Gem::Net::HTTPGenericRequest, req
-        throw :block_called
-      end
-    end
-  end
-
   def test_yaml_error_on_size
     use_ui @stub_ui do
       fetcher = Gem::RemoteFetcher.new nil

test/rubygems/test_gem_remote_fetcher_s3.rb

@@ -49,8 +49,8 @@ def initialize(uri, method)
       super
     end
 
-    def ec2_iam_request(uri, verb)
-      fake_s3_request = FakeGemRequest.new(uri, verb, nil, nil)
+    def ec2_iam_request(uri, verb, headers)
+      fake_s3_request = FakeGemRequest.new(uri, verb, nil, nil, headers)
       @aws_iam_calls << fake_s3_request
 
       case uri.to_s
@@ -90,7 +90,7 @@ def res.body = FakeS3URISigner.instance_profile
   class FakeGemFetcher < Gem::RemoteFetcher
     attr_reader :fetched_uri, :last_s3_uri_signer
 
-    def request(uri, request_class, last_modified = nil)
+    def request(uri, ...)
       @fetched_uri = uri
       res = Gem::Net::HTTPOK.new nil, 200, nil
       def res.body = "success"