feat: support bootstrap buildkit with registry (#1656)

Signed-off-by: Kevin Su <pingsutw@apache.org>
tensorchord · Jun 16, 2023 · 45dffe7 · 45dffe7
1 parent 1143cbf
commit 45dffe7
Show file tree

Hide file tree

Showing 7 changed files with 29 additions and 14 deletions.
diff --git a/pkg/app/bootstrap.go b/pkg/app/bootstrap.go
@@ -70,6 +70,11 @@ var CommandBootstrap = &cli.Command{
 			Usage: "Use HTTP instead of HTTPS for the registry",
 			Value: false,
 		},
+		&cli.StringFlag{
+			Name:    "registry",
+			Usage:   "Specify the registry to pull the image from",
+			Aliases: []string{"r"},
+		},
 	},
 
 	Action: bootstrap,
@@ -287,20 +292,21 @@ func buildkit(clicontext *cli.Context) error {
 	mirror := clicontext.String("dockerhub-mirror")
 	setRegistryCA := clicontext.IsSet("registry-ca-keypair")
 	useHTTP := clicontext.Bool("use-http")
+	registry := clicontext.String("registry")
 
 	if setRegistryCA && useHTTP {
 		return errors.New("cannot use both registry CA and HTTP")
 	}
 
 	if c.Builder == types.BuilderTypeMoby {
 		bkClient, err = buildkitd.NewMobyClient(clicontext.Context,
-			c.Builder, c.BuilderAddress, mirror, setRegistryCA, useHTTP)
+			c.Builder, c.BuilderAddress, mirror, registry, setRegistryCA, useHTTP)
 		if err != nil {
 			return errors.Wrap(err, "failed to create moby buildkit client")
 		}
 	} else {
 		bkClient, err = buildkitd.NewClient(clicontext.Context,
-			c.Builder, c.BuilderAddress, mirror, setRegistryCA, useHTTP)
+			c.Builder, c.BuilderAddress, mirror, registry, setRegistryCA, useHTTP)
 		if err != nil {
 			return errors.Wrap(err, "failed to create buildkit client")
 		}

diff --git a/pkg/app/prune.go b/pkg/app/prune.go
@@ -81,13 +81,13 @@ func prune(clicontext *cli.Context) error {
 	var bkClient buildkitd.Client
 	if c.Builder == types.BuilderTypeMoby {
 		bkClient, err = buildkitd.NewMobyClient(clicontext.Context,
-			c.Builder, c.BuilderAddress, "", false, false)
+			c.Builder, c.BuilderAddress, "", "", false, false)
 		if err != nil {
 			return errors.Wrap(err, "failed to create moby buildkit client")
 		}
 	} else {
 		bkClient, err = buildkitd.NewClient(clicontext.Context,
-			c.Builder, c.BuilderAddress, "", false, false)
+			c.Builder, c.BuilderAddress, "", "", false, false)
 		if err != nil {
 			return errors.Wrap(err, "failed to create buildkit client")
 		}

diff --git a/pkg/builder/build.go b/pkg/builder/build.go
@@ -90,13 +90,13 @@ func New(ctx context.Context, opt Options) (Builder, error) {
 	var cli buildkitd.Client
 	if c.Builder == types.BuilderTypeMoby {
 		cli, err = buildkitd.NewMobyClient(ctx,
-			c.Builder, c.BuilderAddress, "", false, false)
+			c.Builder, c.BuilderAddress, "", "", false, false)
 		if err != nil {
 			return nil, errors.Wrap(err, "failed to create moby buildkit client")
 		}
 	} else {
 		cli, err = buildkitd.NewClient(ctx,
-			c.Builder, c.BuilderAddress, "", false, false)
+			c.Builder, c.BuilderAddress, "", "", false, false)
 		if err != nil {
 			return nil, errors.Wrap(err, "failed to create buildkit client")
 		}

diff --git a/pkg/buildkitd/buildkitd.go b/pkg/buildkitd/buildkitd.go
@@ -65,6 +65,7 @@ type generalClient struct {
 	containerName    string
 	image            string
 	mirror           string
+	registry         string
 	enableRegistryCA bool
 	useHTTP          bool
 
@@ -76,11 +77,12 @@ type generalClient struct {
 }
 
 func NewMobyClient(ctx context.Context, driver types.BuilderType,
-	socket, mirror string, enableRegistryCA bool, useHTTP bool) (Client, error) {
+	socket, mirror, registry string, enableRegistryCA bool, useHTTP bool) (Client, error) {
 	logrus.Debug("getting moby buildkit client")
 	c := &generalClient{
 		containerName:    socket,
 		image:            viper.GetString(flag.FlagBuildkitdImage),
+		registry:         registry,
 		mirror:           mirror,
 		enableRegistryCA: enableRegistryCA,
 		useHTTP:          useHTTP,
@@ -113,11 +115,12 @@ func NewMobyClient(ctx context.Context, driver types.BuilderType,
 }
 
 func NewClient(ctx context.Context, driver types.BuilderType,
-	socket, mirror string, enableRegistryCA bool, useHTTP bool) (Client, error) {
+	socket, mirror, registry string, enableRegistryCA bool, useHTTP bool) (Client, error) {
 	c := &generalClient{
 		containerName:    socket,
 		image:            viper.GetString(flag.FlagBuildkitdImage),
 		mirror:           mirror,
+		registry:         registry,
 		enableRegistryCA: enableRegistryCA,
 		useHTTP:          useHTTP,
 		socket:           socket,
@@ -176,7 +179,7 @@ func (c *generalClient) maybeStart(ctx context.Context,
 	}
 
 	if client != nil {
-		if _, err := client.StartBuildkitd(ctx, c.image, c.containerName, c.mirror,
+		if _, err := client.StartBuildkitd(ctx, c.image, c.containerName, c.mirror, c.registry,
 			c.enableRegistryCA, c.useHTTP, runningTimeout); err != nil {
 			return "", err
 		}

diff --git a/pkg/driver/client.go b/pkg/driver/client.go
@@ -25,7 +25,7 @@ import (
 type Client interface {
 	// Load loads the image from the reader to the docker host.
 	Load(ctx context.Context, r io.ReadCloser, quiet bool) error
-	StartBuildkitd(ctx context.Context, tag, name, mirror string, enableRegistryCA, useHTTP bool, timeout time.Duration) (string, error)
+	StartBuildkitd(ctx context.Context, tag, name, mirror, registry string, enableRegistryCA, useHTTP bool, timeout time.Duration) (string, error)
 
 	Exec(ctx context.Context, cname string, cmd []string) error
 

diff --git a/pkg/driver/docker/docker.go b/pkg/driver/docker/docker.go
@@ -47,8 +47,12 @@ const buildkitdHTTP = `
 [registry."docker.io"]
   http = true
 `
-const buildkitdCertPath = "/etc/registry"
 const buildkitdRegistry = `
+[registry."%s"]
+  http = %v
+`
+const buildkitdCertPath = "/etc/registry"
+const buildkitdWithCA = `
 [registry."docker.io"]
   mirrors = ["%s"]
   ca=["/etc/registry/ca.pem"]
@@ -189,7 +193,7 @@ func (c dockerClient) ResumeContainer(ctx context.Context, name string) (string,
 	return name, nil
 }
 
-func (c dockerClient) StartBuildkitd(ctx context.Context, tag, name, mirror string,
+func (c dockerClient) StartBuildkitd(ctx context.Context, tag, name, mirror, registry string,
 	enableRegistryCA, useHTTP bool, timeout time.Duration) (string, error) {
 	logger := logrus.WithFields(logrus.Fields{
 		"tag":       tag,
@@ -225,7 +229,7 @@ func (c dockerClient) StartBuildkitd(ctx context.Context, tag, name, mirror stri
 	var cfg string
 	if mirror != "" {
 		if enableRegistryCA {
-			cfg = fmt.Sprintf(buildkitdRegistry, mirror)
+			cfg = fmt.Sprintf(buildkitdWithCA, mirror)
 			hostConfig.Mounts = append(hostConfig.Mounts, mount.Mount{
 				Type:   mount.TypeBind,
 				Source: fileutil.DefaultConfigDir,
@@ -234,6 +238,8 @@ func (c dockerClient) StartBuildkitd(ctx context.Context, tag, name, mirror stri
 		} else {
 			cfg = fmt.Sprintf(buildkitdMirror, mirror)
 		}
+	} else if registry != "" {
+		cfg = fmt.Sprintf(buildkitdRegistry, registry, useHTTP)
 	} else if useHTTP {
 		cfg = buildkitdHTTP
 	}

diff --git a/pkg/driver/nerdctl/nerdctl.go b/pkg/driver/nerdctl/nerdctl.go
@@ -63,7 +63,7 @@ func (nc *nerdctlClient) Load(ctx context.Context, r io.ReadCloser, quiet bool)
 	return nil
 }
 
-func (nc *nerdctlClient) StartBuildkitd(ctx context.Context, tag, name, mirror string,
+func (nc *nerdctlClient) StartBuildkitd(ctx context.Context, tag, name, mirror, registry string,
 	enableRegistryCA, useHTTP bool, timeout time.Duration) (string, error) {
 	logger := logrus.WithFields(logrus.Fields{
 		"tag":       tag,