diff --git a/main.go b/main.go index f80220b..0422553 100644 --- a/main.go +++ b/main.go @@ -66,6 +66,11 @@ func main() { EnvVar: "PLUGIN_PORT,SSH_PORT", Value: 22, }, + cli.BoolFlag{ + Name: "sync", + Usage: "sync mode", + EnvVar: "PLUGIN_SYNC", + }, cli.DurationFlag{ Name: "timeout,t", Usage: "connection timeout", @@ -195,6 +200,7 @@ func run(c *cli.Context) error { Secrets: c.StringSlice("secrets"), Envs: c.StringSlice("envs"), Debug: c.Bool("debug"), + Sync: c.Bool("sync"), Proxy: easyssh.DefaultConfig{ Key: c.String("proxy.ssh-key"), KeyPath: c.String("proxy.key-path"), diff --git a/plugin.go b/plugin.go index e577565..27f04e4 100644 --- a/plugin.go +++ b/plugin.go @@ -34,6 +34,7 @@ type ( Envs []string Proxy easyssh.DefaultConfig Debug bool + Sync bool } // Plugin structure @@ -42,6 +43,80 @@ type ( } ) +func (p Plugin) exec(host string, wg *sync.WaitGroup, errChannel chan error) { + // Create MakeConfig instance with remote username, server address and path to private key. + ssh := &easyssh.MakeConfig{ + Server: host, + User: p.Config.UserName, + Password: p.Config.Password, + Port: strconv.Itoa(p.Config.Port), + Key: p.Config.Key, + KeyPath: p.Config.KeyPath, + Timeout: p.Config.Timeout, + Proxy: easyssh.DefaultConfig{ + Server: p.Config.Proxy.Server, + User: p.Config.Proxy.User, + Password: p.Config.Proxy.Password, + Port: p.Config.Proxy.Port, + Key: p.Config.Proxy.Key, + KeyPath: p.Config.Proxy.KeyPath, + Timeout: p.Config.Proxy.Timeout, + }, + } + + p.log(host, "======CMD======") + p.log(host, strings.Join(p.Config.Script, "\n")) + p.log(host, "======END======") + + env := []string{} + for _, key := range p.Config.Envs { + key = strings.ToUpper(key) + val := os.Getenv(key) + val = strings.Replace(val, " ", "", -1) + env = append(env, key+"="+val) + } + + p.Config.Script = append(env, p.Config.Script...) + + if p.Config.Debug { + p.log(host, "======ENV======") + p.log(host, strings.Join(env, "\n")) + p.log(host, "======END======") + } + + stdoutChan, stderrChan, doneChan, errChan, err := ssh.Stream(strings.Join(p.Config.Script, "\n"), p.Config.CommandTimeout) + if err != nil { + errChannel <- err + } else { + // read from the output channel until the done signal is passed + isTimeout := true + loop: + for { + select { + case isTimeout = <-doneChan: + break loop + case outline := <-stdoutChan: + p.log(host, "out:", outline) + case errline := <-stderrChan: + p.log(host, "err:", errline) + case err = <-errChan: + } + } + + // get exit code or command error. + if err != nil { + errChannel <- err + } + + // command time out + if !isTimeout { + errChannel <- fmt.Errorf(commandTimeOut) + } + } + + wg.Done() +} + func (p Plugin) log(host string, message ...interface{}) { if count := len(p.Config.Host); count == 1 { fmt.Printf("%s", fmt.Sprintln(message...)) @@ -69,85 +144,19 @@ func (p Plugin) Exec() error { errChannel := make(chan error, 1) finished := make(chan bool, 1) for _, host := range p.Config.Host { - go func(host string) { - // Create MakeConfig instance with remote username, server address and path to private key. - ssh := &easyssh.MakeConfig{ - Server: host, - User: p.Config.UserName, - Password: p.Config.Password, - Port: strconv.Itoa(p.Config.Port), - Key: p.Config.Key, - KeyPath: p.Config.KeyPath, - Timeout: p.Config.Timeout, - Proxy: easyssh.DefaultConfig{ - Server: p.Config.Proxy.Server, - User: p.Config.Proxy.User, - Password: p.Config.Proxy.Password, - Port: p.Config.Proxy.Port, - Key: p.Config.Proxy.Key, - KeyPath: p.Config.Proxy.KeyPath, - Timeout: p.Config.Proxy.Timeout, - }, - } - - p.log(host, "======CMD======") - p.log(host, strings.Join(p.Config.Script, "\n")) - p.log(host, "======END======") - - env := []string{} - for _, key := range p.Config.Envs { - key = strings.ToUpper(key) - val := os.Getenv(key) - val = strings.Replace(val, " ", "", -1) - env = append(env, key+"="+val) - } - - p.Config.Script = append(env, p.Config.Script...) - - if p.Config.Debug { - p.log(host, "======ENV======") - p.log(host, strings.Join(env, "\n")) - p.log(host, "======END======") - } - - stdoutChan, stderrChan, doneChan, errChan, err := ssh.Stream(strings.Join(p.Config.Script, "\n"), p.Config.CommandTimeout) - if err != nil { - errChannel <- err - } else { - // read from the output channel until the done signal is passed - isTimeout := true - loop: - for { - select { - case isTimeout = <-doneChan: - break loop - case outline := <-stdoutChan: - p.log(host, "out:", outline) - case errline := <-stderrChan: - p.log(host, "err:", errline) - case err = <-errChan: - } - } - - // get exit code or command error. - if err != nil { - errChannel <- err - } - - // command time out - if !isTimeout { - errChannel <- fmt.Errorf(commandTimeOut) - } - } - - wg.Done() - }(host) + if p.Config.Sync { + p.exec(host, &wg, errChannel) + } else { + go p.exec(host, &wg, errChannel) + } } - go func() { - wg.Wait() - close(finished) - }() + if !p.Config.Sync { + go func() { + wg.Wait() + close(finished) + }() + } select { case <-finished: diff --git a/plugin_test.go b/plugin_test.go index 833e4fd..83c5934 100644 --- a/plugin_test.go +++ b/plugin_test.go @@ -256,3 +256,20 @@ func TestSetENV(t *testing.T) { err := plugin.Exec() assert.Nil(t, err) } + +func TestSyncMode(t *testing.T) { + plugin := Plugin{ + Config: Config{ + Host: []string{"localhost", "127.0.0.1"}, + UserName: "drone-scp", + Port: 22, + KeyPath: "./tests/.ssh/id_rsa", + Script: []string{"whoami", "for i in {1..3}; do echo ${i}; sleep 1; done", "echo 'done'"}, + CommandTimeout: 60, + Sync: true, + }, + } + + err := plugin.Exec() + assert.Nil(t, err) +}