From 0b72453262711ec359004904363a8f2817a08b1f Mon Sep 17 00:00:00 2001
From: Chunyi Lyu <clyu@pivotal.io>
Date: Fri, 10 Aug 2018 12:08:40 +0100
Subject: [PATCH] Create create a dailFunc in ssh package once

[#159619865]

Signed-off-by: Jake Klein <jklein@pivotal.io>
---
 ssh/connection.go | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/ssh/connection.go b/ssh/connection.go
index 3c1fbde5f..2e65d0142 100644
--- a/ssh/connection.go
+++ b/ssh/connection.go
@@ -53,6 +53,7 @@ func NewConnectionWithServerAliveInterval(hostName, userName, privateKey string,
 		},
 		logger:              logger,
 		serverAliveInterval: serverAliveInterval,
+		dailFunc:            createDialFunc(),
 	}
 
 	return conn, nil
@@ -63,6 +64,7 @@ type Connection struct {
 	sshConfig           *ssh.ClientConfig
 	logger              Logger
 	serverAliveInterval time.Duration
+	dailFunc            boshhttp.DialFunc
 }
 
 func (c Connection) Run(cmd string) (stdout, stderr []byte, exitCode int, err error) {
@@ -106,14 +108,7 @@ func (w *sessionClosingOnErrorWriter) Write(data []byte) (int, error) {
 }
 
 func (c Connection) getConnection() (*ssh.Client, error) {
-	dialFunc := net.Dial
-
-	if os.Getenv("BOSH_ALL_PROXY") != "" {
-		socksProxy := proxy.NewSocks5Proxy(proxy.NewHostKey(), log.New(os.Stdout, "sock5-proxy", log.LstdFlags))
-		dialFunc = boshhttp.SOCKS5DialFuncFromEnvironment(net.Dial, socksProxy)
-	}
-
-	conn, err := dialFunc("tcp", c.host)
+	conn, err := c.dailFunc("tcp", c.host)
 	if err != nil {
 		return nil, err
 	}
@@ -126,6 +121,15 @@ func (c Connection) getConnection() (*ssh.Client, error) {
 	return ssh.NewClient(client, chans, reqs), nil
 }
 
+func createDialFunc() boshhttp.DialFunc {
+	dialFunc := net.Dial
+	if os.Getenv("BOSH_ALL_PROXY") != "" {
+		socksProxy := proxy.NewSocks5Proxy(proxy.NewHostKey(), log.New(os.Stdout, "sock5-proxy", log.LstdFlags))
+		dialFunc = boshhttp.SOCKS5DialFuncFromEnvironment(net.Dial, socksProxy)
+	}
+	return dialFunc
+}
+
 func (c Connection) runInSession(cmd string, stdout, stderr io.Writer, stdin io.Reader) (int, error) {
 	connection, err := c.getConnection()
 	if err != nil {