Add sync mode. (#101)

* Add sync mode.

Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>

* close channel in sync mode.

Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>

* close channel in sync mode.

Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>

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"),

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,79 +144,11 @@ 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() {

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)
+}