Repository: tensorchord/envd (language: python)

Issue:
refactor: bootstrap buildkit config file
### Describe the feature

1. no need to pass these arguments layer by layer since it's only used by the buildkit part
2. generate the config file with a template engine to make the logic clean


### Why do you need this feature?

_No response_

### Additional context

_No response_


Discussion / hints:

Fix patch:
diff --git a/pkg/app/bootstrap.go b/pkg/app/bootstrap.go
index defc64f8e..2dc421937 100644
--- a/pkg/app/bootstrap.go
+++ b/pkg/app/bootstrap.go
@@ -30,6 +30,7 @@ import (
 	"github.com/tensorchord/envd/pkg/home"
 	sshconfig "github.com/tensorchord/envd/pkg/ssh/config"
 	"github.com/tensorchord/envd/pkg/types"
+	"github.com/tensorchord/envd/pkg/util/buildkitutil"
 	"github.com/tensorchord/envd/pkg/util/fileutil"
 )
 
@@ -289,24 +290,22 @@ func buildkit(clicontext *cli.Context) error {
 
 	logrus.Debug("bootstrap the buildkitd container")
 	var bkClient buildkitd.Client
-	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")
+	config := buildkitutil.BuildkitConfig{
+		Registry: clicontext.String("registry"),
+		Mirror:   clicontext.String("dockerhub-mirror"),
+		UseHTTP:  clicontext.Bool("use-http"),
+		SetCA:    clicontext.IsSet("registry-ca-keypair"),
 	}
 
 	if c.Builder == types.BuilderTypeMoby {
 		bkClient, err = buildkitd.NewMobyClient(clicontext.Context,
-			c.Builder, c.BuilderAddress, mirror, registry, setRegistryCA, useHTTP)
+			c.Builder, c.BuilderAddress, &config)
 		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, registry, setRegistryCA, useHTTP)
+			c.Builder, c.BuilderAddress, &config)
 		if err != nil {
 			return errors.Wrap(err, "failed to create buildkit client")
 		}
diff --git a/pkg/app/prune.go b/pkg/app/prune.go
index ab4da99be..6f212be0a 100644
--- 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, nil)
 		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, nil)
 		if err != nil {
 			return errors.Wrap(err, "failed to create buildkit client")
 		}
diff --git a/pkg/builder/build.go b/pkg/builder/build.go
index 2957b2f92..39e1cc53b 100644
--- a/pkg/builder/build.go
+++ b/pkg/builder/build.go
@@ -39,6 +39,7 @@ import (
 	"github.com/tensorchord/envd/pkg/lang/version"
 	"github.com/tensorchord/envd/pkg/progress/progresswriter"
 	"github.com/tensorchord/envd/pkg/types"
+	"github.com/tensorchord/envd/pkg/util/buildkitutil"
 )
 
 func New(ctx context.Context, opt Options) (Builder, error) {
@@ -88,15 +89,16 @@ func New(ctx context.Context, opt Options) (Builder, error) {
 	}
 
 	var cli buildkitd.Client
+	bc := buildkitutil.BuildkitConfig{}
 	if c.Builder == types.BuilderTypeMoby {
 		cli, err = buildkitd.NewMobyClient(ctx,
-			c.Builder, c.BuilderAddress, "", "", false, false)
+			c.Builder, c.BuilderAddress, &bc)
 		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, &bc)
 		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
index 97b0bcc9c..a54c4d062 100644
--- a/pkg/buildkitd/buildkitd.go
+++ b/pkg/buildkitd/buildkitd.go
@@ -37,6 +37,7 @@ import (
 	"github.com/tensorchord/envd/pkg/driver/nerdctl"
 	"github.com/tensorchord/envd/pkg/flag"
 	"github.com/tensorchord/envd/pkg/types"
+	"github.com/tensorchord/envd/pkg/util/buildkitutil"
 	"github.com/tensorchord/envd/pkg/util/envutil"
 )
 
@@ -62,12 +63,9 @@ type Client interface {
 }
 
 type generalClient struct {
-	containerName    string
-	image            string
-	mirror           string
-	registry         string
-	enableRegistryCA bool
-	useHTTP          bool
+	containerName  string
+	image          string
+	buildkitConfig *buildkitutil.BuildkitConfig
 
 	driver types.BuilderType
 	socket string
@@ -77,17 +75,14 @@ type generalClient struct {
 }
 
 func NewMobyClient(ctx context.Context, driver types.BuilderType,
-	socket, mirror, registry string, enableRegistryCA bool, useHTTP bool) (Client, error) {
+	socket string, config *buildkitutil.BuildkitConfig) (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,
-		socket:           socket,
-		driver:           driver,
+		containerName:  socket,
+		image:          viper.GetString(flag.FlagBuildkitdImage),
+		buildkitConfig: config,
+		socket:         socket,
+		driver:         driver,
 	}
 	c.logger = logrus.WithFields(logrus.Fields{
 		"container": c.containerName,
@@ -115,16 +110,13 @@ func NewMobyClient(ctx context.Context, driver types.BuilderType,
 }
 
 func NewClient(ctx context.Context, driver types.BuilderType,
-	socket, mirror, registry string, enableRegistryCA bool, useHTTP bool) (Client, error) {
+	socket string, config *buildkitutil.BuildkitConfig) (Client, error) {
 	c := &generalClient{
-		containerName:    socket,
-		image:            viper.GetString(flag.FlagBuildkitdImage),
-		mirror:           mirror,
-		registry:         registry,
-		enableRegistryCA: enableRegistryCA,
-		useHTTP:          useHTTP,
-		socket:           socket,
-		driver:           driver,
+		containerName:  socket,
+		image:          viper.GetString(flag.FlagBuildkitdImage),
+		buildkitConfig: config,
+		socket:         socket,
+		driver:         driver,
 	}
 	c.logger = logrus.WithFields(logrus.Fields{
 		"container": c.containerName,
@@ -179,8 +171,7 @@ func (c *generalClient) maybeStart(ctx context.Context,
 	}
 
 	if client != nil {
-		if _, err := client.StartBuildkitd(ctx, c.image, c.containerName, c.mirror, c.registry,
-			c.enableRegistryCA, c.useHTTP, runningTimeout); err != nil {
+		if _, err := client.StartBuildkitd(ctx, c.image, c.containerName, c.buildkitConfig, runningTimeout); err != nil {
 			return "", err
 		}
 	}
diff --git a/pkg/driver/client.go b/pkg/driver/client.go
index 3b22d73bf..349ed9564 100644
--- a/pkg/driver/client.go
+++ b/pkg/driver/client.go
@@ -20,12 +20,14 @@ import (
 	"time"
 
 	"github.com/docker/docker/api/types"
+
+	"github.com/tensorchord/envd/pkg/util/buildkitutil"
 )
 
 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, registry string, enableRegistryCA, useHTTP bool, timeout time.Duration) (string, error)
+	StartBuildkitd(ctx context.Context, tag, name string, config *buildkitutil.BuildkitConfig, 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
index b17373ed4..67f0bbc42 100644
--- a/pkg/driver/docker/docker.go
+++ b/pkg/driver/docker/docker.go
@@ -36,30 +36,11 @@ import (
 	"github.com/sirupsen/logrus"
 
 	"github.com/tensorchord/envd/pkg/driver"
+	"github.com/tensorchord/envd/pkg/util/buildkitutil"
 	"github.com/tensorchord/envd/pkg/util/fileutil"
 )
 
-const buildkitdMirror = `
-[registry."docker.io"]
-  mirrors = ["%s"]
-`
-const buildkitdHTTP = `
-[registry."docker.io"]
-  http = true
-`
-const buildkitdRegistry = `
-[registry."%s"]
-  http = %v
-`
 const buildkitdCertPath = "/etc/registry"
-const buildkitdWithCA = `
-[registry."docker.io"]
-  mirrors = ["%s"]
-  ca=["/etc/registry/ca.pem"]
-  [[registry."docker.io".keypair]]
-    key="/etc/registry/key.pem"
-    cert="/etc/registry/cert.pem"
-`
 
 var (
 	anchoredIdentifierRegexp = regexp.MustCompile(`^([a-f0-9]{64})$`)
@@ -193,12 +174,11 @@ func (c dockerClient) ResumeContainer(ctx context.Context, name string) (string,
 	return name, nil
 }
 
-func (c dockerClient) StartBuildkitd(ctx context.Context, tag, name, mirror, registry string,
-	enableRegistryCA, useHTTP bool, timeout time.Duration) (string, error) {
+func (c dockerClient) StartBuildkitd(ctx context.Context, tag, name string, bc *buildkitutil.BuildkitConfig, timeout time.Duration) (string, error) {
 	logger := logrus.WithFields(logrus.Fields{
-		"tag":       tag,
-		"container": name,
-		"mirror":    mirror,
+		"tag":             tag,
+		"container":       name,
+		"buildkit-config": bc,
 	})
 	logger.Debug("starting buildkitd")
 	if _, _, err := c.ImageInspectWithRaw(ctx, tag); err != nil {
@@ -226,22 +206,17 @@ func (c dockerClient) StartBuildkitd(ctx context.Context, tag, name, mirror, reg
 		Privileged: true,
 		AutoRemove: true,
 	}
-	var cfg string
-	if mirror != "" {
-		if enableRegistryCA {
-			cfg = fmt.Sprintf(buildkitdWithCA, mirror)
-			hostConfig.Mounts = append(hostConfig.Mounts, mount.Mount{
-				Type:   mount.TypeBind,
-				Source: fileutil.DefaultConfigDir,
-				Target: buildkitdCertPath,
-			})
-		} else {
-			cfg = fmt.Sprintf(buildkitdMirror, mirror)
-		}
-	} else if registry != "" {
-		cfg = fmt.Sprintf(buildkitdRegistry, registry, useHTTP)
-	} else if useHTTP {
-		cfg = buildkitdHTTP
+
+	if bc.SetCA {
+		hostConfig.Mounts = append(hostConfig.Mounts, mount.Mount{
+			Type:   mount.TypeBind,
+			Source: fileutil.DefaultConfigDir,
+			Target: buildkitdCertPath,
+		})
+	}
+	cfg, err := bc.String()
+	if err != nil {
+		return "", errors.Wrap(err, "failed to generate buildkit config")
 	}
 	config.Entrypoint = []string{
 		"/bin/sh",
diff --git a/pkg/driver/nerdctl/nerdctl.go b/pkg/driver/nerdctl/nerdctl.go
index 591a99805..3f14fc7e4 100644
--- a/pkg/driver/nerdctl/nerdctl.go
+++ b/pkg/driver/nerdctl/nerdctl.go
@@ -28,6 +28,7 @@ import (
 	"github.com/sirupsen/logrus"
 
 	"github.com/tensorchord/envd/pkg/driver"
+	"github.com/tensorchord/envd/pkg/util/buildkitutil"
 )
 
 type nerdctlClient struct {
@@ -63,13 +64,12 @@ func (nc *nerdctlClient) Load(ctx context.Context, r io.ReadCloser, quiet bool)
 	return nil
 }
 
-func (nc *nerdctlClient) StartBuildkitd(ctx context.Context, tag, name, mirror, registry string,
-	enableRegistryCA, useHTTP bool, timeout time.Duration) (string, error) {
+func (nc *nerdctlClient) StartBuildkitd(ctx context.Context, tag, name string, bc *buildkitutil.BuildkitConfig, timeout time.Duration) (string, error) {
 	logger := logrus.WithFields(logrus.Fields{
-		"tag":       tag,
-		"container": name,
-		"mirror":    mirror,
-		"driver":    "nerdctl",
+		"tag":             tag,
+		"container":       name,
+		"buildkit-config": bc,
+		"driver":          "nerdctl",
 	})
 	logger.Debug("starting buildkitd")
 
@@ -83,10 +83,11 @@ func (nc *nerdctlClient) StartBuildkitd(ctx context.Context, tag, name, mirror,
 	if !existed {
 		buildkitdCmd := "buildkitd"
 		// TODO: support mirror CA keypair
-		if mirror != "" {
-			cfg := fmt.Sprintf(`
-[registry."docker.io"]
-	mirrors = ["%s"]`, mirror)
+		if bc.Registry != "" || bc.Mirror != "" || bc.UseHTTP {
+			cfg, err := bc.String()
+			if err != nil {
+				return "", errors.Wrap(err, "failed to generate buildkit config")
+			}
 			buildkitdCmd = fmt.Sprintf("mkdir /etc/buildkit && echo '%s' > /etc/buildkit/buildkitd.toml && buildkitd", cfg)
 			logger.Debugf("setting buildkit config: %s", cfg)
 		}
diff --git a/pkg/util/buildkitutil/buildkit.go b/pkg/util/buildkitutil/buildkit.go
new file mode 100644
index 000000000..161dbd966
--- /dev/null
+++ b/pkg/util/buildkitutil/buildkit.go
@@ -0,0 +1,48 @@
+// Copyright 2023 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package buildkitutil
+
+import (
+	"strings"
+	"text/template"
+)
+
+const buildkitConfigTemplate = `
+[registry."{{ if .Registry }}{{ .Registry }}{{ else }}docker.io{{ end }}"]{{ if .Mirror }}
+  mirrors = ["{{ .Mirror }}"]{{ end }}
+  http = {{ .UseHTTP }}
+  {{ if .SetCA}}ca=["/etc/registry/ca.pem"]
+  [[registry."{{ if .Registry }}{{ .Registry }}{{ else }}docker.io{{ end }}".keypair]]
+	key="/etc/registry/key.pem"
+	cert="/etc/registry/cert.pem"
+  {{ end }}
+`
+
+type BuildkitConfig struct {
+	Registry string
+	Mirror   string
+	UseHTTP  bool
+	SetCA    bool
+}
+
+func (c *BuildkitConfig) String() (string, error) {
+	tmpl, err := template.New("buildkitConfig").Parse(buildkitConfigTemplate)
+	if err != nil {
+		return "", err
+	}
+	var config strings.Builder
+	err = tmpl.Execute(&config, c)
+	return config.String(), err
+}


Test patch:
diff --git a/pkg/util/buildkitutil/buildkit_test.go b/pkg/util/buildkitutil/buildkit_test.go
new file mode 100644
index 000000000..f613e1d16
--- /dev/null
+++ b/pkg/util/buildkitutil/buildkit_test.go
@@ -0,0 +1,69 @@
+// Copyright 2023 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package buildkitutil
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBuildkitWithRegistry(t *testing.T) {
+	testCases := []struct {
+		bc       BuildkitConfig
+		expected string
+	}{
+		{
+			BuildkitConfig{
+				Registry: "registry.example.com",
+				Mirror:   "https://mirror.example.com",
+				UseHTTP:  true,
+			},
+			`
+[registry."registry.example.com"]
+  mirrors = ["https://mirror.example.com"]
+  http = true
+`,
+		},
+		{
+			BuildkitConfig{
+				Registry: "registry.example.com",
+				SetCA:    true,
+			},
+			`
+[registry."registry.example.com"]
+  http = false
+  ca=["/etc/registry/ca.pem"]
+  [[registry."registry.example.com".keypair]]
+	key="/etc/registry/key.pem"
+	cert="/etc/registry/cert.pem"
+`,
+		},
+		{
+			BuildkitConfig{},
+			`
+[registry."docker.io"]
+  http = false
+			`,
+		},
+	}
+
+	for _, tc := range testCases {
+		config, err := tc.bc.String()
+		assert.NoError(t, err)
+		assert.Equal(t, strings.TrimSpace(tc.expected), strings.TrimSpace(config))
+	}
+}