Move saving of hosts file to after all ports are forwarded

This fixes #74
txn2 · Aug 21, 2019 · cb2bed3 · cb2bed3
1 parent 6a86353
commit cb2bed3
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 19 deletions.
diff --git a/cmd/kubefwd/services/services.go b/cmd/kubefwd/services/services.go
@@ -217,10 +217,17 @@ Try:
 			}
 		}
 
+		log.Printf("Saving hosts file\n")
+		err = hostFile.Save()
+		if err != nil {
+			log.Error("Error saving hosts file", err)
+		}
+
 		wg.Wait()
 
 		log.Printf("Done...\n")
 
+		log.Printf("Restoring hosts file\n")
 		err = hostFile.Save()
 		if err != nil {
 			log.Fatalf("Error saving hostfile: %s\n", err.Error())

diff --git a/pkg/fwdport/fwdport.go b/pkg/fwdport/fwdport.go
@@ -10,8 +10,6 @@ import (
 	"strconv"
 	"strings"
 
-	"github.com/prometheus/common/log"
-
 	"github.com/txn2/txeh"
 
 	"github.com/txn2/kubefwd/pkg/fwdpub"
@@ -113,28 +111,22 @@ func PortForward(pfo *PortForwardOpts) error {
 
 	}
 
-	err = pfo.Hostfile.Save()
-	if err != nil {
-		return err
+	cleanupHostfile := func() {
+		if pfo.Remote == false {
+			if pfo.Domain != "" {
+		 		pfo.Hostfile.RemoveHost(localServiceName + "." + pfo.Domain)
+		 		pfo.Hostfile.RemoveHost(nsServiceName + "." + pfo.Domain)
+		 	}
+		 	pfo.Hostfile.RemoveHost(localServiceName)
+		 	pfo.Hostfile.RemoveHost(nsServiceName)
+		}
+		pfo.Hostfile.RemoveHost(fullServiceName)
 	}
 
 	go func() {
 		<-signals
 		if stopChannel != nil {
-			if pfo.Remote == false {
-				if pfo.Domain != "" {
-					pfo.Hostfile.RemoveHost(localServiceName + "." + pfo.Domain)
-					pfo.Hostfile.RemoveHost(nsServiceName + "." + pfo.Domain)
-				}
-				pfo.Hostfile.RemoveHost(localServiceName)
-				pfo.Hostfile.RemoveHost(nsServiceName)
-			}
-			pfo.Hostfile.RemoveHost(fullServiceName)
-			err = pfo.Hostfile.Save()
-			if err != nil {
-				log.Error("Error saving hosts file", err)
-			}
-
+			cleanupHostfile()
 			close(stopChannel)
 		}
 	}()
@@ -146,6 +138,7 @@ func PortForward(pfo *PortForwardOpts) error {
 	fw, err := portforward.New(dialer, fwdPorts, stopChannel, readyChannel, &p, &p)
 	if err != nil {
 		signal.Stop(signals)
+		cleanupHostfile()
 		return err
 	}
 
@@ -154,6 +147,7 @@ func PortForward(pfo *PortForwardOpts) error {
 	err = fw.ForwardPorts()
 	if err != nil {
 		signal.Stop(signals)
+		cleanupHostfile()
 		return err
 	}