From ed93932c3bb6659d637090c9c9c6e0e57dd4d2fb Mon Sep 17 00:00:00 2001
From: Shaya Potter <shaya@redislabs.com>
Date: Sun, 5 Feb 2023 23:59:55 +0200
Subject: [PATCH] add the ability to control socks5 resolution

not all servers support remote resolution, and require one to send only ip addresses

this add a new SOCKS5 dialer creation function that lets one pass in a custom *net.Resolver.

if a resolver is set on the dialer, it will use it it convert non ip addresses into ip addresses.  For instance, can simply be used with net.DefaultResolver to do local resolution
---
 internal/socks/client.go | 10 ++++++++++
 internal/socks/socks.go  |  6 ++++--
 proxy/socks5.go          | 12 ++++++++++--
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/internal/socks/client.go b/internal/socks/client.go
index 3d6f516a59..65ddd06914 100644
--- a/internal/socks/client.go
+++ b/internal/socks/client.go
@@ -23,6 +23,16 @@ func (d *Dialer) connect(ctx context.Context, c net.Conn, address string) (_ net
 	if err != nil {
 		return nil, err
 	}
+
+	if d.Resolver != nil {
+		if ip := net.ParseIP(host); ip == nil {
+			addresses, err := d.Resolver.LookupHost(ctx, host)
+			if err != nil {
+				return nil, err
+			}
+			host = addresses[0]
+		}
+	}
 	if deadline, ok := ctx.Deadline(); ok && !deadline.IsZero() {
 		c.SetDeadline(deadline)
 		defer c.SetDeadline(noDeadline)
diff --git a/internal/socks/socks.go b/internal/socks/socks.go
index 97db2340ec..84a8abba72 100644
--- a/internal/socks/socks.go
+++ b/internal/socks/socks.go
@@ -136,6 +136,8 @@ type Dialer struct {
 	// function. It must be non-nil when AuthMethods is not empty.
 	// It must return an error when the authentication is failed.
 	Authenticate func(context.Context, io.ReadWriter, AuthMethod) error
+
+	Resolver *net.Resolver
 }
 
 // DialContext connects to the provided address on the provided
@@ -266,8 +268,8 @@ func (d *Dialer) pathAddrs(address string) (proxy, dst net.Addr, err error) {
 
 // NewDialer returns a new Dialer that dials through the provided
 // proxy server's network and address.
-func NewDialer(network, address string) *Dialer {
-	return &Dialer{proxyNetwork: network, proxyAddress: address, cmd: CmdConnect}
+func NewDialer(network, address string, resolver *net.Resolver) *Dialer {
+	return &Dialer{proxyNetwork: network, proxyAddress: address, cmd: CmdConnect, Resolver: resolver}
 }
 
 const (
diff --git a/proxy/socks5.go b/proxy/socks5.go
index c91651f96d..6e3c2944a6 100644
--- a/proxy/socks5.go
+++ b/proxy/socks5.go
@@ -14,8 +14,9 @@ import (
 // SOCKS5 returns a Dialer that makes SOCKSv5 connections to the given
 // address with an optional username and password.
 // See RFC 1928 and RFC 1929.
-func SOCKS5(network, address string, auth *Auth, forward Dialer) (Dialer, error) {
-	d := socks.NewDialer(network, address)
+
+func socks5Internal(network, address string, auth *Auth, forward Dialer, resolver *net.Resolver) (Dialer, error) {
+	d := socks.NewDialer(network, address, resolver)
 	if forward != nil {
 		if f, ok := forward.(ContextDialer); ok {
 			d.ProxyDial = func(ctx context.Context, network string, address string) (net.Conn, error) {
@@ -40,3 +41,10 @@ func SOCKS5(network, address string, auth *Auth, forward Dialer) (Dialer, error)
 	}
 	return d, nil
 }
+func SOCKS5(network, address string, auth *Auth, forward Dialer) (Dialer, error) {
+	return socks5Internal(network, address, auth, forward, nil)
+}
+
+func SOCKS5WithResolver(network, address string, auth *Auth, forward Dialer, resolver *net.Resolver) (Dialer, error) {
+	return socks5Internal(network, address, auth, forward, resolver)
+}