fix(ssh): Do not create the ssh key if the key exists (#211)

* fix: Fix log

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

* fix(builder): Fix interval and timeout

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

* fix: Update mock

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

Author: gaocegege

pkg/buildkitd/buildkitd.go

@@ -30,14 +30,15 @@ import (
 )
 
 var (
-	interval = time.Second * 1
+	interval          = time.Second * 1
+	timeoutConnection = time.Second * 5
+	timeoutRun        = time.Second * 3
 )
 
 // Client is a client for the buildkitd daemon.
 // It's up to the caller to close the client.
 type Client interface {
 	BuildkitdAddr() string
-	Bootstrap(ctx context.Context) (string, error)
 	// Solve calls Solve on the controller.
 	Solve(ctx context.Context, def *llb.Definition, opt client.SolveOpt, statusChan chan *client.SolveStatus) (*client.SolveResponse, error)
 	Close() error
@@ -67,14 +68,15 @@ func NewClient(ctx context.Context) (Client, error) {
 	}
 	c.Client = cli
 
-	if _, err := c.Bootstrap(ctx); err != nil {
+	if _, err := c.Bootstrap(ctx, timeoutRun, timeoutConnection); err != nil {
 		return nil, errors.Wrap(err, "failed to bootstrap the buildkitd")
 	}
 	return c, nil
 }
 
-func (c *generalClient) Bootstrap(ctx context.Context) (string, error) {
-	address, err := c.maybeStart(ctx, time.Second*100, time.Second*100)
+func (c *generalClient) Bootstrap(ctx context.Context,
+	runningTimeout, connectingTimeout time.Duration) (string, error) {
+	address, err := c.maybeStart(ctx, runningTimeout, connectingTimeout)
 	if err != nil {
 		return "", err
 	}

pkg/buildkitd/mock/mock.go

@@ -19,7 +19,6 @@ import (
 	"crypto/rsa"
 	"crypto/x509"
 	"encoding/pem"
-	"fmt"
 	"os"
 
 	"github.com/adrg/xdg"
@@ -39,19 +38,19 @@ func KeyExists(public, private string) bool {
 	// public, private := getKeyPaths()
 	publicKeyExists, _ := fileutil.FileExists(public)
 	if !publicKeyExists {
-		logrus.Infof("%s doesn't exist", public)
+		logrus.Debugf("%s doesn't exist", public)
 		return false
 	}
 
-	logrus.Infof("%s already present", public)
+	logrus.Debugf("%s already present", public)
 
 	privateKeyExists, _ := fileutil.FileExists(private)
 	if !privateKeyExists {
-		logrus.Infof("%s doesn't exist", private)
+		logrus.Debugf("%s doesn't exist", private)
 		return false
 	}
 
-	logrus.Infof("%s already present", private)
+	logrus.Debugf("%s already present", private)
 	return true
 }
 
@@ -65,27 +64,31 @@ func GenerateKeys() error {
 }
 
 func generateKeys(public, private string, bitSize int) error {
+	if KeyExists(public, private) {
+		return nil
+	}
+
 	privateKey, err := generatePrivateKey(bitSize)
 	if err != nil {
-		return fmt.Errorf("failed to generate private SSH key: %s", err)
+		return errors.Wrap(err, "failed to generate private SSH key")
 	}
 
 	publicKeyBytes, err := generatePublicKey(&privateKey.PublicKey)
 	if err != nil {
-		return fmt.Errorf("failed to generate public SSH key: %s", err)
+		return errors.Wrap(err, "failed to generate public SSH key")
 	}
 
 	privateKeyBytes := encodePrivateKeyToPEM(privateKey)
 
 	if err := os.WriteFile(public, publicKeyBytes, 0600); err != nil {
-		return fmt.Errorf("failed to write public SSH key: %s", err)
+		return errors.Wrap(err, "failed to write public SSH key")
 	}
 
 	if err := os.WriteFile(private, privateKeyBytes, 0600); err != nil {
-		return fmt.Errorf("failed to write private SSH key: %s", err)
+		return errors.Wrap(err, "failed to write private SSH key")
 	}
 
-	logrus.Infof("created ssh keypair at  %s and %s", public, private)
+	logrus.Debugf("created ssh keypair at  %s and %s", public, private)
 	return nil
 }