From cd1e5cf6109c4ec77a563ce73c91934cefeca3bc Mon Sep 17 00:00:00 2001
From: galal-hussein <hussein.galal.ahmed.11@gmail.com>
Date: Sat, 3 Feb 2018 03:04:53 +0200
Subject: [PATCH] Custom dialers and remove local fixes

---
 cmd/remove.go           | 13 ++++++++++++-
 hosts/dialer.go         |  2 +-
 hosts/local.go          |  2 +-
 services/healthcheck.go |  6 +++---
 4 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/cmd/remove.go b/cmd/remove.go
index e2ea15e1..820b2629 100644
--- a/cmd/remove.go
+++ b/cmd/remove.go
@@ -98,6 +98,17 @@ func clusterRemoveFromCli(ctx *cli.Context) error {
 
 func clusterRemoveLocal(ctx *cli.Context) error {
 	var rkeConfig *v3.RancherKubernetesEngineConfig
-	rkeConfig = cluster.GetLocalRKEConfig()
+	clusterFile, filePath, err := resolveClusterFile(ctx)
+	if err != nil {
+		log.Infof(context.Background(), "Failed to resolve cluster file, using default cluster instead")
+		rkeConfig = cluster.GetLocalRKEConfig()
+	} else {
+		clusterFilePath = filePath
+		rkeConfig, err = cluster.ParseConfig(clusterFile)
+		if err != nil {
+			return fmt.Errorf("Failed to parse cluster file: %v", err)
+		}
+		rkeConfig.Nodes = []v3.RKEConfigNode{*cluster.GetLocalRKENodeConfig()}
+	}
 	return ClusterRemove(context.Background(), rkeConfig, nil, true, "")
 }
diff --git a/hosts/dialer.go b/hosts/dialer.go
index 47801b66..ae442767 100644
--- a/hosts/dialer.go
+++ b/hosts/dialer.go
@@ -73,7 +73,7 @@ func (d *dialer) DialLocalConn(network, addr string) (net.Conn, error) {
 	if err != nil {
 		return nil, fmt.Errorf("Failed to dial ssh using address [%s]: %v", sshAddr, err)
 	}
-	remote, err := conn.Dial("tcp", fmt.Sprintf("localhost:%d", d.host.LocalConnPort))
+	remote, err := conn.Dial(network, addr)
 	if err != nil {
 		return nil, fmt.Errorf("Failed to dial to Local Port [%d] on host [%s]: %v", d.host.LocalConnPort, d.host.Address, err)
 	}
diff --git a/hosts/local.go b/hosts/local.go
index 5364ee48..88f3e7c8 100644
--- a/hosts/local.go
+++ b/hosts/local.go
@@ -13,7 +13,7 @@ func LocalHealthcheckFactory(h *Host) (func(network, address string) (net.Conn,
 }
 
 func (d *dialer) DialHealthcheckLocally(network, addr string) (net.Conn, error) {
-	conn, err := net.Dial("tcp", fmt.Sprintf("localhost:%d", d.host.LocalConnPort))
+	conn, err := net.Dial(network, addr)
 	if err != nil {
 		return nil, fmt.Errorf("Failed to dial address [%s]: %v", d.host.Address, err)
 	}
diff --git a/services/healthcheck.go b/services/healthcheck.go
index d4d4a80a..11f883ff 100644
--- a/services/healthcheck.go
+++ b/services/healthcheck.go
@@ -27,7 +27,7 @@ func runHealthcheck(ctx context.Context, host *hosts.Host, port int, useTLS bool
 		return fmt.Errorf("Failed to initiate new HTTP client for service [%s] for host [%s]", serviceName, host.Address)
 	}
 	for retries := 0; retries < 10; retries++ {
-		if err = getHealthz(client, useTLS, serviceName, host.Address); err != nil {
+		if err = getHealthz(client, useTLS, serviceName, host.Address, port); err != nil {
 			logrus.Debugf("[healthcheck] %v", err)
 			time.Sleep(5 * time.Second)
 			continue
@@ -58,12 +58,12 @@ func getHealthCheckHTTPClient(host *hosts.Host, port int, localConnDialerFactory
 	}, nil
 }
 
-func getHealthz(client *http.Client, useTLS bool, serviceName, hostAddress string) error {
+func getHealthz(client *http.Client, useTLS bool, serviceName, hostAddress string, port int) error {
 	proto := HTTPProtoPrefix
 	if useTLS {
 		proto = HTTPSProtoPrefix
 	}
-	resp, err := client.Get(fmt.Sprintf("%s%s%s", proto, HealthzAddress, HealthzEndpoint))
+	resp, err := client.Get(fmt.Sprintf("%s%s:%d%s", proto, HealthzAddress, port, HealthzEndpoint))
 	if err != nil {
 		return fmt.Errorf("Failed to check %s for service [%s] on host [%s]: %v", HealthzEndpoint, serviceName, hostAddress, err)
 	}