From 3d8984c385a6024705cfc2f0e8ca066523bffc5f Mon Sep 17 00:00:00 2001 From: Aaron Patterson Date: Wed, 15 Oct 2025 10:27:55 -0700 Subject: [PATCH] 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 --- lib/rubygems/remote_fetcher.rb | 12 ++++------- lib/rubygems/request.rb | 9 ++++----- lib/rubygems/s3_uri_signer.rb | 22 ++++++++++----------- test/rubygems/test_gem_remote_fetcher.rb | 18 +++-------------- test/rubygems/test_gem_remote_fetcher_s3.rb | 6 +++--- 5 files changed, 24 insertions(+), 43 deletions(-) diff --git a/lib/rubygems/remote_fetcher.rb b/lib/rubygems/remote_fetcher.rb index 6ed0842963e6..40615376832e 100644 --- a/lib/rubygems/remote_fetcher.rb +++ b/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) diff --git a/lib/rubygems/request.rb b/lib/rubygems/request.rb index 9116785231ea..c2f32251ec59 100644 --- a/lib/rubygems/request.rb +++ b/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 diff --git a/lib/rubygems/s3_uri_signer.rb b/lib/rubygems/s3_uri_signer.rb index 148cba38c47a..eb22285897d1 100644 --- a/lib/rubygems/s3_uri_signer.rb +++ b/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) diff --git a/test/rubygems/test_gem_remote_fetcher.rb b/test/rubygems/test_gem_remote_fetcher.rb index 5c1d89fad693..335d0391cce3 100644 --- a/test/rubygems/test_gem_remote_fetcher.rb +++ b/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 diff --git a/test/rubygems/test_gem_remote_fetcher_s3.rb b/test/rubygems/test_gem_remote_fetcher_s3.rb index 4a5acc5a8666..9fe82af5b408 100644 --- a/test/rubygems/test_gem_remote_fetcher_s3.rb +++ b/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"