feat(ssh): Add/Remove SSH config entry (#172)

* add ssh config entry

* address comment

Co-authored-by: Jinjing.Zhou <allenzhou@tensorchord.ai>

cmd/envd-ssh/main.go

@@ -27,6 +27,7 @@ import (
  cli "github.com/urfave/cli/v2"
 
  "github.com/tensorchord/envd/pkg/remote/sshd"
+ "github.com/tensorchord/envd/pkg/ssh"
  "github.com/tensorchord/envd/pkg/version"
 )
 
@@ -88,7 +89,7 @@ func sshServer(c *cli.Context) error {
  logrus.Fatal(err.Error())
  }
 
- port := 2222
+ port := ssh.DefaultSSHPort
  // TODO(gaocegege): Set it as a flag.
  if p, ok := os.LookupEnv(envPort); ok {
  var err error

cmd/envd/destroy.go

@@ -22,6 +22,7 @@ import (
  cli "github.com/urfave/cli/v2"
 
  "github.com/tensorchord/envd/pkg/docker"
+ "github.com/tensorchord/envd/pkg/ssh"
  "github.com/tensorchord/envd/pkg/util/fileutil"
 )
 
@@ -57,6 +58,11 @@ func destroy(clicontext *cli.Context) error {
  if err := dockerClient.Destroy(clicontext.Context, ctr); err != nil {
  return errors.Wrapf(err, "failed to destroy the environment: %s", ctr)
  }
+
+ if err = ssh.RemoveEntry(ctr); err != nil {
+ logrus.Infof("failed to remove entry %s from your SSH config file: %s", ctr, err)
+ return errors.Wrap(err, "failed to remove entry from your SSH config file")
+ }
  logrus.Info("envd environment destroyed")
  return nil
 }

cmd/envd/up.go

@@ -146,9 +146,15 @@ func up(clicontext *cli.Context) error {
  }
  logrus.Debugf("container %s is running", containerID)
 
+ logrus.Debugf("Add entry %s to SSH config. at %s", buildContext, containerIP)
+ if err = ssh.AddEntry(ctr, containerIP, ssh.DefaultSSHPort); err != nil {
+ logrus.Infof("failed to add entry %s to your SSH config file: %s", ctr, err)
+ return errors.Wrap(err, "failed to add entry to your SSH config file")
+ }
+
  if !detach {
  sshClient, err := ssh.NewClient(
- containerIP, "root", 2222, clicontext.Bool("auth"), clicontext.Path("private-key"), "")
+ containerIP, "root", ssh.DefaultSSHPort, clicontext.Bool("auth"), clicontext.Path("private-key"), "")
  if err != nil {
  return err
  }

pkg/ssh/ssh.go

@@ -36,6 +36,8 @@ import (
  "golang.org/x/term"
 )
 
+const DefaultSSHPort = 2222
+
 type Client interface {
  Attach() error
 }