From 6f9fd4431aff58576604932afa5ed0cc26023de3 Mon Sep 17 00:00:00 2001
From: Phu Kieu <pskieu@gmail.com>
Date: Wed, 12 Oct 2016 08:02:39 -0700
Subject: [PATCH] Use same SSH tunnel as kubelet

---
 pkg/ssh/ssh.go | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/pkg/ssh/ssh.go b/pkg/ssh/ssh.go
index dcb039c51c4..5b0db13f875 100644
--- a/pkg/ssh/ssh.go
+++ b/pkg/ssh/ssh.go
@@ -32,6 +32,7 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"strings"
 	"sync"
 	"time"
 
@@ -391,19 +392,27 @@ func (l *SSHTunnelList) Dial(net, addr string) (net.Conn, error) {
 	defer func() {
 		glog.Infof("[%x: %v] Dialed in %v.", id, addr, time.Now().Sub(start))
 	}()
-	tunnel, err := l.pickRandomTunnel()
+	tunnel, err := l.pickTunnel(strings.Split(addr, ":")[0])
 	if err != nil {
 		return nil, err
 	}
 	return tunnel.Dial(net, addr)
 }
 
-func (l *SSHTunnelList) pickRandomTunnel() (tunnel, error) {
+func (l *SSHTunnelList) pickTunnel(addr string) (tunnel, error) {
 	l.tunnelsLock.Lock()
 	defer l.tunnelsLock.Unlock()
 	if len(l.entries) == 0 {
 		return nil, fmt.Errorf("No SSH tunnels currently open. Were the targets able to accept an ssh-key for user %q?", l.user)
 	}
+	// Prefer same tunnel as kubelet
+	// TODO: Change l.entries to a map of address->tunnel
+	for _, entry := range l.entries {
+		if entry.Address == addr {
+			return entry.Tunnel, nil
+		}
+	}
+	glog.Warningf("SSH tunnel not found for address %q, picking random node", addr)
 	n := mathrand.Intn(len(l.entries))
 	return l.entries[n].Tunnel, nil
 }