From 0d0effd983e27d5073fd86d2c6bb62aabe5119bc Mon Sep 17 00:00:00 2001
From: Zack <zack.scholl@gmail.com>
Date: Thu, 13 Jun 2024 08:27:55 -0700
Subject: [PATCH] fix: croc hangs with broken dns Fixes #735

---
 src/models/constants.go | 26 +++++++++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

diff --git a/src/models/constants.go b/src/models/constants.go
index bb25e8a..1622518 100644
--- a/src/models/constants.go
+++ b/src/models/constants.go
@@ -6,8 +6,10 @@ import (
 	"net"
 	"os"
 	"path"
+	"time"
 
 	"github.com/schollz/croc/v10/src/utils"
+	log "github.com/schollz/logger"
 )
 
 // TCP_BUFFER_SIZE is the maximum packet size
@@ -54,6 +56,8 @@ func getConfigFile(requireValidPath bool) (fname string, err error) {
 }
 
 func init() {
+	log.SetLevel("info")
+	log.SetOutput(os.Stderr)
 	doRemember := false
 	for _, flag := range os.Args {
 		if flag == "--internal-dns" {
@@ -78,6 +82,7 @@ func init() {
 			INTERNAL_DNS = utils.Exists(fname)
 		}
 	}
+	log.Trace("Using internal DNS: ", INTERNAL_DNS)
 	var err error
 	var addr string
 	addr, err = lookup(DEFAULT_RELAY)
@@ -86,17 +91,20 @@ func init() {
 	} else {
 		DEFAULT_RELAY = ""
 	}
+	log.Tracef("Default ipv4 relay: %s", addr)
 	addr, err = lookup(DEFAULT_RELAY6)
 	if err == nil {
 		DEFAULT_RELAY6 = net.JoinHostPort(addr, DEFAULT_PORT)
 	} else {
 		DEFAULT_RELAY6 = ""
 	}
+	log.Tracef("Default ipv6 relay: %s", addr)
 }
 
 // Resolve a hostname to an IP address using DNS.
 func lookup(address string) (ipaddress string, err error) {
 	if !INTERNAL_DNS {
+		log.Tracef("Using local DNS to resolve %s", address)
 		return localLookupIP(address)
 	}
 	type Result struct {
@@ -108,11 +116,13 @@ func lookup(address string) (ipaddress string, err error) {
 		go func(dns string) {
 			var r Result
 			r.s, r.err = remoteLookupIP(address, dns)
+			log.Tracef("Resolved %s to %s using %s", address, r.s, dns)
 			result <- r
 		}(dns)
 	}
 	for i := 0; i < len(publicDNS); i++ {
 		ipaddress = (<-result).s
+		log.Tracef("Resolved %s to %s", address, ipaddress)
 		if ipaddress != "" {
 			return
 		}
@@ -121,9 +131,16 @@ func lookup(address string) (ipaddress string, err error) {
 	return
 }
 
-// localLookupIP returns a host's IP address based on the local resolver.
+// localLookupIP returns a host's IP address using the local DNS configuration.
 func localLookupIP(address string) (ipaddress string, err error) {
-	ip, err := net.LookupHost(address)
+	// Create a context with a 500 millisecond timeout
+	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+	defer cancel()
+
+	r := &net.Resolver{}
+
+	// Use the context with timeout in the LookupHost function
+	ip, err := r.LookupHost(ctx, address)
 	if err != nil {
 		return
 	}
@@ -133,6 +150,9 @@ func localLookupIP(address string) (ipaddress string, err error) {
 
 // remoteLookupIP returns a host's IP address based on a remote DNS server.
 func remoteLookupIP(address, dns string) (ipaddress string, err error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+	defer cancel()
+
 	r := &net.Resolver{
 		PreferGo: true,
 		Dial: func(ctx context.Context, network, _ string) (net.Conn, error) {
@@ -140,7 +160,7 @@ func remoteLookupIP(address, dns string) (ipaddress string, err error) {
 			return d.DialContext(ctx, network, dns+":53")
 		},
 	}
-	ip, err := r.LookupHost(context.Background(), address)
+	ip, err := r.LookupHost(ctx, address)
 	if err != nil {
 		return
 	}