server/proxy: allow customizing the DNS server (#3)

server/example.env

@@ -15,3 +15,6 @@ EMISSARY_ALLOWED_PROXY_TARGETS='[{ "host": "www.google.com", "port": 443 }]'
 #  - `data` is the key base64 encoded
 #  - Every key in this list can be used to authenticate against emissary
 EMISSARY_AUTH_KEYS='[{ "kid": 1, "data": "c29tZSBzdXBlciBzZWNyZXQgcmFuZG9taXNlZCBrZXkgaGVyZS4KClRoaXMgaXMgc2ltcGx5IGFuIGV4YW1wbGUga2V5" }]'
+
+# The DNS servers to use, as a space-separated list.
+EMISSARY_DNS_SERVERS='1.1.1.1 8.8.8.8'

server/proxy/allow_list.go

@@ -1,9 +1,10 @@
 package proxy
 
 import (
+	"context"
+
 	"github.com/armon/go-socks5"
 	"github.com/rs/zerolog/log"
-	"golang.org/x/net/context"
 )
 
 type AllowedHost struct {

server/proxy/config.go

@@ -17,6 +17,7 @@ type Config struct {
 	TcpPort             int                 // What port should this server listen for raw TCP connections on (0 == disabled)
 	AuthKeys            auth.Keys           // What auth keys can be used when talking with this Emissary server
 	AllowedProxyTargets AllowedProxyTargets // What proxy targets are allowed through this Emissary server
+	DNSServers          []string            // The DNS server IPs to use; nil means the system default
 }
 
 // LoadConfig performs setup for the proxy layer and returns an error if we cannot initialise
@@ -78,5 +79,6 @@ func LoadConfig(_ context.Context) (*Config, error) {
 		TcpPort:             viper.GetInt("tcp_port"),
 		AuthKeys:            authKeys,
 		AllowedProxyTargets: allowedProxyTargets,
+		DNSServers:          viper.GetStringSlice("dns_servers"),
 	}, nil
 }

server/proxy/proxy.go

@@ -1,9 +1,11 @@
 package proxy
 
 import (
+	"context"
 	"crypto/rand"
 	golog "log"
 	"net"
+	"time"
 
 	"github.com/armon/go-socks5"
 	"github.com/cockroachdb/errors"
@@ -35,11 +37,17 @@ func (cfg *Config) ServeConn(conn net.Conn) error {
 		return errors.Wrap(err, "unable to send connect message")
 	}
 
+	var resolver socks5.NameResolver = socks5.DNSResolver{}
+	if len(cfg.DNSServers) > 0 {
+		resolver = customDNSResolver{ServerIPs: cfg.DNSServers, Fallback: socks5.DNSResolver{}}
+	}
+
 	// Set up our SOCKS5 server
 	server, err := socks5.New(&socks5.Config{
 		AuthMethods: newAuthenticator(cfg, nonce),
 		Rules:       cfg.AllowedProxyTargets,
 		Logger:      golog.New(log.Logger, "", golog.Lshortfile),
+		Resolver:    resolver,
 	})
 	if err != nil {
 		return errors.Wrap(err, "unable to setup socks5 proxy server")
@@ -52,3 +60,49 @@ func (cfg *Config) ServeConn(conn net.Conn) error {
 
 	return nil
 }
+
+type customDNSResolver struct {
+	// ServerIPs are the IPs to dial to do DNS lookups, in order.
+	ServerIPs []string
+	// Fallback specifies the fallback name resolver to use, if the given resolvers don't return anything.
+	Fallback socks5.NameResolver
+}
+
+func (d customDNSResolver) Resolve(ctx context.Context, name string) (context.Context, net.IP, error) {
+	var lastErr error
+	for _, serverIP := range d.ServerIPs {
+		r := &net.Resolver{
+			PreferGo: true,
+			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+				return (&net.Dialer{}).DialContext(ctx, network, serverIP+":53")
+			},
+		}
+
+		// Use a context with a short timeout so we don't block forever.
+		lookupCtx, cancel := context.WithTimeout(ctx, 1*time.Second)
+		addrs, err := r.LookupIPAddr(lookupCtx, name)
+		cancel()
+		if len(addrs) > 0 {
+			return ctx, addrs[0].IP, nil
+		}
+		if err != nil {
+			lastErr = err
+			continue
+		}
+	}
+
+	if d.Fallback != nil {
+		return d.Fallback.Resolve(ctx, name)
+	}
+
+	// We failed; return the last error, if any.
+	if lastErr != nil {
+		return ctx, nil, lastErr
+	}
+	// Otherwise replicate the default behavior of net.Resolver when there are no addresses.
+	err := &net.AddrError{
+		Addr: name,
+		Err:  "no suitable address found",
+	}
+	return ctx, nil, err
+}

server/proxy/proxy_test.go

@@ -0,0 +1,45 @@
+package proxy
+
+import (
+	"context"
+	"testing"
+
+	"github.com/armon/go-socks5"
+)
+
+func TestCustomDNSServer(t *testing.T) {
+	t.Parallel()
+	dns := customDNSResolver{ServerIPs: []string{"1.1.1.1"}}
+	_, ip, err := dns.Resolve(context.Background(), "google.com")
+	if err != nil {
+		t.Fatalf("unable to resolve with custom resolver: %v", err)
+	} else if len(ip) == 0 {
+		t.Fatalf("got zero ip: %s", ip)
+	}
+}
+
+func TestCustomDNSServerMultiple(t *testing.T) {
+	t.Parallel()
+
+	// Start with an invalid DNS server to ensure we try the second, valid one.
+	dns := customDNSResolver{ServerIPs: []string{"127.127.127.127", "1.1.1.1"}}
+	_, ip, err := dns.Resolve(context.Background(), "google.com")
+	if err != nil {
+		t.Fatalf("unable to resolve with custom resolver: %v", err)
+	} else if len(ip) == 0 {
+		t.Fatalf("got zero ip: %s", ip)
+	}
+}
+
+func TestCustomDNSServerFallback(t *testing.T) {
+	t.Parallel()
+
+	// Start with an invalid DNS server to ensure we try the fallback
+	dns := customDNSResolver{ServerIPs: []string{"127.127.127.127"}, Fallback: socks5.DNSResolver{}}
+	_, ip, err := dns.Resolve(context.Background(), "google.com")
+	if err != nil {
+		t.Fatalf("unable to resolve with custom resolver: %v", err)
+	} else if len(ip) == 0 {
+		t.Fatalf("got zero ip: %s", ip)
+	}
+}