From f7676e73ba39bd2a343cb27a90576a89ba7df753 Mon Sep 17 00:00:00 2001
From: Andrew <cybershell@cybershell.xyz>
Date: Mon, 13 Mar 2023 20:25:27 -0500
Subject: [PATCH] Make host connections close last

---
 pkg/backy/backy.go | 15 ++++++++++++++-
 pkg/backy/ssh.go   |  1 +
 pkg/backy/types.go |  1 +
 3 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/pkg/backy/backy.go b/pkg/backy/backy.go
index 7d5a737..1c626a2 100644
--- a/pkg/backy/backy.go
+++ b/pkg/backy/backy.go
@@ -328,7 +328,9 @@ func (config *BackyConfigFile) ExecuteCmds(opts *BackyConfigOpts) {
 
 func (c *BackyConfigFile) closeHostConnections() {
 	for _, host := range c.Hosts {
-
+		if host.isProxyHost {
+			continue
+		}
 		if host.SshClient != nil {
 			if _, err := host.SshClient.NewSession(); err == nil {
 				c.Logger.Info().Msgf("Closing host connection %s", host.HostName)
@@ -336,6 +338,9 @@ func (c *BackyConfigFile) closeHostConnections() {
 			}
 		}
 		for _, proxyHost := range host.ProxyHost {
+			if proxyHost.isProxyHost {
+				continue
+			}
 			if proxyHost.SshClient != nil {
 				if _, err := host.SshClient.NewSession(); err == nil {
 					c.Logger.Info().Msgf("Closing connection to proxy host %s", host.HostName)
@@ -344,4 +349,12 @@ func (c *BackyConfigFile) closeHostConnections() {
 			}
 		}
 	}
+	for _, host := range c.Hosts {
+		if host.SshClient != nil {
+			if _, err := host.SshClient.NewSession(); err == nil {
+				c.Logger.Info().Msgf("Closing proxy host connection %s", host.HostName)
+				host.SshClient.Close()
+			}
+		}
+	}
 }
diff --git a/pkg/backy/ssh.go b/pkg/backy/ssh.go
index 13570b5..721fae0 100644
--- a/pkg/backy/ssh.go
+++ b/pkg/backy/ssh.go
@@ -392,6 +392,7 @@ func (remoteConfig *Host) GetProxyJumpConfig(hosts map[string]*Host) error {
 	remoteConfig.GetHostName()
 	remoteConfig.CombineHostNameWithPort()
 	remoteConfig.GetSshUserFromConfig()
+	remoteConfig.isProxyHost = true
 	if remoteConfig.HostName == "" {
 		return errors.Errorf("No hostname found or specified for host %s", remoteConfig.Host)
 	}
diff --git a/pkg/backy/types.go b/pkg/backy/types.go
index 89d627d..6d6c73f 100644
--- a/pkg/backy/types.go
+++ b/pkg/backy/types.go
@@ -50,6 +50,7 @@ type (
 		PrivateKeyPassword string `yaml:"privatekeypassword,omitempty"`
 		useDefaultConfig   bool
 		User               string `yaml:"user,omitempty"`
+		isProxyHost        bool
 		// ProxyHost holds the configuration for a ProxyJump host
 		ProxyHost []*Host
 		// CertPath           string `yaml:"cert_path,omitempty"`