vm/isolated: allow to specify ssh port for target machines

1 files changed, 40 insertions, 13 deletions

diff --git a/vm/isolated/isolated.go b/vm/isolated/isolated.go
index 581e33f64..540ce9bce 100644
--- a/vm/isolated/isolated.go
+++ b/vm/isolated/isolated.go
@@ -9,6 +9,8 @@ import (
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"strconv"
+	"strings"
 	"time"
 
 	"github.com/google/syzkaller/pkg/config"
@@ -22,7 +24,7 @@ func init() {
 }
 
 type Config struct {
-	Targets       []string // target machines
+	Targets       []string // target machines: (hostname|ip)(:port)?
 	Target_Dir    string   // directory to copy/run on target
 	Target_Reboot bool     // reboot target on repair
 }
@@ -33,12 +35,13 @@ type Pool struct {
 }
 
 type instance struct {
-	cfg    *Config
-	target string
-	closed chan bool
-	debug  bool
-	sshkey string
-	port   int
+	cfg        *Config
+	target     string
+	targetPort int
+	closed     chan bool
+	debug      bool
+	sshkey     string
+	port       int
 }
 
 func ctor(env *vmimpl.Env) (vmimpl.Pool, error) {
@@ -56,6 +59,11 @@ func ctor(env *vmimpl.Env) (vmimpl.Pool, error) {
 	if env.SshKey != "" && !osutil.IsExist(env.SshKey) {
 		return nil, fmt.Errorf("ssh key '%v' does not exist", env.SshKey)
 	}
+	for _, target := range cfg.Targets {
+		if _, _, err := splitTargetPort(target); err != nil {
+			return nil, fmt.Errorf("bad target %q: %v", target, err)
+		}
+	}
 	if env.Debug {
 		cfg.Targets = cfg.Targets[:1]
 	}
@@ -71,12 +79,14 @@ func (pool *Pool) Count() int {
 }
 
 func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
+	target, targetPort, _ := splitTargetPort(pool.cfg.Targets[index])
 	inst := &instance{
-		cfg:    pool.cfg,
-		target: pool.env.SshUser + "@" + pool.cfg.Targets[index],
-		closed: make(chan bool),
-		debug:  pool.env.Debug,
-		sshkey: pool.env.SshKey,
+		cfg:        pool.cfg,
+		target:     pool.env.SshUser + "@" + target,
+		targetPort: targetPort,
+		closed:     make(chan bool),
+		debug:      pool.env.Debug,
+		sshkey:     pool.env.SshKey,
 	}
 	closeInst := inst
 	defer func() {
@@ -339,7 +349,7 @@ func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command strin
 
 func (inst *instance) sshArgs(portArg string) []string {
 	args := []string{
-		portArg, "22",
+		portArg, fmt.Sprint(inst.targetPort),
 		"-o", "ConnectionAttempts=10",
 		"-o", "ConnectTimeout=10",
 		"-o", "BatchMode=yes",
@@ -356,3 +366,20 @@ func (inst *instance) sshArgs(portArg string) []string {
 	}
 	return args
 }
+
+func splitTargetPort(addr string) (string, int, error) {
+	target := addr
+	port := 22
+	if colonPos := strings.Index(addr, ":"); colonPos != -1 {
+		p, err := strconv.ParseUint(addr[colonPos+1:], 10, 16)
+		if err != nil {
+			return "", 0, err
+		}
+		target = addr[:colonPos]
+		port = int(p)
+	}
+	if target == "" {
+		return "", 0, fmt.Errorf("target is empty")
+	}
+	return target, port, nil
+}