feat: Support cli argument for host key (#992)

* feat: Support cli argument for host key Signed-off-by: Ce Gao <[email protected]> * fix: Update Signed-off-by: Ce Gao <[email protected]> * fix: Remove comments Signed-off-by: Ce Gao <[email protected]> Signed-off-by: Ce Gao <[email protected]>
tensorchord · Oct 9, 2022 · b2a9018 · b2a9018
1 parent f249e4b
commit b2a9018
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 35 deletions.
diff --git a/cmd/envd-sshd/main.go b/cmd/envd-sshd/main.go
@@ -18,12 +18,14 @@ package main
 
 import (
 	"fmt"
+	"io/ioutil"
 	"os"
 
 	"github.com/cockroachdb/errors"
-	rawssh "github.com/gliderlabs/ssh"
+	"github.com/gliderlabs/ssh"
 	"github.com/sirupsen/logrus"
 	"github.com/urfave/cli/v2"
+	rawssh "golang.org/x/crypto/ssh"
 
 	"github.com/tensorchord/envd/pkg/config"
 	"github.com/tensorchord/envd/pkg/remote/sshd"
@@ -36,6 +38,7 @@ const (
 	flagNoAuth  = "no-auth"
 	flagPort    = "port"
 	flagShell   = "shell"
+	flagHostKey = "hostkey"
 )
 
 func main() {
@@ -59,14 +62,20 @@ func main() {
 			EnvVars: []string{"ENVD_AUTHORIZED_KEYS_PATH"},
 			Aliases: []string{"a"},
 		},
+		&cli.StringFlag{
+			Name:    flagHostKey,
+			Usage:   "path to the host key",
+			EnvVars: []string{"ENVD_HOST_KEY"},
+		},
 		&cli.BoolFlag{
 			Name:  flagNoAuth,
 			Usage: "disable authentication",
 			Value: false,
 		},
 		&cli.IntFlag{
-			Name:  flagPort,
-			Usage: "port to listen on",
+			Name:    flagPort,
+			Usage:   "port to listen on",
+			Aliases: []string{"p"},
 		},
 		&cli.StringFlag{
 			Name:  flagShell,
@@ -107,7 +116,7 @@ func sshServer(c *cli.Context) error {
 	}
 
 	noAuth := c.Bool(flagNoAuth)
-	var keys []rawssh.PublicKey
+	var keys []ssh.PublicKey
 	if !noAuth {
 		var err error
 		path := c.String(flagAuthKey)
@@ -125,10 +134,27 @@ func sshServer(c *cli.Context) error {
 		logrus.Warn("no authentication enabled")
 	}
 
+	var hostKey ssh.Signer = nil
+	if c.String(flagHostKey) != "" {
+		// read private key file
+		pemBytes, err := ioutil.ReadFile(c.String(flagHostKey))
+		if err != nil {
+			return errors.Wrapf(
+				err, "reading private key %s failed", c.String(flagHostKey))
+		}
+		if privateKey, err := rawssh.ParsePrivateKey(pemBytes); err != nil {
+			return err
+		} else {
+			logrus.Debugf("load host key from %s", c.String(flagHostKey))
+			hostKey = privateKey
+		}
+	}
+
 	srv := sshd.Server{
 		Port:           port,
 		Shell:          shell,
 		AuthorizedKeys: keys,
+		Hostkey:        hostKey,
 	}
 
 	logrus.Infof("ssh server %s started in 0.0.0.0:%d", version.GetVersion().String(), srv.Port)

diff --git a/pkg/app/ssh.go b/pkg/app/ssh.go
@@ -58,6 +58,7 @@ func sshc(clicontext *cli.Context) error {
 	opt.User = it
 	opt.PrivateKeyPath = clicontext.Path("private-key")
 	opt.Port = 2222
+	opt.AgentForwarding = false
 	sshClient, err := ssh.NewClient(opt)
 	if err != nil {
 		return errors.Wrap(err, "failed to create the ssh client")

diff --git a/pkg/remote/sshd/sshd.go b/pkg/remote/sshd/sshd.go
@@ -68,9 +68,11 @@ func LoadAuthorizedKeys(path string) ([]ssh.PublicKey, error) {
 
 // Server holds the ssh server configuration.
 type Server struct {
-	Port           int
-	Shell          string
+	Port  int
+	Shell string
+
 	AuthorizedKeys []ssh.PublicKey
+	Hostkey        ssh.Signer
 }
 
 // ListenAndServe starts the SSH server using port
@@ -117,6 +119,10 @@ func (srv *Server) getServer() (*ssh.Server, error) {
 		server.PasswordHandler = nil
 	}
 
+	if srv.Hostkey != nil {
+		server.AddHostKey(srv.Hostkey)
+	}
+
 	return server, nil
 }
 

diff --git a/pkg/ssh/ssh.go b/pkg/ssh/ssh.go
@@ -47,20 +47,22 @@ type Client interface {
 }
 
 type Options struct {
-	Server         string
-	User           string
-	Port           int
-	Auth           bool
-	PrivateKeyPath string
-	PrivateKeyPwd  string
+	AgentForwarding bool
+	Server          string
+	User            string
+	Port            int
+	Auth            bool
+	PrivateKeyPath  string
+	PrivateKeyPwd   string
 }
 
 func DefaultOptions() Options {
 	return Options{
-		Server:        "localhost",
-		User:          "envd",
-		Auth:          true,
-		PrivateKeyPwd: "",
+		Server:          "localhost",
+		User:            "envd",
+		Auth:            true,
+		PrivateKeyPwd:   "",
+		AgentForwarding: true,
 	}
 }
 
@@ -82,6 +84,7 @@ func GetOptions(entry string) (*Options, error) {
 
 type generalClient struct {
 	cli *ssh.Client
+	opt *Options
 }
 
 func NewClient(opt Options) (Client, error) {
@@ -120,26 +123,31 @@ func NewClient(opt Options) (Client, error) {
 	}
 	cli = conn
 
-	// open connection to the local agent
-	socketLocation := os.Getenv("SSH_AUTH_SOCK")
-	if socketLocation != "" {
-		agentConn, err := net.Dial("unix", socketLocation)
-		if err != nil {
-			return nil, errors.Wrap(err, "could not connect to local agent socket")
-		}
-		// create agent and add in auth
-		forwardingAgent := agent.NewClient(agentConn)
-		// add callback for forwarding agent to SSH config
-		// XXX - might want to handle reconnects appending multiple callbacks
-		auth := ssh.PublicKeysCallback(forwardingAgent.Signers)
-		config.Auth = append(config.Auth, auth)
-		if err := agent.ForwardToAgent(cli, forwardingAgent); err != nil {
-			return nil, errors.Wrap(err, "forwarding agent to client failed")
+	if opt.AgentForwarding {
+		// open connection to the local agent
+		socketLocation := os.Getenv("SSH_AUTH_SOCK")
+		if socketLocation != "" {
+			agentConn, err := net.Dial("unix", socketLocation)
+			if err != nil {
+				return nil, errors.Wrap(err, "could not connect to local agent socket")
+			}
+			// create agent and add in auth
+			forwardingAgent := agent.NewClient(agentConn)
+			// add callback for forwarding agent to SSH config
+			// might want to handle reconnects appending multiple callbacks
+			auth := ssh.PublicKeysCallback(forwardingAgent.Signers)
+			config.Auth = append(config.Auth, auth)
+			if err := agent.ForwardToAgent(cli, forwardingAgent); err != nil {
+				return nil, errors.Wrap(err, "forwarding agent to client failed")
+			}
+		} else {
+			logrus.Warn("failed to get the environment variable SSH_AUTH_SOCK")
 		}
 	}
 
 	return &generalClient{
 		cli: cli,
+		opt: &opt,
 	}, nil
 }
 
@@ -157,8 +165,10 @@ func (c generalClient) ExecWithOutput(cmd string) ([]byte, error) {
 	}
 	defer session.Close()
 
-	if err := agent.RequestAgentForwarding(session); err != nil {
-		return nil, errors.Wrap(err, "requesting agent forwarding failed")
+	if c.opt.AgentForwarding {
+		if err := agent.RequestAgentForwarding(session); err != nil {
+			return nil, errors.Wrap(err, "requesting agent forwarding failed")
+		}
 	}
 
 	return session.CombinedOutput(cmd)
@@ -172,8 +182,10 @@ func (c generalClient) Attach() error {
 	}
 	defer session.Close()
 
-	if err := agent.RequestAgentForwarding(session); err != nil {
-		return errors.Wrap(err, "requesting agent forwarding failed")
+	if c.opt.AgentForwarding {
+		if err := agent.RequestAgentForwarding(session); err != nil {
+			return errors.Wrap(err, "requesting agent forwarding failed")
+		}
 	}
 
 	modes := ssh.TerminalModes{