Merge pull request #11 from AkihiroSuda/dev

support specifying port

Author: AkihiroSuda

cmd/sshocker/run.go

@@ -1,8 +1,10 @@
 package main
 
 import (
+	"net"
 	"os/user"
 	"path/filepath"
+	"strconv"
 	"strings"
 
 	"github.com/AkihiroSuda/sshocker/pkg/mount"
@@ -43,6 +45,23 @@ var (
 	}
 )
 
+func parseHost(s string) (string, int, error) {
+	if !strings.Contains(s, ":") {
+		// FIXME: this check is not valid for IPv6!
+		return s, 22, nil
+	}
+	host, portStr, err := net.SplitHostPort(s)
+	if err != nil {
+		return "", 0, err
+	}
+	port, err := strconv.Atoi(portStr)
+	if err != nil {
+		return "", 0, err
+	}
+	// host may contain "user@" prefix.
+	return host, port, nil
+}
+
 func runAction(clicontext *cli.Context) error {
 	if clicontext.NArg() < 1 {
 		return errors.New("no host specified")
@@ -51,9 +70,14 @@ func runAction(clicontext *cli.Context) error {
 		ConfigFile: clicontext.String("ssh-config"),
 		Persist:    clicontext.Bool("ssh-persist"),
 	}
+	host, port, err := parseHost(clicontext.Args().First())
+	if err != nil {
+		return err
+	}
 	x := &sshocker.Sshocker{
 		SSHConfig: sshConfig,
-		Host:      clicontext.Args().First(),
+		Host:      host,
+		Port:      port,
 		Command:   clicontext.Args().Tail(),
 	}
 	if len(x.Command) > 0 && x.Command[0] == "--" {

pkg/reversesshfs/reversesshfs.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"strconv"
 
 	"github.com/AkihiroSuda/sshocker/pkg/ssh"
 	"github.com/AkihiroSuda/sshocker/pkg/util"
@@ -19,6 +20,7 @@ type ReverseSSHFS struct {
 	*ssh.SSHConfig
 	LocalPath  string
 	Host       string
+	Port       int
 	RemotePath string
 	Readonly   bool
 	sshCmd     *exec.Cmd
@@ -30,7 +32,7 @@ func (rsf *ReverseSSHFS) Prepare() error {
 	if !filepath.IsAbs(rsf.RemotePath) {
 		return errors.Errorf("unexpected relative path: %q", rsf.RemotePath)
 	}
-	sshArgs = append(sshArgs, rsf.Host, "--", "mkdir", "-p", rsf.RemotePath)
+	sshArgs = append(sshArgs, "-p", strconv.Itoa(rsf.Port), rsf.Host, "--", "mkdir", "-p", rsf.RemotePath)
 	sshCmd := exec.Command(sshBinary, sshArgs...)
 	logrus.Debugf("executing ssh for preparing sshfs: %s %v", sshCmd.Path, sshCmd.Args)
 	out, err := sshCmd.CombinedOutput()
@@ -49,7 +51,7 @@ func (rsf *ReverseSSHFS) Start() error {
 	if !filepath.IsAbs(rsf.RemotePath) {
 		return errors.Errorf("unexpected relative path: %q", rsf.RemotePath)
 	}
-	sshArgs = append(sshArgs, rsf.Host, "--", "sshfs", ":"+rsf.LocalPath, rsf.RemotePath, "-o", "slave")
+	sshArgs = append(sshArgs, "-p", strconv.Itoa(rsf.Port), rsf.Host, "--", "sshfs", ":"+rsf.LocalPath, rsf.RemotePath, "-o", "slave")
 	if rsf.Readonly {
 		sshArgs = append(sshArgs, "-o", "ro")
 	}
@@ -139,7 +141,7 @@ done
 	}
 	script := b.String()
 	logrus.Debugf("generated script %q with map %v: %q", scriptName, m, script)
-	stdout, stderr, err := ssh.ExecuteScript(rsf.Host, rsf.SSHConfig, script, scriptName)
+	stdout, stderr, err := ssh.ExecuteScript(rsf.Host, rsf.Port, rsf.SSHConfig, script, scriptName)
 	logrus.Debugf("executed script %q, stdout=%q, stderr=%q, err=%v", scriptName, stdout, stderr, err)
 	return err
 }

pkg/ssh/ssh.go

@@ -38,17 +38,17 @@ func (c *SSHConfig) Args() []string {
 }
 
 // ExitMaster executes `ssh -O exit`
-func ExitMaster(host string, c *SSHConfig) error {
+func ExitMaster(host string, port int, c *SSHConfig) error {
 	if c == nil {
 		return errors.New("got nil SSHConfig")
 	}
 	args := c.Args()
-	args = append(args, "-O", "exit", host)
+	args = append(args, "-O", "exit", "-p", strconv.Itoa(port), host)
 	cmd := exec.Command(c.Binary(), args...)
 	logrus.Debugf("executing ssh for exiting the master: %s %v", cmd.Path, cmd.Args)
 	out, err := cmd.CombinedOutput()
 	if err != nil {
-		return errors.Wrapf(err, "failed to execute `%s -O exit %s`, out=%q", c.Binary(), host, string(out))
+		return errors.Wrapf(err, "failed to execute `%s -O exit -p %d %s`, out=%q", c.Binary(), port, host, string(out))
 	}
 	return nil
 }
@@ -78,7 +78,7 @@ func parseScriptInterpreter(script string) (string, error) {
 // Returns stdout and stderr.
 //
 // scriptName is used only for readability of error strings.
-func ExecuteScript(host string, c *SSHConfig, script, scriptName string) (string, string, error) {
+func ExecuteScript(host string, port int, c *SSHConfig, script, scriptName string) (string, string, error) {
 	if c == nil {
 		return "", "", errors.New("got nil SSHConfig")
 	}
@@ -88,7 +88,7 @@ func ExecuteScript(host string, c *SSHConfig, script, scriptName string) (string
 	}
 	sshBinary := c.Binary()
 	sshArgs := c.Args()
-	sshArgs = append(sshArgs, host, "--", interpreter)
+	sshArgs = append(sshArgs, "-p", strconv.Itoa(port), host, "--", interpreter)
 	sshCmd := exec.Command(sshBinary, sshArgs...)
 	sshCmd.Stdin = strings.NewReader(script)
 	var stderr bytes.Buffer

pkg/sshocker/sshocker.go

@@ -3,6 +3,7 @@ package sshocker
 import (
 	"os"
 	"os/exec"
+	"strconv"
 
 	"github.com/AkihiroSuda/sshocker/pkg/mount"
 	"github.com/AkihiroSuda/sshocker/pkg/reversesshfs"
@@ -14,6 +15,7 @@ import (
 type Sshocker struct {
 	*ssh.SSHConfig
 	Host      string   // Required
+	Port      int      // Required
 	Command   []string // Optional
 	Mounts    []mount.Mount
 	LForwards []string
@@ -28,7 +30,7 @@ func (x *Sshocker) Run() error {
 	for _, l := range x.LForwards {
 		args = append(args, "-L", l)
 	}
-	args = append(args, x.Host, "--")
+	args = append(args, "-p", strconv.Itoa(x.Port), x.Host, "--")
 	if len(x.Command) > 0 {
 		args = append(args, x.Command...)
 	}
@@ -44,6 +46,7 @@ func (x *Sshocker) Run() error {
 				SSHConfig:  x.SSHConfig,
 				LocalPath:  m.Source,
 				Host:       x.Host,
+				Port:       x.Port,
 				RemotePath: m.Destination,
 				Readonly:   m.Readonly,
 			}
@@ -66,7 +69,7 @@ func (x *Sshocker) Run() error {
 	}
 	defer func() {
 		if x.SSHConfig.Persist {
-			if emErr := ssh.ExitMaster(x.Host, x.SSHConfig); emErr != nil {
+			if emErr := ssh.ExitMaster(x.Host, x.Port, x.SSHConfig); emErr != nil {
 				logrus.WithError(emErr).Error("failed to exit the master")
 			}
 		}