From 6a014d5e22e6a0b7c1fcb65f59872e4dd1227111 Mon Sep 17 00:00:00 2001
From: Alexander Pevzner <pzz@apevzner.com>
Date: Sat, 27 Apr 2019 19:36:46 +0300
Subject: [PATCH] Add DialWithDialFunc to specify dial function used for both
 control and data connections (#140)

Add DialWithDialFunc to specify dial function used for both control and data connections

If used DialWithNetConn, DialWithNetConn takes precedence for
the control connection, while data connections will be established
using function specified with the DialWithDialFunc option
---
 ftp.go | 36 ++++++++++++++++++++++++++++++------
 1 file changed, 30 insertions(+), 6 deletions(-)

diff --git a/ftp.go b/ftp.go
index 62429f1..9746e05 100644
--- a/ftp.go
+++ b/ftp.go
@@ -54,6 +54,7 @@ type dialOptions struct {
 	disableEPSV bool
 	location    *time.Location
 	debugOutput io.Writer
+	dialFunc    func(network, address string) (net.Conn, error)
 }
 
 // Entry describes a file and is returned by List().
@@ -84,17 +85,23 @@ func Dial(addr string, options ...DialOption) (*ServerConn, error) {
 
 	tconn := do.conn
 	if tconn == nil {
-		ctx := do.context
+		var err error
 
-		if ctx == nil {
-			ctx = context.Background()
+		if do.dialFunc != nil {
+			tconn, err = do.dialFunc("tcp", addr)
+		} else {
+			ctx := do.context
+
+			if ctx == nil {
+				ctx = context.Background()
+			}
+
+			tconn, err = do.dialer.DialContext(ctx, "tcp", addr)
 		}
 
-		conn, err := do.dialer.DialContext(ctx, "tcp", addr)
 		if err != nil {
 			return nil, err
 		}
-		tconn = conn
 	}
 
 	// Use the resolved IP address in case addr contains a domain name
@@ -192,6 +199,18 @@ func DialWithDebugOutput(w io.Writer) DialOption {
 	}}
 }
 
+// DialWithDialFunc returns a DialOption that configures the ServerConn to use the
+// specified function to establish both control and data connections
+//
+// If used together with the DialWithNetConn option, the DialWithNetConn
+// takes precedence for the control connection, while data connections will
+// be established using function specified with the DialWithDialFunc option
+func DialWithDialFunc(f func(network, address string) (net.Conn, error)) DialOption {
+	return DialOption{func(do *dialOptions) {
+		do.dialFunc = f
+	}}
+}
+
 // Connect is an alias to Dial, for backward compatibility
 func Connect(addr string) (*ServerConn, error) {
 	return Dial(addr)
@@ -387,7 +406,12 @@ func (c *ServerConn) openDataConn() (net.Conn, error) {
 		return nil, err
 	}
 
-	return c.options.dialer.Dial("tcp", net.JoinHostPort(host, strconv.Itoa(port)))
+	addr := net.JoinHostPort(host, strconv.Itoa(port))
+	if c.options.dialFunc != nil {
+		return c.options.dialFunc("tcp", addr)
+	}
+
+	return c.options.dialer.Dial("tcp", addr)
 }
 
 // cmd is a helper function to execute a command and check for the expected FTP