Skip to content

Commit

Permalink
fix(ssh): Do not create the ssh key if the key exists (#211)
Browse files Browse the repository at this point in the history
* 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>
  • Loading branch information
gaocegege authored May 29, 2022
1 parent 40c7d06 commit b8a5516
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 30 deletions.
12 changes: 7 additions & 5 deletions pkg/buildkitd/buildkitd.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
}
Expand Down
15 changes: 0 additions & 15 deletions pkg/buildkitd/mock/mock.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

23 changes: 13 additions & 10 deletions pkg/ssh/config/key.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ import (
"crypto/rsa"
"crypto/x509"
"encoding/pem"
"fmt"
"os"

"github.com/adrg/xdg"
Expand All @@ -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
}

Expand All @@ -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
}

Expand Down

0 comments on commit b8a5516

Please sign in to comment.