Add support for chunked downloads

.rubocop.yml

@@ -19,12 +19,18 @@ Metrics/BlockLength:
 Metrics/ClassLength:
   Enabled: false
 
+Metrics/CyclomaticComplexity:
+  Enabled: false
+
 Metrics/MethodLength:
   Enabled: false
 
 Naming/MethodParameterName:
   Enabled: false
 
+Metrics/PerceivedComplexity:
+  Enabled: false
+
 Layout/CaseIndentation:
   EnforcedStyle: end
 

README.md

@@ -86,11 +86,19 @@ end
 resolver = proc do |hostname|
   [IPAddr.new('2001:500:8f::53')] # Static resolver
 end
-SsrfFilter.get('https://www.example.com', resolver: resolver) do |request|
+request_proc = proc do |request|
   # Do some extra processing on the request
   request['content-type'] = 'application/json'
   request.basic_auth('username', 'password')
 end
+SsrfFilter.get('https://www.example.com', resolver: resolver, request_proc: request_proc)
+
+# Stream response
+SsrfFilter.get('https://www.example.com') do |response|
+  response.read_body do |chunk|
+    puts chunk
+  end
+end
 ```
 
 ### Changelog

lib/ssrf_filter/ssrf_filter.rb

@@ -127,16 +127,8 @@ class CRLFInjection < Error
         public_addresses = ip_addresses.reject(&method(:unsafe_ip_address?))
         raise PrivateIPAddress, "Hostname '#{hostname}' has no public ip addresses" if public_addresses.empty?
 
-        response = fetch_once(uri, public_addresses.sample.to_s, method, options, &block)
-
-        case response
-        when ::Net::HTTPRedirection
-          url = response['location']
-          # Handle relative redirects
-          url = "#{uri.scheme}://#{hostname}:#{uri.port}#{url}" if url.start_with?('/')
-        else
-          return response
-        end
+        response, url = fetch_once(uri, public_addresses.sample.to_s, method, options, &block)
+        return response if url.nil?
       end
 
       raise TooManyRedirects, "Got #{max_redirects} redirects fetching #{original_url}"
@@ -189,15 +181,26 @@ def self.fetch_once(uri, ip, verb, options, &block)
 
     request.body = options[:body] if options[:body]
 
-    block.call(request) if block_given?
+    options[:request_proc].call(request) if options[:request_proc].respond_to?(:call)
     validate_request(request)
 
     http_options = options[:http_options] || {}
     http_options[:use_ssl] = (uri.scheme == 'https')
 
     with_forced_hostname(hostname) do
       ::Net::HTTP.start(uri.hostname, uri.port, **http_options) do |http|
-        http.request(request)
+        http.request(request) do |response|
+          case response
+          when ::Net::HTTPRedirection
+            url = response['location']
+            # Handle relative redirects
+            url = "#{uri.scheme}://#{hostname}:#{uri.port}#{url}" if url.start_with?('/')
+            return nil, url
+          else
+            block&.call(response)
+            return response, nil
+          end
+        end
       end
     end
   end

spec/lib/ssrf_filter_spec.rb

@@ -81,49 +81,55 @@
     it 'sets the host header' do
       stub_request(:post, "https://#{public_ipv4}").with(headers: {host: 'www.example.com'})
         .to_return(status: 200, body: 'response body')
-      response = described_class.fetch_once(URI('https://www.example.com'), public_ipv4.to_s, :post, {})
+      response, url = described_class.fetch_once(URI('https://www.example.com'), public_ipv4.to_s, :post, {})
       expect(response.code).to eq('200')
       expect(response.body).to eq('response body')
+      expect(url).to be_nil
     end
 
     it 'does not send the port in the host header for default ports (http)' do
       stub_request(:post, "http://#{public_ipv4}").with(headers: {host: 'www.example.com'})
         .to_return(status: 200, body: 'response body')
-      response = described_class.fetch_once(URI('http://www.example.com'), public_ipv4.to_s, :post, {})
+      response, url = described_class.fetch_once(URI('http://www.example.com'), public_ipv4.to_s, :post, {})
       expect(response.code).to eq('200')
       expect(response.body).to eq('response body')
+      expect(url).to be_nil
     end
 
     it 'sends the port in the host header for non-default ports' do
       stub_request(:post, "https://#{public_ipv4}:80").with(headers: {host: 'www.example.com:80'})
         .to_return(status: 200, body: 'response body')
-      response = described_class.fetch_once(URI('https://www.example.com:80'), public_ipv4.to_s, :post, {})
+      response, url = described_class.fetch_once(URI('https://www.example.com:80'), public_ipv4.to_s, :post, {})
       expect(response.code).to eq('200')
       expect(response.body).to eq('response body')
+      expect(url).to be_nil
     end
 
     it 'passes headers, params, and blocks' do
       stub_request(:get, "https://#{public_ipv4}/?key=value").with(headers:
         {host: 'www.example.com', header: 'value', header2: 'value2'}).to_return(status: 200, body: 'response body')
       options = {
         headers: {'header' => 'value'},
-        params: {'key' => 'value'}
+        params: {'key' => 'value'},
+        request_proc: proc do |req|
+          req['header2'] = 'value2'
+        end
       }
       uri = URI('https://www.example.com/?key=value')
-      response = described_class.fetch_once(uri, public_ipv4.to_s, :get, options) do |req|
-        req['header2'] = 'value2'
-      end
+      response, url = described_class.fetch_once(uri, public_ipv4.to_s, :get, options)
       expect(response.code).to eq('200')
       expect(response.body).to eq('response body')
+      expect(url).to be_nil
     end
 
     it 'merges params' do
       stub_request(:get, "https://#{public_ipv4}/?key=value&key2=value2")
         .with(headers: {host: 'www.example.com'}).to_return(status: 200, body: 'response body')
       uri = URI('https://www.example.com/?key=value')
-      response = described_class.fetch_once(uri, public_ipv4.to_s, :get, params: {'key2' => 'value2'})
+      response, url = described_class.fetch_once(uri, public_ipv4.to_s, :get, params: {'key2' => 'value2'})
       expect(response.code).to eq('200')
       expect(response.body).to eq('response body')
+      expect(url).to be_nil
     end
 
     it 'does not use tls for http urls', only: true do
@@ -270,7 +276,7 @@ def inject_custom_trust_store(*certificates)
           expect(response['X-Host']).to eq("#{hostname}:#{port}")
         end
       ensure
-        web_server_thread.kill
+        web_server_thread&.kill
       end
     end
 
@@ -317,7 +323,55 @@ def inject_custom_trust_store(*certificates)
           expect(response['X-Host']).to eq("virtualhost:#{port}")
         end
       ensure
-        web_server_thread.kill
+        web_server_thread&.kill
+      end
+    end
+
+    it 'supports chunked responses' do
+      hostname = 'ssrf-filter.example.com'
+      port = 8443
+
+      private_key, certificate = make_keypair("CN=#{hostname}")
+      inject_custom_trust_store(certificate)
+      stub_const('SsrfFilter::IPV4_BLACKLIST', [])
+
+      begin
+        queue = Queue.new # Used as a semaphore
+
+        chunks = ['chunk 1', 'chunk 2', 'chunk 3']
+
+        web_server_thread = Thread.new do
+          server = make_web_server(port, private_key, certificate) do
+            queue.push(nil)
+          end
+
+          server.mount_proc '/chunked' do |_, res|
+            res.status = 200
+            res.chunked = true
+            res.body = proc do |chunked_wrapper|
+              chunks.each { |chunk| chunked_wrapper.write(chunk) }
+            end
+          end
+
+          server.start
+        end
+
+        Timeout.timeout(2) do
+          queue.pop
+
+          chunk_index = 0
+          url = "https://#{hostname}:#{port}/chunked"
+          described_class.get(url, resolver: proc { [IPAddr.new('127.0.0.1')] }) do |response|
+            expect(response.code).to eq('200')
+            response.read_body do |chunk|
+              expect(chunk).to eq(chunks[chunk_index])
+              chunk_index += 1
+            end
+          end
+          expect(chunk_index).to eq(chunks.length)
+        end
+      ensure
+        web_server_thread&.kill
       end
     end
   end

spec/spec_helper.rb

@@ -19,11 +19,11 @@
 
 def allow_net_connections_for_context(context)
   context.before :all do
-    WebMock.allow_net_connect!
+    WebMock.disable!
   end
 
   context.after :all do
-    WebMock.disable_net_connect!
+    WebMock.enable!
   end
 end