feat: support mirror ca for buildkitd (#1602)

Signed-off-by: Keming <[email protected]>

pkg/app/bootstrap.go

@@ -54,6 +54,11 @@ var CommandBootstrap = &cli.Command{
 			Usage:   "Dockerhub mirror to use",
 			Aliases: []string{"m"},
 		},
+		&cli.StringFlag{
+			Name:    "registry-ca-keypair",
+			Usage:   "Specify the ca/key/cert file path for the private registry (format: 'ca=/etc/config/ca.pem,key=/etc/config/key.pem,cert=/etc/config/cert.pem')",
+			Aliases: []string{"ca"},
+		},
 		&cli.StringSliceFlag{
 			Name: "ssh-keypair",
 			Usage: fmt.Sprintf("Manually specify ssh key pair as `publicKey,privateKey`. Envd will generate a keypair at %s and %s if not specified",
@@ -72,6 +77,9 @@ func bootstrap(clicontext *cli.Context) error {
 	}{{
 		"SSH Key",
 		sshKey,
+	}, {
+		"registry CA keypair",
+		registryCA,
 	}, {
 		"autocomplete",
 		autocomplete,
@@ -92,6 +100,61 @@ func bootstrap(clicontext *cli.Context) error {
 	return nil
 }
 
+func registryCA(clicontext *cli.Context) error {
+	ca := clicontext.String("registry-ca-keypair")
+	if len(ca) == 0 {
+		return nil
+	}
+	mirror := clicontext.String("dockerhub-mirror")
+	if len(mirror) == 0 {
+		return errors.New("`registry-ca-keypair` should be used with `dockerhub-mirror`")
+	}
+
+	// parse ca/key/cert
+	kvPairs := strings.Split(ca, ",")
+	if len(kvPairs) != 3 {
+		return errors.New("`registry-ca-keypair` requires ca/key/cert 3 part separated by ','")
+	}
+	names := []string{"ca", "cert", "key"}
+	for _, pair := range kvPairs {
+		kv := strings.SplitN(pair, "=", 2)
+		index := -1
+		for i, name := range names {
+			if name == kv[0] {
+				index = i
+				break
+			}
+		}
+		if index == -1 {
+			return errors.Newf("parse error: `%s` is not a valid ca/key/cert key or it's duplicated")
+		}
+		exist, err := fileutil.FileExists(kv[1])
+		if err != nil {
+			return errors.Wrap(err, fmt.Sprintf("failed to parse file path %s", pair))
+		}
+		if !exist {
+			return errors.Newf("file %s doesn't exist", kv[1])
+		}
+		path, err := fileutil.ConfigFile(fmt.Sprintf("registry_%s.pem", kv[0]))
+		if err != nil {
+			return errors.Wrap(err, "failed to get the envd config file path")
+		}
+		content, err := os.ReadFile(kv[1])
+		if err != nil {
+			return errors.Wrap(err, "failed to read the file")
+		}
+		if err = os.WriteFile(path, content, 0644); err != nil {
+			return errors.Wrap(err, "failed to store the CA file")
+		}
+		names = append(names[:index], names[index+1:]...)
+	}
+
+	if len(names) != 0 {
+		return errors.Newf("registry %s are not provided", names)
+	}
+	return nil
+}
+
 func sshKey(clicontext *cli.Context) error {
 	sshKeyPair := clicontext.StringSlice("ssh-keypair")
 
@@ -218,13 +281,15 @@ func buildkit(clicontext *cli.Context) error {
 	var bkClient buildkitd.Client
 	if c.Builder == types.BuilderTypeMoby {
 		bkClient, err = buildkitd.NewMobyClient(clicontext.Context,
-			c.Builder, c.BuilderAddress, clicontext.String("dockerhub-mirror"))
+			c.Builder, c.BuilderAddress, clicontext.String("dockerhub-mirror"),
+			clicontext.IsSet("registry-ca-keypair"))
 		if err != nil {
 			return errors.Wrap(err, "failed to create moby buildkit client")
 		}
 	} else {
 		bkClient, err = buildkitd.NewClient(clicontext.Context,
-			c.Builder, c.BuilderAddress, clicontext.String("dockerhub-mirror"))
+			c.Builder, c.BuilderAddress, clicontext.String("dockerhub-mirror"),
+			clicontext.IsSet("registry-ca-keypair"))
 		if err != nil {
 			return errors.Wrap(err, "failed to create buildkit client")
 		}

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, "")
+			c.Builder, c.BuilderAddress, "", 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, "")
+			c.Builder, c.BuilderAddress, "", false)
 		if err != nil {
 			return errors.Wrap(err, "failed to create buildkit client")
 		}

pkg/builder/build.go

@@ -88,13 +88,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, "")
+			c.Builder, c.BuilderAddress, "", 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, "")
+			c.Builder, c.BuilderAddress, "", false)
 		if err != nil {
 			return nil, errors.Wrap(err, "failed to create buildkit client")
 		}

pkg/buildkitd/buildkitd.go

@@ -62,9 +62,10 @@ type Client interface {
 }
 
 type generalClient struct {
-	containerName string
-	image         string
-	mirror        string
+	containerName    string
+	image            string
+	mirror           string
+	enableRegistryCA bool
 
 	driver types.BuilderType
 	socket string
@@ -74,15 +75,16 @@ type generalClient struct {
 }
 
 func NewMobyClient(ctx context.Context, driver types.BuilderType,
-	socket, mirror string) (Client, error) {
+	socket, mirror string, enableRegistryCA bool) (Client, error) {
 	logrus.Debug("getting moby buildkit client")
 	c := &generalClient{
-		containerName: socket,
-		image:         viper.GetString(flag.FlagBuildkitdImage),
-		mirror:        mirror,
+		containerName:    socket,
+		image:            viper.GetString(flag.FlagBuildkitdImage),
+		mirror:           mirror,
+		enableRegistryCA: enableRegistryCA,
+		socket:           socket,
+		driver:           driver,
 	}
-	c.socket = socket
-	c.driver = driver
 	c.logger = logrus.WithFields(logrus.Fields{
 		"container": c.containerName,
 		"image":     c.image,
@@ -109,14 +111,15 @@ func NewMobyClient(ctx context.Context, driver types.BuilderType,
 }
 
 func NewClient(ctx context.Context, driver types.BuilderType,
-	socket, mirror string) (Client, error) {
+	socket, mirror string, enableRegistryCA bool) (Client, error) {
 	c := &generalClient{
-		containerName: socket,
-		image:         viper.GetString(flag.FlagBuildkitdImage),
-		mirror:        mirror,
+		containerName:    socket,
+		image:            viper.GetString(flag.FlagBuildkitdImage),
+		mirror:           mirror,
+		enableRegistryCA: enableRegistryCA,
+		socket:           socket,
+		driver:           driver,
 	}
-	c.socket = socket
-	c.driver = driver
 	c.logger = logrus.WithFields(logrus.Fields{
 		"container": c.containerName,
 		"image":     c.image,
@@ -170,8 +173,8 @@ func (c *generalClient) maybeStart(ctx context.Context,
 	}
 
 	if client != nil {
-		if _, err := client.StartBuildkitd(ctx,
-			c.image, c.containerName, c.mirror, runningTimeout); err != nil {
+		if _, err := client.StartBuildkitd(ctx, c.image, c.containerName, c.mirror,
+			c.enableRegistryCA, runningTimeout); err != nil {
 			return "", err
 		}
 	}

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, timeout time.Duration) (string, error)
+	StartBuildkitd(ctx context.Context, tag, name, mirror string, enableRegistryCA bool, timeout time.Duration) (string, error)
 
 	Exec(ctx context.Context, cname string, cmd []string) error
 

pkg/driver/docker/docker.go

@@ -29,14 +29,30 @@ import (
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/api/types/container"
 	"github.com/docker/docker/api/types/filters"
+	"github.com/docker/docker/api/types/mount"
 	"github.com/docker/docker/client"
 	"github.com/docker/docker/pkg/jsonmessage"
 	"github.com/moby/term"
 	"github.com/sirupsen/logrus"
 
 	"github.com/tensorchord/envd/pkg/driver"
+	"github.com/tensorchord/envd/pkg/util/fileutil"
 )
 
+const buildkitdMirror = `
+[registry."docker.io"]
+  mirrors = ["%s"]
+`
+const buildkitdCertPath = "/etc/registry"
+const buildkitdRegistry = `
+[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})$`)
 	waitingInterval          = 1 * time.Second
@@ -169,8 +185,8 @@ func (c dockerClient) ResumeContainer(ctx context.Context, name string) (string,
 	return name, nil
 }
 
-func (c dockerClient) StartBuildkitd(ctx context.Context,
-	tag, name, mirror string, timeout time.Duration) (string, error) {
+func (c dockerClient) StartBuildkitd(ctx context.Context, tag, name, mirror string,
+	enableRegistryCA bool, timeout time.Duration) (string, error) {
 	logger := logrus.WithFields(logrus.Fields{
 		"tag":       tag,
 		"container": name,
@@ -198,20 +214,29 @@ func (c dockerClient) StartBuildkitd(ctx context.Context,
 	config := &container.Config{
 		Image: tag,
 	}
+	hostConfig := &container.HostConfig{
+		Privileged: true,
+		AutoRemove: true,
+	}
 	if mirror != "" {
-		cfg := fmt.Sprintf(`
-[registry."docker.io"]
-	mirrors = ["%s"]`, mirror)
+		var cfg string
+		if enableRegistryCA {
+			cfg = fmt.Sprintf(buildkitdRegistry, mirror)
+			hostConfig.Mounts = append(hostConfig.Mounts, mount.Mount{
+				Type:   mount.TypeBind,
+				Source: fileutil.DefaultConfigDir,
+				Target: buildkitdCertPath,
+			})
+		} else {
+			cfg = fmt.Sprintf(buildkitdMirror, mirror)
+		}
 		config.Entrypoint = []string{
 			"/bin/sh",
 			"-c",
 			fmt.Sprintf("mkdir /etc/buildkit && echo '%s' > /etc/buildkit/buildkitd.toml && buildkitd", cfg),
 		}
 		logger.Debugf("setting buildkit config: %s", cfg)
 	}
-	hostConfig := &container.HostConfig{
-		Privileged: true,
-	}
 	created, _ := c.Exists(ctx, name)
 	if created {
 		err := c.ContainerStart(ctx, name, types.ContainerStartOptions{})

pkg/driver/nerdctl/nerdctl.go

@@ -63,8 +63,8 @@ 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, timeout time.Duration) (string, error) {
+func (nc *nerdctlClient) StartBuildkitd(ctx context.Context, tag, name, mirror string,
+	enableRegistryCA bool, timeout time.Duration) (string, error) {
 	logger := logrus.WithFields(logrus.Fields{
 		"tag":       tag,
 		"container": name,
@@ -82,6 +82,7 @@ func (nc *nerdctlClient) StartBuildkitd(ctx context.Context,
 	existed, _ := nc.containerExists(ctx, name)
 	if !existed {
 		buildkitdCmd := "buildkitd"
+		// TODO: support mirror CA keypair
 		if mirror != "" {
 			cfg := fmt.Sprintf(`
 [registry."docker.io"]