Optionally allow specification of custom DNS resolver

getlantern · Oct 5, 2023 · 106c0be · 106c0be
1 parent 08c7942
commit 106c0be
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 3 deletions.
diff --git a/http-proxy/main.go b/http-proxy/main.go
@@ -184,6 +184,8 @@ var (
 	broflakeAddr = flag.String("broflake-addr", "", "Address at which to listen for broflake connections.")
 
 	track = flag.String("track", "", "The track this proxy is running on")
+
+	dnsServer = flag.String("dns-server", "172.16.0.53", "Optional DNS server to use for DNS lookups (in place of system resolver)")
 )
 
 const (
@@ -468,6 +470,7 @@ func main() {
 		BroflakeAddr:                       *broflakeAddr,
 		BroflakeCert:                       os.Getenv("BROFLAKE_CERT"),
 		BroflakeKey:                        os.Getenv("BROFLAKE_KEY"),
+		DNSServer:                          *dnsServer,
 	}
 	if *maxmindLicenseKey != "" {
 		log.Debug("Will use Maxmind for geolocating clients")

diff --git a/http_proxy.go b/http_proxy.go
@@ -29,6 +29,7 @@ import (
 	shadowsocks "github.com/getlantern/http-proxy-lantern/v2/shadowsocks"
 	"github.com/getlantern/http-proxy-lantern/v2/starbridge"
 	"github.com/getlantern/kcpwrapper"
+	"github.com/getlantern/netx"
 
 	"github.com/xtaci/smux"
 
@@ -183,6 +184,7 @@ type Proxy struct {
 	BroflakeCert string
 	BroflakeKey  string
 
+	DNSServer      string
 	throttleConfig throttle.Config
 	instrument     instrument.Instrument
 }
@@ -277,6 +279,35 @@ func (p *Proxy) ListenAndServe(ctx context.Context) error {
 	// Throttle connections when signaled
 	srv.AddListenerWrappers(lanternlisteners.NewBitrateListener, bwReporting.wrapper)
 
+	if p.DNSServer != "" {
+		log.Debugf("Will resolve DNS using %v", p.DNSServer)
+		host, port, err := net.SplitHostPort(p.DNSServer)
+		if err != nil {
+			log.Fatalf("invalid dns-server address %v: %v", p.DNSServer, err)
+		}
+		r := &net.Resolver{
+			PreferGo: true,
+			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+				log.Debug("Dialing custom resolver")
+				return netx.DialContext(ctx, host, port)
+			},
+		}
+		netx.OverrideResolveIPs(func(host string) ([]net.IP, error) {
+			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer cancel()
+
+			addrs, err := r.LookupIPAddr(ctx, host)
+			if err != nil {
+				return nil, err
+			}
+			ips := make([]net.IP, 0, len(addrs))
+			for _, addr := range addrs {
+				ips = append(ips, addr.IP)
+			}
+			return ips, nil
+		})
+	}
+
 	allListeners := make([]net.Listener, 0)
 	listenerProtocols := make([]string, 0)
 	addListenerIfNecessary := func(proto, addr string, fn listenerBuilderFN) error {
@@ -613,7 +644,7 @@ func (p *Proxy) createFilterChain(bl *blacklist.Blacklist) (filters.Chain, proxy
 	_dialer := preferIPV4Dialer(timeoutToDialOriginSite)
 	dialer := func(ctx context.Context, network, addr string) (net.Conn, error) {
 		// resolve separately so that we can track the DNS resolution time
-		resolvedAddr, resolveErr := net.ResolveTCPAddr(network, addr)
+		resolvedAddr, resolveErr := netx.Resolve(network, addr)
 		if resolveErr != nil {
 			return nil, resolveErr
 		}

diff --git a/throttle_integration_test.go b/throttle_integration_test.go
@@ -17,11 +17,12 @@ import (
 	"testing"
 	"time"
 
-	"github.com/getlantern/golog/testlog"
-	. "github.com/getlantern/waitforserver"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/getlantern/golog/testlog"
+	. "github.com/getlantern/waitforserver"
+
 	"github.com/getlantern/http-proxy-lantern/v2/common"
 	"github.com/getlantern/http-proxy-lantern/v2/internal/testutil"
 	"github.com/getlantern/http-proxy-lantern/v2/throttle"