From 344a34b7bb2e0d92e1c8402602f958aa1efa056f Mon Sep 17 00:00:00 2001
From: Amir Zarrinkafsh <nightah@me.com>
Date: Thu, 28 Jul 2022 13:36:40 +1000
Subject: [PATCH] fix: multi-resolver querying and http client calls

This change resolves two key bugs:

1. All HTTP requests to detect a lancache heartbeat were incorrectly utilising the system resolver, we now utilise an HTTP transport with the appropriate resolver.
2. Works around a current Go limitation with Windows where all requests would utilise the system resolver even if attempting to manually override, we now use a Pure Go dns client library to make and parse said requests.
---
 const.go |  3 ++
 types.go |  4 +--
 utils.go | 92 +++++++++++++++++++++++++++++++++++++++++++++++---------
 3 files changed, 82 insertions(+), 17 deletions(-)

diff --git a/const.go b/const.go
index 0ce9be7..a3aa19f 100644
--- a/const.go
+++ b/const.go
@@ -17,6 +17,9 @@ const (
 
 	testPrefix     = "lancachetest."
 	wildcardPrefix = "*."
+
+	portHTTP = ":80"
+	portDNS  = ":53"
 )
 
 var (
diff --git a/types.go b/types.go
index f5da362..e9ba649 100644
--- a/types.go
+++ b/types.go
@@ -1,7 +1,5 @@
 package main
 
-import "net"
-
 type CDN struct {
 	Name string
 	File string
@@ -10,7 +8,7 @@ type CDN struct {
 type Lookup struct {
 	Resolver    string
 	Hostname    string
-	Address     []net.IP
+	Address     []string
 	ContainerID string
 	Time        string
 }
diff --git a/utils.go b/utils.go
index 4bd115b..4180f0a 100644
--- a/utils.go
+++ b/utils.go
@@ -11,6 +11,8 @@ import (
 	"reflect"
 	"strings"
 	"time"
+
+	"github.com/miekg/dns"
 )
 
 func getInterfaceAddresses(logger io.Writer) {
@@ -132,24 +134,42 @@ func lookupHostnames(host string, hostnames []string, iterations int, servers []
 
 func processHostnames(hostname, resolver string, logger io.Writer) (success, failed []Lookup) {
 	var (
-		ips []net.IP
-		err error
+		dialer    net.Dialer
+		transport http.Transport
+		ips       []string
+		err       error
 	)
 
 	if resolver != "system" {
-		r := &net.Resolver{
-			PreferGo: true,
-			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
-				d := net.Dialer{
-					Timeout: 1 * time.Second,
-				}
-				return d.DialContext(ctx, network, resolver)
-			},
+		// Fixed in Go 1.19: https://github.com/golang/go/issues/33097
+		//r := &net.Resolver{
+		//	PreferGo: true,
+		//	Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+		//		dialer = net.Dialer{
+		//			Timeout: 1 * time.Second,
+		//		}
+		//		return dialer.DialContext(ctx, network, resolver+portDNS)
+		//	},
+		//}
+		//
+		//ips, err = r.LookupHost(context.Background(), hostname)
+
+		ips, err = resolveIP(hostname, resolver+portDNS)
+		if len(ips) == 0 {
+			failed = append(failed, Lookup{
+				Resolver: resolver,
+				Hostname: hostname,
+				Time:     time.Now().Format(time.RFC822),
+			})
+			return success, failed
 		}
 
-		ips, err = r.LookupIP(context.Background(), "ip", hostname)
+		transport = http.Transport{DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
+			addr = ips[0] + portHTTP
+			return dialer.DialContext(ctx, network, addr)
+		}}
 	} else {
-		ips, err = net.LookupIP(hostname)
+		ips, err = net.LookupHost(hostname)
 	}
 
 	if err != nil {
@@ -157,13 +177,14 @@ func processHostnames(hostname, resolver string, logger io.Writer) (success, fai
 		failed = append(failed, Lookup{
 			Resolver: resolver,
 			Hostname: hostname,
-			Time:     time.Now().String(),
+			Time:     time.Now().Format(time.RFC822),
 		})
 		return success, failed
 	}
 
 	client := &http.Client{
-		Timeout: 1 * time.Second,
+		Timeout:   1 * time.Second,
+		Transport: &transport,
 		CheckRedirect: func(req *http.Request, via []*http.Request) error {
 			return http.ErrUseLastResponse
 		},
@@ -223,6 +244,49 @@ func parseCDN(name, file string, logger io.Writer) (hostnames []string) {
 	return hostnames
 }
 
+func resolveIP(name, resolver string) ([]string, error) {
+	var addresses []string
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+
+	client := new(dns.Client)
+	parsed := net.ParseIP(name)
+	if parsed != nil {
+		addresses = append(addresses, name)
+		return addresses, nil
+	}
+
+	messageAAAA := new(dns.Msg)
+	messageAAAA.SetQuestion(dns.Fqdn(name), dns.TypeAAAA)
+	inAAAA, _, err := client.ExchangeContext(ctx, messageAAAA, resolver)
+
+	if err != nil {
+		return nil, err
+	}
+
+	for _, record := range inAAAA.Answer {
+		if t, ok := record.(*dns.AAAA); ok {
+			addresses = append(addresses, t.AAAA.String())
+		}
+	}
+
+	messageA := new(dns.Msg)
+	messageA.SetQuestion(dns.Fqdn(name), dns.TypeA)
+	inA, _, err := client.ExchangeContext(ctx, messageA, resolver)
+
+	if err != nil {
+		return nil, err
+	}
+
+	for _, record := range inA.Answer {
+		if t, ok := record.(*dns.A); ok {
+			addresses = append(addresses, t.A.String())
+		}
+	}
+
+	return addresses, nil
+}
+
 func isLookupInSliceEqual(a []Lookup) []Lookup {
 	var l []Lookup