Merge branch 'andysoiron/patch-net-buffered-io-readuntil-with-timeout' into 'master'

Add Gitlab::BufferedIo with header read timeout

See merge request gitlab-org/gitlab!78065

config/feature_flags/development/header_read_timeout_buffered_io.yml

+---
+name: header_read_timeout_buffered_io
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/78065
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/350233
+milestone: '14.8'
+type: development
+group: group::integrations
+default_enabled: false

lib/gitlab/buffered_io.rb

+# frozen_string_literal: true
+
+module Gitlab
+  # Net::BufferedIO is overwritten by webmock but in order to test this class, it needs to inherit from the original BufferedIO.
+  # https://github.com/bblimke/webmock/blob/867f4b290fd133658aa9530cba4ba8b8c52c0d35/lib/webmock/http_lib_adapters/net_http.rb#L266
+  parent_class = if const_defined?('WebMock::HttpLibAdapters::NetHttpAdapter::OriginalNetBufferedIO') && Rails.env.test?
+                   WebMock::HttpLibAdapters::NetHttpAdapter::OriginalNetBufferedIO
+                 else
+                   Net::BufferedIO
+                 end
+
+  class BufferedIo < parent_class
+    extend ::Gitlab::Utils::Override
+
+    HEADER_READ_TIMEOUT = 20
+
+    # rubocop: disable Style/RedundantReturn
+    # rubocop: disable Cop/LineBreakAfterGuardClauses
+    # rubocop: disable Layout/EmptyLineAfterGuardClause
+
+    # Original method:
+    # https://github.com/ruby/ruby/blob/cdb7d699d0641e8f081d590d06d07887ac09961f/lib/net/protocol.rb#L190-L200
+    override :readuntil
+    def readuntil(terminator, ignore_eof = false)
+      start_time = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      check_timeout = Feature.enabled?(:header_read_timeout_buffered_io)
+
+      begin
+        until idx = @rbuf.index(terminator)
+          if check_timeout && (elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - start_time) > HEADER_READ_TIMEOUT
+            raise Gitlab::HTTP::HeaderReadTimeout, "Request timed out after reading headers for #{elapsed} seconds"
+          end
+
+          rbuf_fill
+        end
+
+        return rbuf_consume(idx + terminator.size)
+      rescue EOFError
+        raise unless ignore_eof
+        return rbuf_consume(@rbuf.size)
+      end
+    end
+    # rubocop: enable Style/RedundantReturn
+    # rubocop: enable Cop/LineBreakAfterGuardClauses
+    # rubocop: enable Layout/EmptyLineAfterGuardClause
+  end
+end

lib/gitlab/http.rb

@@ -9,6 +9,7 @@ module Gitlab
     BlockedUrlError = Class.new(StandardError)
     RedirectionTooDeep = Class.new(StandardError)
     ReadTotalTimeout = Class.new(Net::ReadTimeout)
+    HeaderReadTimeout = Class.new(Net::ReadTimeout)
 
     HTTP_TIMEOUT_ERRORS = [
       Net::OpenTimeout, Net::ReadTimeout, Net::WriteTimeout, Gitlab::HTTP::ReadTotalTimeout

lib/gitlab/http_connection_adapter.rb

 # frozen_string_literal: true
 
-# This class is part of the Gitlab::HTTP wrapper. Depending on the value
-# of the global setting allow_local_requests_from_web_hooks_and_services this adapter
-# will allow/block connection to internal IPs and/or urls.
+# This class is part of the Gitlab::HTTP wrapper. It handles local requests and header timeouts
 #
-# This functionality can be overridden by providing the setting the option
-# allow_local_requests = true in the request. For example:
-# Gitlab::HTTP.get('http://www.gitlab.com', allow_local_requests: true)
+# 1. Local requests
+#   Depending on the value of the global setting allow_local_requests_from_web_hooks_and_services,
+#   this adapter will allow/block connection to internal IPs and/or urls.
 #
-# This option will take precedence over the global setting.
+#   This functionality can be overridden by providing the setting the option
+#   allow_local_requests = true in the request. For example:
+#   Gitlab::HTTP.get('http://www.gitlab.com', allow_local_requests: true)
+#
+#   This option will take precedence over the global setting.
+#
+# 2. Header timeouts
+#   When the use_read_total_timeout option is used, that means the receiver
+#   of the HTTP request cannot be trusted. Gitlab::BufferedIo will be used,
+#   to read header data. It is a modified version of Net::BufferedIO that
+#   raises a timeout error if reading header data takes too much time.
+
 module Gitlab
   class HTTPConnectionAdapter < HTTParty::ConnectionAdapter
     extend ::Gitlab::Utils::Override
@@ -17,9 +26,20 @@ module Gitlab
     def connection
       @uri, hostname = validate_url!(uri)
 
-      super.tap do |http|
-        http.hostname_override = hostname if hostname
+      http = super
+      http.hostname_override = hostname if hostname
+
+      if options[:use_read_total_timeout]
+        gitlab_http = Gitlab::NetHttpAdapter.new(http.address, http.port)
+
+        http.instance_variables.each do |variable|
+          gitlab_http.instance_variable_set(variable, http.instance_variable_get(variable))
+        end
+
+        return gitlab_http
       end
+
+      http
     end
 
     private

lib/gitlab/net_http_adapter.rb

+# frozen_string_literal: true
+
+module Gitlab
+  # Webmock overwrites the Net::HTTP#request method with
+  # https://github.com/bblimke/webmock/blob/867f4b290fd133658aa9530cba4ba8b8c52c0d35/lib/webmock/http_lib_adapters/net_http.rb#L74
+  # Net::HTTP#request usually calls Net::HTTP#connect but the Webmock overwrite doesn't.
+  # This makes sure that, in a test environment, the superclass is the Webmock overwrite.
+  parent_class = if defined?(WebMock) && Rails.env.test?
+                   WebMock::HttpLibAdapters::NetHttpAdapter.instance_variable_get('@webMockNetHTTP')
+                 else
+                   Net::HTTP
+                 end
+
+  class NetHttpAdapter < parent_class
+    extend ::Gitlab::Utils::Override
+
+    private
+
+    override :connect
+    def connect
+      result = super
+
+      @socket = Gitlab::BufferedIo.new(@socket.io,
+        read_timeout: @socket.read_timeout,
+        write_timeout: @socket.write_timeout,
+        continue_timeout: @socket.continue_timeout,
+        debug_output: @socket.debug_output)
+
+      result
+    end
+  end
+end

spec/lib/gitlab/buffered_io_spec.rb

+# rubocop:disable Style/FrozenStringLiteralComment
+require 'spec_helper'
+
+RSpec.describe Gitlab::BufferedIo do
+  describe '#readuntil' do
+    let(:never_ending_tcp_socket) do
+      Class.new do
+        def initialize(*_)
+          @read_counter = 0
+        end
+
+        def setsockopt(*_); end
+
+        def closed?
+          false
+        end
+
+        def close
+          true
+        end
+
+        def to_io
+          StringIO.new('Hello World!')
+        end
+
+        def write_nonblock(data, *_)
+          data.size
+        end
+
+        def read_nonblock(buffer_size, *_)
+          sleep 0.01
+          @read_counter += 1
+
+          raise 'Test did not raise HeaderReadTimeout' if @read_counter > 10
+
+          'H' * buffer_size
+        end
+      end
+    end
+
+    before do
+      stub_const('Gitlab::BufferedIo::HEADER_READ_TIMEOUT', 0.1)
+    end
+
+    subject(:readuntil) do
+      Gitlab::BufferedIo.new(never_ending_tcp_socket.new).readuntil('a')
+    end
+
+    it 'raises a timeout error' do
+      expect { readuntil }.to raise_error(Gitlab::HTTP::HeaderReadTimeout, /Request timed out after reading headers for 0\.[0-9]+ seconds/)
+    end
+
+    context 'when the header_read_timeout feature is disabled' do
+      before do
+        stub_feature_flags(header_read_timeout_buffered_io: false)
+      end
+
+      it 'does not raise a timeout error' do
+        expect { readuntil }.to raise_error(RuntimeError, 'Test did not raise HeaderReadTimeout')
+      end
+    end
+  end
+end
+# rubocop:enable Style/FrozenStringLiteralComment

spec/lib/gitlab/http_connection_adapter_spec.rb

@@ -15,6 +15,18 @@ RSpec.describe Gitlab::HTTPConnectionAdapter do
       stub_all_dns('https://example.org', ip_address: '93.184.216.34')
     end
 
+    context 'with use_read_total_timeout option' do
+      let(:options) { { use_read_total_timeout: true } }
+
+      it 'sets up the connection using the Gitlab::NetHttpAdapter' do
+        expect(connection).to be_a(Gitlab::NetHttpAdapter)
+        expect(connection.address).to eq('93.184.216.34')
+        expect(connection.hostname_override).to eq('example.org')
+        expect(connection.addr_port).to eq('example.org')
+        expect(connection.port).to eq(443)
+      end
+    end
+
     context 'when local requests are allowed' do
       let(:options) { { allow_local_requests: true } }
 

spec/lib/gitlab/http_spec.rb

@@ -28,7 +28,7 @@ RSpec.describe Gitlab::HTTP do
   end
 
   context 'when reading the response is too slow' do
-    before do
+    before(:all) do
       # Override Net::HTTP to add a delay between sending each response chunk
       mocked_http = Class.new(Net::HTTP) do
         def request(*)
@@ -51,8 +51,17 @@ RSpec.describe Gitlab::HTTP do
       end
 
       @original_net_http = Net.send(:remove_const, :HTTP)
+      @webmock_net_http = WebMock::HttpLibAdapters::NetHttpAdapter.instance_variable_get('@webMockNetHTTP')
+
       Net.send(:const_set, :HTTP, mocked_http)
+      WebMock::HttpLibAdapters::NetHttpAdapter.instance_variable_set('@webMockNetHTTP', mocked_http)
+
+      # Reload Gitlab::NetHttpAdapter
+      Gitlab.send(:remove_const, :NetHttpAdapter)
+      load "#{Rails.root}/lib/gitlab/net_http_adapter.rb"
+    end
 
+    before do
       stub_const("#{described_class}::DEFAULT_READ_TOTAL_TIMEOUT", 0.001.seconds)
 
       WebMock.stub_request(:post, /.*/).to_return do |request|
@@ -60,9 +69,14 @@ RSpec.describe Gitlab::HTTP do
       end
     end
 
-    after do
+    after(:all) do
       Net.send(:remove_const, :HTTP)
       Net.send(:const_set, :HTTP, @original_net_http)
+      WebMock::HttpLibAdapters::NetHttpAdapter.instance_variable_set('@webMockNetHTTP', @webmock_net_http)
+
+      # Reload Gitlab::NetHttpAdapter
+      Gitlab.send(:remove_const, :NetHttpAdapter)
+      load "#{Rails.root}/lib/gitlab/net_http_adapter.rb"
     end
 
     let(:options) { {} }

spec/lib/gitlab/net_http_adapter_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::NetHttpAdapter do
+  describe '#connect' do
+    let(:url) { 'https://example.org' }
+    let(:net_http_adapter) { described_class.new(url) }
+
+    subject(:connect) { net_http_adapter.send(:connect) }
+
+    before do
+      allow(TCPSocket).to receive(:open).and_return(Socket.new(:INET, :STREAM))
+    end
+
+    it 'uses a Gitlab::BufferedIo instance as @socket' do
+      connect
+
+      expect(net_http_adapter.instance_variable_get(:@socket)).to be_a(Gitlab::BufferedIo)
+    end
+  end
+end