alajmo · alajmo · Aug 12, 2022 · Aug 5, 2022 · Aug 7, 2022 · Aug 12, 2022
diff --git a/Makefile b/Makefile
@@ -37,6 +37,9 @@ update-golden-files:
 mock-ssh:
 	cd ./test && docker-compose up
 
+mock-performance-ssh:
+	cd ./test && docker-compose -f docker-compose-performance.yaml up
+
 build:
 	CGO_ENABLED=0 go build \
 	-ldflags "-s -w -X '${PACKAGE}/cmd.version=${VERSION}' -X '${PACKAGE}/cmd.commit=${GIT}' -X '${PACKAGE}/cmd.date=${DATE}'" \

diff --git a/core/dao/server.go b/core/dao/server.go
@@ -27,6 +27,7 @@ type Server struct {
 	WorkDir      string
 	IdentityFile *string
 	Password     *string
+	AuthMethod   string
 
 	context     string // config path
 	contextLine int    // defined at
@@ -45,6 +46,7 @@ type ServerYAML struct {
 	WorkDir      string    `yaml:"work_dir"`
 	IdentityFile *string   `yaml:"identity_file"`
 	Password     *string   `yaml:"password"`
+	AuthMethod   string    `yaml:"-"`
 }
 
 func (s Server) GetValue(key string, _ int) string {
@@ -188,6 +190,16 @@ func (c *ConfigYAML) ParseServersYAML() ([]Server, []ResourceErrors[Server]) {
 			server.Password = serverYAML.Password
 		}
 
+		if server.IdentityFile != nil && server.Password != nil {
+			server.AuthMethod = "password-key"
+		} else if server.IdentityFile != nil {
+			server.AuthMethod = "key"
+		} else if server.Password != nil {
+			server.AuthMethod = "password"
+		} else {
+			server.AuthMethod = "none"
+		}
+
 		servers = append(servers, *server)
 	}
 

diff --git a/core/run/client.go b/core/run/client.go
@@ -28,3 +28,7 @@ type ErrConnect struct {
 	Port   uint16
 	Reason string
 }
+
+func (e *ErrConnect) Error() string {
+	return ""
+}
diff --git a/core/run/exec.go b/core/run/exec.go
@@ -5,7 +5,6 @@ import (
 	"golang.org/x/crypto/ssh"
 	"os"
 	"os/signal"
-	"os/user"
 	"path/filepath"
 	"strconv"
 	"strings"
@@ -161,45 +160,6 @@ func (run *Run) SetClients(
 	clientCh chan Client,
 	errCh chan ErrConnect,
 ) ([]ErrConnect, error) {
-	globalIdentityFile, globalPassword := getGlobalIdentity(runFlags)
-
-	// Iterate through servers and create a singleton AuthMethod, which is used for
-	// connecting to all hosts using a identity key. Servers which only use a password
-	// are not included here and are handled separately.
-	var identities []Identity
-	for _, server := range run.Servers {
-		if server.Local {
-			continue
-		}
-
-		var pass *string
-		if server.Password != nil {
-			pw, err := dao.EvaluatePassword(*server.Password)
-			pass = &pw
-			if err != nil {
-				errConnect := &ErrConnect{
-					Name:   server.Name,
-					User:   server.User,
-					Host:   server.Host,
-					Port:   server.Port,
-					Reason: err.Error(),
-				}
-				return []ErrConnect{*errConnect}, nil
-			}
-		}
-
-		identities = append(identities, Identity{
-			IdentityFile: server.IdentityFile,
-			Password:     pass,
-		})
-	}
-
-	// VerifyHost
-	authMethod, err := InitAuthMethod(globalIdentityFile, globalPassword, identities)
-	if err != nil {
-		return []ErrConnect{}, err
-	}
-
 	createLocalClient := func(server dao.Server, wg *sync.WaitGroup, mu *sync.Mutex) {
 		defer wg.Done()
 
@@ -212,37 +172,15 @@ func (run *Run) SetClients(
 		clientCh <- local
 	}
 
-	createRemoteClient := func(server dao.Server, wg *sync.WaitGroup, mu *sync.Mutex) {
+	createRemoteClient := func(authMethod []ssh.AuthMethod, server dao.Server, wg *sync.WaitGroup, mu *sync.Mutex) {
 		defer wg.Done()
 
-		var auth ssh.AuthMethod
-		if server.IdentityFile == nil && server.Password != nil {
-			// Password only logic
-			password, err := dao.EvaluatePassword(*server.Password)
-			if err != nil {
-				errConnect := &ErrConnect{
-					Name:   server.Name,
-					User:   server.User,
-					Host:   server.Host,
-					Port:   server.Port,
-					Reason: err.Error(),
-				}
-				errCh <- *errConnect
-			}
-
-			passwordAuth := ssh.Password(password)
-			auth = passwordAuth
-		} else {
-			// Identity key logic
-			auth = authMethod
-		}
-
 		remote := &SSHClient{
 			Name:       server.Name,
 			User:       server.User,
 			Host:       server.Host,
 			Port:       server.Port,
-			AuthMethod: auth,
+			AuthMethod: authMethod,
 		}
 
 		if err := remote.Connect(run.Config.DisableVerifyHost, run.Config.KnownHostsFile, mu); err != nil {
@@ -255,13 +193,77 @@ func (run *Run) SetClients(
 
 	var wg sync.WaitGroup
 	var mu sync.Mutex
+
+	agentSigners, err := GetSSHAgentSigners()
+	if err != nil {
+		return []ErrConnect{}, err
+	}
+
+	globalSigner, err := GetGlobalIdentitySigner(runFlags)
+	if err != nil {
+		return []ErrConnect{}, err
+	}
+
+	identities := make(map[string]ssh.Signer)
+	passwordAuthMethods := make(map[string]ssh.AuthMethod)
+	for _, server := range run.Servers {
+		if server.AuthMethod == "password-key" {
+			_, found := identities[*server.IdentityFile]
+			if !found {
+				signer, err := GetPassworIdentitySigner(server)
+				if err != nil {
+					return []ErrConnect{*err}, nil
+				}
+				identities[*server.IdentityFile] = signer
+			}
+		} else if server.AuthMethod == "key" {
+			_, found := identities[*server.IdentityFile]
+			if !found {
+				signer, err := GetIdentity(server)
+				if err != nil {
+					return []ErrConnect{*err}, nil
+				}
+				identities[*server.IdentityFile] = signer
+			}
+		} else if server.AuthMethod == "password" {
+			_, found := passwordAuthMethods[*server.Password]
+			if !found {
+				passAuthMethod, err := GetPasswordAuth(server)
+				if err != nil {
+					return []ErrConnect{*err}, nil
+				}
+				passwordAuthMethods[*server.Password] = passAuthMethod
+			}
+		}
+	}
+
 	for _, server := range run.Servers {
 		wg.Add(1)
 		go createLocalClient(server, &wg, &mu)
 
 		if !server.Local {
 			wg.Add(1)
-			go createRemoteClient(server, &wg, &mu)
+
+			var authMethods []ssh.AuthMethod
+			var signers []ssh.Signer
+
+			if globalSigner != nil {
+				signers = append(signers, globalSigner)
+			} else if server.AuthMethod == "password" {
+				pwAuth := passwordAuthMethods[*server.Password]
+				authMethods = append(authMethods, pwAuth)
+			} else if server.AuthMethod == "key" || server.AuthMethod == "password-key" {
+				identitySigner := identities[*server.IdentityFile]
+				signers = append(signers, identitySigner)
+			} else if agentSigners != nil {
+				signers = append(signers, agentSigners...)
+			}
+
+			if len(signers) > 0 {
+				authMethods = append(authMethods, ssh.PublicKeys(signers...))
+			}
+
+			go createRemoteClient(authMethods, server, &wg, &mu)
 		}
 	}
 	wg.Wait()
@@ -412,40 +414,6 @@ func (run *Run) setKnownHostsFile(knownHostsFileFlag string) error {
 	return nil
 }
 
-func getGlobalIdentity(runFlags *core.RunFlags) (string, string) {
-	var identityFile string
-	var password string
-
-	if runFlags.IdentityFile != "" {
-		identityFile = runFlags.IdentityFile
-	} else {
-		value, found := os.LookupEnv("SAKE_IDENTITY_FILE")
-		if found {
-			if strings.HasPrefix(value, "~/") {
-				usr, err := user.Current()
-				if err != nil {
-					panic(err)
-				}
-				dir := usr.HomeDir
-				identityFile = filepath.Join(dir, value[2:])
-			} else {
-				identityFile = value
-			}
-		}
-	}
-
-	if runFlags.Password != "" {
-		password = runFlags.Password
-	} else {
-		value, found := os.LookupEnv("SAKE_PASSWORD")
-		if found {
-			password = value
-		}
-	}
-
-	return identityFile, password
-}
-
 func getWorkDir(cmd dao.TaskCmd, server dao.Server) string {
 	if cmd.Local || server.Local {
 		rootDir := os.ExpandEnv(cmd.RootDir)