diff --git a/lib/final_destination/ssrf_detector.rb b/lib/final_destination/ssrf_detector.rb index 0d06306ce8f..d9953044ded 100644 --- a/lib/final_destination/ssrf_detector.rb +++ b/lib/final_destination/ssrf_detector.rb @@ -78,6 +78,14 @@ class FinalDestination ips end + def self.allow_ip_lookups_in_test! + @allow_ip_lookups_in_test = true + end + + def self.disallow_ip_lookups_in_test! + @allow_ip_lookups_in_test = false + end + private def self.ip_in_ranges?(ip, ranges) @@ -85,7 +93,7 @@ class FinalDestination end def self.lookup_ips(name, timeout: nil) - if Rails.env.test? + if Rails.env.test? && !@allow_ip_lookups_in_test ["1.2.3.4"] else FinalDestination::Resolver.lookup(name, timeout: timeout) diff --git a/spec/lib/final_destination/http_spec.rb b/spec/lib/final_destination/http_spec.rb index 3388a2c835b..9ef4977826b 100644 --- a/spec/lib/final_destination/http_spec.rb +++ b/spec/lib/final_destination/http_spec.rb @@ -7,9 +7,14 @@ describe FinalDestination::HTTP do Socket.stubs(:tcp).never TCPSocket.stubs(:open).never Addrinfo.stubs(:getaddrinfo).never + + FinalDestination::SSRFDetector.allow_ip_lookups_in_test! end - after { WebMock.enable! } + after do + WebMock.enable! + FinalDestination::SSRFDetector.disallow_ip_lookups_in_test! + end def expect_tcp_and_abort(stub_addr, &blk) success = Class.new(StandardError) @@ -21,7 +26,12 @@ describe FinalDestination::HTTP do end def stub_ip_lookup(stub_addr, ips) - FinalDestination::SSRFDetector.stubs(:lookup_ips).with { |addr| stub_addr == addr }.returns(ips) + Addrinfo + .stubs(:getaddrinfo) + .with { |addr, _| addr == stub_addr } + .returns( + ips.map { |ip| Addrinfo.new([IPAddr.new(ip).ipv6? ? "AF_INET6" : "AF_INET", 80, nil, ip]) }, + ) end def stub_tcp_to_raise(stub_addr, exception)