浏览代码

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
Alexander Pevzner 6 年之前
父节点
当前提交
6a014d5e22
共有 1 个文件被更改,包括 30 次插入6 次删除
  1. 30 6
      ftp.go

+ 30 - 6
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