feat(CLI): Download midi-ssh instead of the built-in one (#73)

* feat: Move ssh into cache

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Simplify the build progress logging

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

cmd/midi/bootstrap.go

@@ -28,6 +28,7 @@ import (
 	cli "github.com/urfave/cli/v2"
 
 	"github.com/tensorchord/MIDI/pkg/buildkitd"
+	"github.com/tensorchord/MIDI/pkg/home"
 	"github.com/tensorchord/MIDI/pkg/util/fileutil"
 )
 
@@ -50,7 +51,7 @@ var CommandBootstrap = &cli.Command{
 			Name:  "midi-ssh",
 			Usage: "url to download midi-ssh binary",
 			// TODO(gaocegege): Use version.Version to generate the right URL.
-			Value: "https://github.com/tensorchord/MIDI/releases/download/v0.0.1-alpha.1/midi-ssh_0.0.1-alpha.1_Linux_x86_64",
+			Value: "https://objects.githubusercontent.com/github-production-release-asset-2e65be/480303698/93f6c5c3-b1b9-4302-a943-9ad2294520d2?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20220427%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20220427T074231Z&X-Amz-Expires=300&X-Amz-Signature=96baa2ada9295bd45beeaa318f1f3427247d22ebf7d2737fcbf2144930f50b9f&X-Amz-SignedHeaders=host&actor_id=5100735&key_id=0&repo_id=480303698&response-content-disposition=attachment%3B%20filename%3Dmidi-ssh_0.0.1-alpha.1_Linux_x86_64&response-content-type=application%2Foctet-stream",
 		},
 	},
 
@@ -75,16 +76,22 @@ func bootstrap(clicontext *cli.Context) error {
 		logrus.Info("You may have to restart your shell for autocomplete to get initialized (e.g. run \"exec $SHELL\")\n")
 	}
 
+	sshURL := clicontext.String("midi-ssh")
+	if err := home.GetManager().DownloadOrCacheSSHBinary(sshURL); err != nil {
+		return errors.Wrap(err, "failed o download midi-ssh")
+	}
+
 	buildkit := clicontext.Bool("buildkit")
 
 	if buildkit {
 		logrus.Debug("bootstrap the buildkitd container")
 		bkClient := buildkitd.NewClient()
 		defer bkClient.Close()
-		err := bkClient.Bootstrap(clicontext.Context)
+		addr, err := bkClient.Bootstrap(clicontext.Context)
 		if err != nil {
 			return errors.Wrap(err, "failed to bootstrap buildkit")
 		}
+		logrus.Infof("The buildkit is running at %s", addr)
 	}
 	return nil
 }

cmd/midi/up.go

@@ -47,11 +47,16 @@ var CommandUp = &cli.Command{
 			Aliases: []string{"f"},
 			Value:   "./build.MIDI",
 		},
+		&cli.BoolFlag{
+			Name:  "auth",
+			Usage: "Enable authentication for ssh",
+			Value: false,
+		},
 		&cli.PathFlag{
 			Name:    "private-key",
 			Usage:   "Path to the private key",
 			Aliases: []string{"k"},
-			Value:   "./examples/ssh_keypairs/key.pem",
+			Value:   "~/.ssh/id_rsa",
 		},
 	},
 
@@ -88,7 +93,8 @@ func up(clicontext *cli.Context) error {
 	}
 	logrus.Debugf("container %s is running", containerID)
 
-	sshClient, err := ssh.NewClient(containerIP, "root", 2222, clicontext.Path("private-key"), "")
+	sshClient, err := ssh.NewClient(
+		containerIP, "root", 2222, clicontext.Bool("auth"), clicontext.Path("private-key"), "")
 	if err != nil {
 		return err
 	}

pkg/builder/builder.go

@@ -93,23 +93,15 @@ func (b generalBuilder) Build(ctx context.Context) error {
 	defer cancel()
 	eg, ctx := errgroup.WithContext(ctx)
 
-	pw, err := progresswriter.NewPrinter(context.TODO(), os.Stderr, b.progressMode)
+	pw, err := progresswriter.NewPrinter(context.TODO(), os.Stdout, b.progressMode)
 	if err != nil {
 		return err
 	}
-	mw := progresswriter.NewMultiWriter(pw)
-
-	var writers []progresswriter.Writer
-	w := mw.WithPrefix("", false)
-	writers = append(writers, w)
-
 	// Create a pipe to load the image into the docker host.
 	pipeR, pipeW := io.Pipe()
 	eg.Go(func() error {
 		defer func() {
-			for _, w := range writers {
-				close(w.Status())
-			}
+			close(pw.Status())
 		}()
 		defer pipeW.Close()
 		wd, err := os.Getwd()
@@ -133,7 +125,7 @@ func (b generalBuilder) Build(ctx context.Context) error {
 				flag.FlagContextDir: wd,
 				flag.FlagCacheDir:   home.GetManager().CacheDir(),
 			},
-		}, progresswriter.ResetTime(mw.WithPrefix("", false)).Status())
+		}, pw.Status())
 		if err != nil {
 			err = errors.Wrap(err, "failed to solve LLB")
 			b.logger.Error(err)

pkg/buildkitd/buildkitd.go

@@ -20,7 +20,7 @@ var (
 // Client is a client for the buildkitd daemon.
 // It's up to the caller to close the client.
 type Client interface {
-	Bootstrap(ctx context.Context) error
+	Bootstrap(ctx context.Context) (string, error)
 	Close() error
 }
 
@@ -44,12 +44,12 @@ func NewClient() Client {
 	return c
 }
 
-func (c *generalClient) Bootstrap(ctx context.Context) error {
-	_, err := c.maybeStart(ctx, time.Second*100, time.Second*100)
+func (c *generalClient) Bootstrap(ctx context.Context) (string, error) {
+	address, err := c.maybeStart(ctx, time.Second*100, time.Second*100)
 	if err != nil {
-		return err
+		return "", err
 	}
-	return nil
+	return address, nil
 }
 
 func (c generalClient) Close() error {

pkg/home/cache.go

@@ -0,0 +1,50 @@
+package home
+
+import (
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+
+	"github.com/cockroachdb/errors"
+	"github.com/sirupsen/logrus"
+)
+
+func (m generalManager) SSHBinaryFile() string {
+	return "midi-ssh"
+}
+
+func (m generalManager) DownloadOrCacheSSHBinary(url string) error {
+	//TODO(gaocegege): Check the checksums.txt first.
+	// https://github.com/tensorchord/MIDI/releases/download/v0.0.1-alpha.1/checksums.txt
+
+	filename := filepath.Join(m.CacheDir(), m.SSHBinaryFile())
+	logger := logrus.WithFields(logrus.Fields{
+		"url":      url,
+		"filename": filename,
+	})
+	out, err := os.Create(filename)
+
+	if err != nil {
+		return err
+	}
+	defer out.Close()
+
+	resp, err := http.Get(url)
+	if err != nil {
+		return errors.Wrap(err, "failed to download ssh binary")
+	}
+	logger.Debugf("downloading %s", filename)
+
+	defer resp.Body.Close()
+	_, err = io.Copy(out, resp.Body)
+	if err != nil {
+		return errors.Wrap(err, "failed to download ssh binary")
+	}
+
+	if err := os.Chmod(filename, 0755); err != nil {
+		return errors.Wrap(err, "failed to chmod ssh binary")
+	}
+
+	return nil
+}

pkg/home/manager.go

@@ -18,6 +18,8 @@ type Manager interface {
 	HomeDir() string
 	CacheDir() string
 	ConfigFile() string
+	DownloadOrCacheSSHBinary(url string) error
+	SSHBinaryFile() string
 }
 
 type generalManager struct {

pkg/lang/ir/graph.go

@@ -24,6 +24,7 @@ import (
 	"github.com/moby/buildkit/client/llb"
 	"github.com/sirupsen/logrus"
 	"github.com/tensorchord/MIDI/pkg/flag"
+	"github.com/tensorchord/MIDI/pkg/home"
 	"github.com/tensorchord/MIDI/pkg/vscode"
 )
 
@@ -215,11 +216,9 @@ func (g Graph) compileSystemPackages(root llb.State) llb.State {
 }
 
 func (g Graph) copyMidiSSHServer() llb.State {
-	run := llb.Scratch().File(llb.Copy(llb.Local(flag.FlagContextDir),
-		"examples/ssh_keypairs/public.pub", "/var/midi/remote/authorized_keys",
-		&llb.CopyInfo{CreateDestPath: true})).
-		File(llb.Copy(llb.Local(flag.FlagContextDir),
-			"bin/midi-ssh", "/var/midi/bin/midi-ssh",
+	run := llb.Scratch().
+		File(llb.Copy(llb.Local(flag.FlagCacheDir),
+			home.GetManager().SSHBinaryFile(), "/var/midi/bin/midi-ssh",
 			&llb.CopyInfo{CreateDestPath: true}))
 	return run
 }

pkg/ssh/ssh.go

@@ -44,28 +44,31 @@ type generalClient struct {
 }
 
 func NewClient(server, user string,
-	port int, privateKeyPath, privateKeyPwd string) (Client, error) {
-	// read private key file
-	pemBytes, err := ioutil.ReadFile(privateKeyPath)
-	if err != nil {
-		return nil, fmt.Errorf("reading private key file failed %v", err)
-	}
-	// create signer
-	signer, err := signerFromPem(pemBytes, []byte(privateKeyPwd))
-	if err != nil {
-		return nil, err
-	}
-
+	port int, auth bool, privateKeyPath, privateKeyPwd string) (Client, error) {
 	config := &ssh.ClientConfig{
 		User: user,
-		Auth: []ssh.AuthMethod{
-			ssh.PublicKeys(signer),
-		},
 		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
 			// use OpenSSH's known_hosts file if you care about host validation
 			return nil
 		},
 	}
+
+	if auth {
+		// read private key file
+		pemBytes, err := ioutil.ReadFile(privateKeyPath)
+		if err != nil {
+			return nil, fmt.Errorf("reading private key file failed %v", err)
+		}
+		// create signer
+		signer, err := signerFromPem(pemBytes, []byte(privateKeyPwd))
+		if err != nil {
+			return nil, err
+		}
+		config.Auth = []ssh.AuthMethod{
+			ssh.PublicKeys(signer),
+		}
+	}
+
 	return &generalClient{
 		config: config,
 		server: fmt.Sprintf("%v:%v", server, port),