chore: inject Picker to resolvers

qdm12 · Jun 24, 2022 · 81f9533 · 81f9533
1 parent 3a2fbd7
commit 81f9533
Show file tree

Hide file tree

Showing 9 changed files with 43 additions and 21 deletions.
diff --git a/pkg/doh/dial.go b/pkg/doh/dial.go
@@ -7,7 +7,6 @@ import (
 	"net"
 	"sync"
 
-	"github.com/qdm12/dns/v2/internal/picker"
 	"github.com/qdm12/dns/v2/internal/server"
 	"github.com/qdm12/dns/v2/pkg/dot"
 	"github.com/qdm12/dns/v2/pkg/provider"
@@ -48,7 +47,7 @@ func newDoHDial(settings ResolverSettings) (
 		},
 	}
 
-	picker := picker.New()
+	picker := settings.Picker
 
 	return func(ctx context.Context, _, _ string) (conn net.Conn, err error) {
 		// Pick DoH server pseudo-randomly from the chosen providers

diff --git a/pkg/doh/interfaces.go b/pkg/doh/interfaces.go
@@ -0,0 +1,9 @@
+package doh
+
+import (
+	"github.com/qdm12/dns/v2/pkg/provider"
+)
+
+type Picker interface {
+	DoHServer(servers []provider.DoHServer) provider.DoHServer
+}
diff --git a/pkg/doh/settings.go b/pkg/doh/settings.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/qdm12/dns/v2/internal/config/defaults"
+	"github.com/qdm12/dns/v2/internal/picker"
 	"github.com/qdm12/dns/v2/pkg/cache"
 	cachenoop "github.com/qdm12/dns/v2/pkg/cache/noop"
 	"github.com/qdm12/dns/v2/pkg/doh/metrics"
@@ -53,6 +54,11 @@ type ResolverSettings struct {
 	// Metrics is the metrics interface to record metric data.
 	// It defaults to a No-Op metrics implementation.
 	Metrics metrics.DialMetrics
+	// Picker is the picker to use for each upstream call to pick
+	// a server from a pool of servers. It must be thread safe.
+	// It defaults to a fast thread safe pseudo random picker
+	// with uniform distribution.
+	Picker Picker
 }
 
 type SelfDNS struct {
@@ -103,6 +109,10 @@ func (s *ResolverSettings) SetDefaults() {
 	if s.Metrics == nil {
 		s.Metrics = metricsnoop.New()
 	}
+
+	if s.Picker == nil {
+		s.Picker = picker.New()
+	}
 }
 
 func (s *SelfDNS) SetDefaults() {

diff --git a/pkg/doh/settings_test.go b/pkg/doh/settings_test.go
@@ -4,6 +4,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/qdm12/dns/v2/internal/picker"
 	cache "github.com/qdm12/dns/v2/pkg/cache/noop"
 	metrics "github.com/qdm12/dns/v2/pkg/doh/metrics/noop"
 	"github.com/qdm12/dns/v2/pkg/filter/mapfilter"
@@ -21,6 +22,7 @@ func Test_ServerSettings_SetDefaults(t *testing.T) {
 	filter := mapfilter.New(mapfilter.Settings{})
 	metrics := metrics.New()
 	logger := log.New()
+	picker := picker.New()
 
 	s := ServerSettings{
 		Cache:   cache,
@@ -30,6 +32,7 @@ func Test_ServerSettings_SetDefaults(t *testing.T) {
 		Resolver: ResolverSettings{
 			Warner:  logger,
 			Metrics: metrics,
+			Picker:  picker,
 		},
 	}
 	s.SetDefaults()
@@ -58,6 +61,7 @@ func Test_ServerSettings_SetDefaults(t *testing.T) {
 			Timeout: 5 * time.Second,
 			Warner:  logger,
 			Metrics: metrics,
+			Picker:  picker,
 		},
 		ListeningAddress: ":53",
 	}

diff --git a/pkg/dot/dial.go b/pkg/dot/dial.go
@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"net"
 
-	"github.com/qdm12/dns/v2/internal/picker"
 	"github.com/qdm12/dns/v2/internal/server"
 	"github.com/qdm12/dns/v2/pkg/dot/metrics"
 	"github.com/qdm12/dns/v2/pkg/log"
@@ -27,7 +26,7 @@ func newDoTDial(settings ResolverSettings) (
 		Timeout: settings.Timeout,
 	}
 
-	picker := picker.New()
+	picker := settings.Picker
 
 	return func(ctx context.Context, _, _ string) (net.Conn, error) {
 		serverName, serverAddress := pickNameAddress(picker,
@@ -75,7 +74,7 @@ func settingsToServers(settings ResolverSettings) (
 	return dotServers, dnsServers, nil
 }
 
-func pickNameAddress(picker picker.DoT, servers []provider.DoTServer,
+func pickNameAddress(picker Picker, servers []provider.DoTServer,
 	ipv6 bool) (name, address string) {
 	server := picker.DoTServer(servers)
 	ip := picker.DoTIP(server, ipv6)
@@ -85,7 +84,7 @@ func pickNameAddress(picker picker.DoT, servers []provider.DoTServer,
 
 func onDialError(ctx context.Context, dialErr error,
 	dotName, dotAddress string, dialer *net.Dialer,
-	picker picker.DNS, ipv6 bool, dnsServers []provider.DNSServer,
+	picker Picker, ipv6 bool, dnsServers []provider.DNSServer,
 	warner log.Warner, metrics metrics.DialMetrics) (
 	conn net.Conn, err error) {
 	warner.Warn(dialErr.Error())
@@ -101,7 +100,7 @@ func onDialError(ctx context.Context, dialErr error,
 }
 
 func dialPlaintext(ctx context.Context, dialer *net.Dialer,
-	picker picker.DNS, ipv6 bool, dnsServers []provider.DNSServer,
+	picker Picker, ipv6 bool, dnsServers []provider.DNSServer,
 	warner log.Warner, metrics metrics.DialDNSMetrics) (
 	conn net.Conn, err error) {
 	dnsServer := picker.DNSServer(dnsServers)

diff --git a/pkg/dot/dial_test.go b/pkg/dot/dial_test.go
@@ -36,7 +36,7 @@ func Test_settingsToServers(t *testing.T) {
 	}, dnsServers)
 }
 
-//go:generate mockgen -destination=mock_picker_test.go -package $GOPACKAGE -mock_names Interface=MockPicker github.com/qdm12/dns/v2/internal/picker Interface
+//go:generate mockgen -destination=mocks_test.go -package $GOPACKAGE . Picker
 //go:generate mockgen -destination=mock_dot_metrics_test.go -package $GOPACKAGE -mock_names Interface=MockDoTMetrics github.com/qdm12/dns/v2/pkg/dot/metrics Interface
 //go:generate mockgen -destination=mock_warner_test.go -package $GOPACKAGE github.com/qdm12/dns/v2/pkg/log Warner
 

diff --git a/internal/picker/interface.go → pkg/dot/interfaces.go b/internal/picker/interface.go → pkg/dot/interfaces.go
@@ -1,24 +1,15 @@
-package picker
+package dot
 
 import (
 	"net"
 
 	"github.com/qdm12/dns/v2/pkg/provider"
 )
 
-type Interface interface {
+type Picker interface {
 	IP(ips []net.IP) net.IP
 	DNSServer(servers []provider.DNSServer) provider.DNSServer
 	DNSIP(server provider.DNSServer, ipv6 bool) net.IP
-	DoT
-}
-
-type DoT interface {
 	DoTServer(servers []provider.DoTServer) provider.DoTServer
 	DoTIP(server provider.DoTServer, ipv6 bool) net.IP
 }
-
-type DNS interface {
-	DNSServer(servers []provider.DNSServer) provider.DNSServer
-	DNSIP(server provider.DNSServer, ipv6 bool) net.IP
-}
diff --git a/pkg/dot/mock_picker_test.go → pkg/dot/mocks_test.go b/pkg/dot/mock_picker_test.go → pkg/dot/mocks_test.go
diff --git a/pkg/dot/settings.go b/pkg/dot/settings.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/qdm12/dns/v2/internal/config/defaults"
+	"github.com/qdm12/dns/v2/internal/picker"
 	"github.com/qdm12/dns/v2/pkg/cache"
 	cachenoop "github.com/qdm12/dns/v2/pkg/cache/noop"
 	"github.com/qdm12/dns/v2/pkg/dot/metrics"
@@ -56,6 +57,11 @@ type ResolverSettings struct {
 	// Metrics is the metrics interface to record metric data.
 	// It defaults to a No-Op metrics implementation.
 	Metrics metrics.DialMetrics
+	// Picker is the picker to use for each upstream call to pick
+	// a server and/or IP address. It must be thread safe.
+	// It defaults to a fast thread safe pseudo random picker
+	// with uniform distribution.
+	Picker Picker
 }
 
 func (s *ServerSettings) SetDefaults() {
@@ -97,6 +103,10 @@ func (s *ResolverSettings) SetDefaults() {
 	if s.Metrics == nil {
 		s.Metrics = metricsnoop.New()
 	}
+
+	if s.Picker == nil {
+		s.Picker = picker.New()
+	}
 }
 
 var (