chore: refactor dns

dns/resolver/factory.go

@@ -0,0 +1,36 @@
+package resolver
+
+import (
+	"context"
+	"net"
+)
+
+type Factory interface {
+	Get(ctx context.Context, host string, qTypes []uint16) (Resolver, error)
+}
+
+type FactoryFunc func(ctx context.Context, host string, qTypes []uint16) (Resolver, error)
+
+func (f FactoryFunc) Get(ctx context.Context, host string, qTypes []uint16) (Resolver, error) {
+	return f(ctx, host, qTypes)
+}
+
+type FactoryResolver struct {
+	f Factory
+}
+
+func NewFactoryResolver(f Factory) *FactoryResolver {
+	return &FactoryResolver{f}
+}
+
+func (c *FactoryResolver) Resolve(ctx context.Context, host string, qTypes []uint16) ([]net.IPAddr, error) {
+	resolver, err := c.f.Get(ctx, host, qTypes)
+	if err != nil {
+		return nil, err
+	}
+	return resolver.Resolve(ctx, host, qTypes)
+}
+
+func (c *FactoryResolver) String() string {
+	return "FactoryResolver"
+}

dns/resolver/handler/err.go

@@ -0,0 +1,27 @@
+package handler
+
+import (
+	"context"
+	"fmt"
+	"net"
+
+	"github.com/xvzc/SpoofDPI/dns/resolver"
+)
+
+type ErrorHandler struct {
+}
+
+func NewErrorHandler() *ErrorHandler {
+	return &ErrorHandler{}
+}
+
+func (e *ErrorHandler) DoHandle(ctx context.Context, host string, qTypes []uint16, next resolver.Resolver) ([]net.IPAddr, error) {
+	addrs, err := next.Resolve(ctx, host, qTypes)
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", next, err)
+	}
+	if len(addrs) == 0 {
+		return nil, fmt.Errorf("could not resolve %s using %s", host, next)
+	}
+	return addrs, nil
+}

dns/resolver/handler/handler.go

@@ -0,0 +1,37 @@
+package handler
+
+import (
+	"context"
+	"fmt"
+	"net"
+
+	"github.com/xvzc/SpoofDPI/dns/resolver"
+)
+
+type DnsHandler interface {
+	DoHandle(ctx context.Context, host string, qTypes []uint16, next resolver.Resolver) ([]net.IPAddr, error)
+}
+type DnsHandlerFunc func(ctx context.Context, host string, qTypes []uint16, r resolver.Resolver) ([]net.IPAddr, error)
+type wrappedResolver struct {
+	next resolver.Resolver
+	h    DnsHandler
+}
+
+func (d *wrappedResolver) Resolve(ctx context.Context, host string, qTypes []uint16) ([]net.IPAddr, error) {
+	return d.h.DoHandle(ctx, host, qTypes, d.next)
+}
+
+func Apply(r resolver.Resolver, h ...DnsHandler) resolver.Resolver {
+	for i := len(h) - 1; i >= 0; i-- {
+		r = wrap(r, h[i])
+	}
+	return r
+}
+
+func wrap(r resolver.Resolver, handler DnsHandler) resolver.Resolver {
+	return &wrappedResolver{r, handler}
+}
+
+func (d *wrappedResolver) String() string {
+	return fmt.Sprintf("DNSWrapper(%v)", d.next)
+}

dns/resolver/handler/logging.go

@@ -0,0 +1,37 @@
+package handler
+
+import (
+	"context"
+	"net"
+	"time"
+
+	"github.com/xvzc/SpoofDPI/dns/resolver"
+	"github.com/xvzc/SpoofDPI/util"
+	"github.com/xvzc/SpoofDPI/util/log"
+)
+
+type LoggingHandler struct {
+}
+
+const scopeDNS = "DNS"
+
+func NewLoggingHandler() *LoggingHandler {
+	return &LoggingHandler{}
+}
+
+func (l *LoggingHandler) DoHandle(ctx context.Context, host string, qTypes []uint16, resolver resolver.Resolver) ([]net.IPAddr, error) {
+	ctx = util.GetCtxWithScope(ctx, scopeDNS)
+	logger := log.GetCtxLogger(ctx)
+	logger.Debug().Msgf("resolving %s using %s", host, resolver)
+	t := time.Now()
+	addrs, err := resolver.Resolve(ctx, host, qTypes)
+	if err != nil {
+		logger.Debug().Msgf("failed to resolve %s using %s", host, resolver)
+		return nil, err
+	}
+	if len(addrs) > 0 {
+		d := time.Since(t).Milliseconds()
+		logger.Debug().Msgf("resolved %s from %s in %d ms", addrs[0].String(), host, d)
+	}
+	return addrs, nil
+}

dns/resolver/resolver.go

@@ -12,6 +12,10 @@ import (
 	"github.com/xvzc/SpoofDPI/dns/addrselect"
 )
 
+type Resolver interface {
+	Resolve(ctx context.Context, host string, qTypes []uint16) ([]net.IPAddr, error)
+}
+
 type exchangeFunc = func(ctx context.Context, msg *dns.Msg) (*dns.Msg, error)
 
 type DNSResult struct {

proxy/dns.go

@@ -0,0 +1,63 @@
+package proxy
+
+import (
+	"context"
+	"net"
+
+	"github.com/xvzc/SpoofDPI/dns/resolver"
+	dnshandler "github.com/xvzc/SpoofDPI/dns/resolver/handler"
+	"github.com/xvzc/SpoofDPI/util"
+	"github.com/xvzc/SpoofDPI/util/log"
+)
+
+type proxyDnsFactory struct {
+	systemClient       *resolver.SystemResolver
+	generalClient      *resolver.GeneralResolver
+	dohClient          *resolver.DOHResolver
+	shouldUseSystemDns func(host string) bool
+	enableDOH          bool
+}
+
+func newProxyDnsResolver(dnsAddr string, dnsPort string, enableDOH bool, shouldUseSystemResolver func(host string) bool) resolver.Resolver {
+	factoryResolver := resolver.NewFactoryResolver(
+		newProxyDnsFactory(
+			dnsAddr,
+			dnsPort,
+			shouldUseSystemResolver,
+			enableDOH,
+		),
+	)
+	return dnshandler.Apply(factoryResolver, dnshandler.NewLoggingHandler(), dnshandler.NewErrorHandler())
+}
+
+func newProxyDnsFactory(
+	dnsAddr string,
+	dnsPort string,
+	shouldUseSystemDnsFunc func(host string) bool,
+	enableDOH bool) *proxyDnsFactory {
+	return &proxyDnsFactory{
+		systemClient:       resolver.NewSystemResolver(),
+		dohClient:          resolver.NewDOHResolver(dnsAddr),
+		generalClient:      resolver.NewGeneralResolver(net.JoinHostPort(dnsAddr, dnsPort)),
+		enableDOH:          enableDOH,
+		shouldUseSystemDns: shouldUseSystemDnsFunc,
+	}
+}
+
+func (p *proxyDnsFactory) Get(ctx context.Context, host string, _ []uint16) (resolver.Resolver, error) {
+	l := util.GetCtxWithScope(ctx, scopeProxy)
+	logger := log.GetCtxLogger(l)
+	r, _ := p.findInternal(host)
+	logger.Debug().Msgf("proxy dns factory returns: %v", r)
+	return r, nil
+}
+
+func (p *proxyDnsFactory) findInternal(host string) (resolver.Resolver, error) {
+	if p.shouldUseSystemDns(host) {
+		return p.systemClient, nil
+	}
+	if p.enableDOH {
+		return p.dohClient, nil
+	}
+	return p.generalClient, nil
+}