DEV: avoid mocking FinalDestination (#20570)

This commit is contained in:
Sam 2023-03-08 12:09:18 +11:00 committed by GitHub
parent f2476d4b80
commit 3f5fa4eb09
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 21 additions and 3 deletions

View File

@ -78,6 +78,14 @@ class FinalDestination
ips ips
end 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 private
def self.ip_in_ranges?(ip, ranges) def self.ip_in_ranges?(ip, ranges)
@ -85,7 +93,7 @@ class FinalDestination
end end
def self.lookup_ips(name, timeout: nil) def self.lookup_ips(name, timeout: nil)
if Rails.env.test? if Rails.env.test? && !@allow_ip_lookups_in_test
["1.2.3.4"] ["1.2.3.4"]
else else
FinalDestination::Resolver.lookup(name, timeout: timeout) FinalDestination::Resolver.lookup(name, timeout: timeout)

View File

@ -7,9 +7,14 @@ describe FinalDestination::HTTP do
Socket.stubs(:tcp).never Socket.stubs(:tcp).never
TCPSocket.stubs(:open).never TCPSocket.stubs(:open).never
Addrinfo.stubs(:getaddrinfo).never Addrinfo.stubs(:getaddrinfo).never
FinalDestination::SSRFDetector.allow_ip_lookups_in_test!
end end
after { WebMock.enable! } after do
WebMock.enable!
FinalDestination::SSRFDetector.disallow_ip_lookups_in_test!
end
def expect_tcp_and_abort(stub_addr, &blk) def expect_tcp_and_abort(stub_addr, &blk)
success = Class.new(StandardError) success = Class.new(StandardError)
@ -21,7 +26,12 @@ describe FinalDestination::HTTP do
end end
def stub_ip_lookup(stub_addr, ips) 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 end
def stub_tcp_to_raise(stub_addr, exception) def stub_tcp_to_raise(stub_addr, exception)