From 5325c8a283bd1a2a248661cf2b7001188c3202d8 Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Thu, 7 Nov 2019 16:40:37 +0100
Subject: [PATCH] p2p/enode: mock DNS resolver in URL parsing test (#20252)

---
 p2p/enode/urlv4.go      | 19 ++++++++++---------
 p2p/enode/urlv4_test.go | 10 ++++++++++
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/p2p/enode/urlv4.go b/p2p/enode/urlv4.go
index 0968694c4..712e1952b 100644
--- a/p2p/enode/urlv4.go
+++ b/p2p/enode/urlv4.go
@@ -31,7 +31,10 @@ import (
 	"github.com/ledgerwatch/turbo-geth/p2p/enr"
 )
 
-var incompleteNodeURL = regexp.MustCompile("(?i)^(?:enode://)?([0-9a-f]+)$")
+var (
+	incompleteNodeURL = regexp.MustCompile("(?i)^(?:enode://)?([0-9a-f]+)$")
+	lookupIPFunc      = net.LookupIP
+)
 
 // MustParseV4 parses a node URL. It panics if the URL is not valid.
 func MustParseV4(rawurl string) *Node {
@@ -107,7 +110,6 @@ func isNewV4(n *Node) bool {
 func parseComplete(rawurl string) (*Node, error) {
 	var (
 		id               *ecdsa.PublicKey
-		ip               net.IP
 		tcpPort, udpPort uint64
 	)
 	u, err := url.Parse(rawurl)
@@ -125,15 +127,14 @@ func parseComplete(rawurl string) (*Node, error) {
 		return nil, fmt.Errorf("invalid public key (%v)", err)
 	}
 	// Parse the IP address.
-	ips, err := net.LookupIP(u.Hostname())
-	if err != nil {
-		if _, ok := err.(*net.DNSError); ok {
-			// backward and cross-platform compatible version
-			return nil, errors.New("no such host")
+	ip := net.ParseIP(u.Hostname())
+	if ip == nil {
+		ips, err := lookupIPFunc(u.Hostname())
+		if err != nil {
+			return nil, err
 		}
-		return nil, err
+		ip = ips[0]
 	}
-	ip = ips[0]
 	// Ensure the IP is 4 bytes long for IPv4 addresses.
 	if ipv4 := ip.To4(); ipv4 != nil {
 		ip = ipv4
diff --git a/p2p/enode/urlv4_test.go b/p2p/enode/urlv4_test.go
index e52062994..56886d5d5 100644
--- a/p2p/enode/urlv4_test.go
+++ b/p2p/enode/urlv4_test.go
@@ -18,6 +18,7 @@ package enode
 
 import (
 	"crypto/ecdsa"
+	"errors"
 	"net"
 	"reflect"
 	"strings"
@@ -27,6 +28,15 @@ import (
 	"github.com/ledgerwatch/turbo-geth/p2p/enr"
 )
 
+func init() {
+	lookupIPFunc = func(name string) ([]net.IP, error) {
+		if name == "node.example.org" {
+			return []net.IP{{33, 44, 55, 66}}, nil
+		}
+		return nil, errors.New("no such host")
+	}
+}
+
 var parseNodeTests = []struct {
 	input      string
 	wantError  string