From 7f0b0485485fe37883b29b45d8857adfd582c6bf Mon Sep 17 00:00:00 2001
From: Arkadiy Tetelman <arkadiyt@gmail.com>
Date: Fri, 1 Sep 2023 11:57:53 -0700
Subject: [PATCH] Resolve conflicts

---
 lib/ssrf_filter/ssrf_filter.rb | 21 ++++++++++++++++-----
 spec/lib/ssrf_filter_spec.rb   | 15 +++++++++++++++
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/lib/ssrf_filter/ssrf_filter.rb b/lib/ssrf_filter/ssrf_filter.rb
index c3bd73e..9dbac54 100644
--- a/lib/ssrf_filter/ssrf_filter.rb
+++ b/lib/ssrf_filter/ssrf_filter.rb
@@ -72,6 +72,7 @@ def self.prefixlen_from_ipaddr(ipaddr)
     ::Resolv.getaddresses(hostname).map { |ip| ::IPAddr.new(ip) }
   end
 
+  DEFAULT_ALLOW_UNFOLLOWED_REDIRECTS = false
   DEFAULT_MAX_REDIRECTS = 10
 
   VERB_MAP = {
@@ -110,10 +111,13 @@ class CRLFInjection < Error
       original_url = url
       scheme_whitelist = options[:scheme_whitelist] || DEFAULT_SCHEME_WHITELIST
       resolver = options[:resolver] || DEFAULT_RESOLVER
+      allow_unfollowed_redirects = options.fetch(:allow_unfollowed_redirects) { DEFAULT_ALLOW_UNFOLLOWED_REDIRECTS }
       max_redirects = options[:max_redirects] || DEFAULT_MAX_REDIRECTS
       url = url.to_s
 
-      (max_redirects + 1).times do
+      redirects = 0
+
+      loop do
         uri = URI(url)
 
         unless scheme_whitelist.include?(uri.scheme)
@@ -128,10 +132,17 @@ class CRLFInjection < Error
         raise PrivateIPAddress, "Hostname '#{hostname}' has no public ip addresses" if public_addresses.empty?
 
         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}"
+        break response if url.nil?
+
+        if max_redirects <= redirects
+          break response if allow_unfollowed_redirects
+
+          raise TooManyRedirects, "Got #{max_redirects} redirects fetching #{original_url}"
+        end
+
+        redirects += 1
+      end
     end
   end
 
@@ -197,7 +208,7 @@ def self.fetch_once(uri, ip, verb, options, &block)
           url = response['location']
           # Handle relative redirects
           url = "#{uri.scheme}://#{hostname}:#{uri.port}#{url}" if url.start_with?('/')
-          return nil, url
+          return response, url
         else
           return response, nil
         end
diff --git a/spec/lib/ssrf_filter_spec.rb b/spec/lib/ssrf_filter_spec.rb
index 4e0a36b..897ae8f 100644
--- a/spec/lib/ssrf_filter_spec.rb
+++ b/spec/lib/ssrf_filter_spec.rb
@@ -460,6 +460,21 @@ def inject_custom_trust_store(*certificates)
       end.to raise_error(described_class::TooManyRedirects)
     end
 
+    it 'returns the last response if there are too many redirects and unfollowed redirects are allowed' do
+      stub_request(:get, "https://#{public_ipv4}").with(headers: {host: 'www.example.com'})
+        .to_return(status: 301, headers: {location: 'https://www.example2.com'})
+      resolver = proc { [public_ipv4] }
+      response =
+        described_class.get(
+          'https://www.example.com',
+          resolver: resolver,
+          allow_unfollowed_redirects: true,
+          max_redirects: 0
+        )
+      expect(response.code).to eq('301')
+      expect(response['location']).to eq('https://www.example2.com')
+    end
+
     it 'fails if the redirected url is not in the scheme whitelist' do
       stub_request(:put, "https://#{public_ipv4}").with(headers: {host: 'www.example.com'})
         .to_return(status: 301, headers: {location: 'ftp://www.example.com'})