Add zstd compression to rules_oci

Author: egor-karavaev-datadog

MODULE.bazel

@@ -25,6 +25,7 @@ use_repo(
     "com_github_containerd_containerd",
     "com_github_containerd_log",
     "com_github_docker_docker_credential_helpers",
+    "com_github_klauspost_compress",
     "com_github_mitchellh_go_homedir",
     "com_github_opencontainers_go_digest",
     "com_github_opencontainers_image_spec",

docs/docs.md

@@ -173,7 +173,8 @@ The config file named after the rule, os, and arch
 <pre>
 load("@rules_oci//oci:defs.bzl", "oci_image_layer")
 
-oci_image_layer(<a href="#oci_image_layer-name">name</a>, <a href="#oci_image_layer-directory">directory</a>, <a href="#oci_image_layer-files">files</a>, <a href="#oci_image_layer-file_map">file_map</a>, <a href="#oci_image_layer-mode_map">mode_map</a>, <a href="#oci_image_layer-owner_map">owner_map</a>, <a href="#oci_image_layer-symlinks">symlinks</a>, <a href="#oci_image_layer-kwargs">kwargs</a>)
+oci_image_layer(<a href="#oci_image_layer-name">name</a>, <a href="#oci_image_layer-directory">directory</a>, <a href="#oci_image_layer-files">files</a>, <a href="#oci_image_layer-file_map">file_map</a>, <a href="#oci_image_layer-mode_map">mode_map</a>, <a href="#oci_image_layer-owner_map">owner_map</a>, <a href="#oci_image_layer-symlinks">symlinks</a>, <a href="#oci_image_layer-compression_method">compression_method</a>,
+                <a href="#oci_image_layer-kwargs">kwargs</a>)
 </pre>
 
 Creates a tarball and an OCI descriptor for it
@@ -190,6 +191,7 @@ Creates a tarball and an OCI descriptor for it
 | <a id="oci_image_layer-mode_map"></a>mode_map |  Dictionary of file location in tarball -> mode int (e.g. 0o755)   |  `None` |
 | <a id="oci_image_layer-owner_map"></a>owner_map |  Dictionary of file location in tarball -> owner:group string (e.g. '501:501')   |  `None` |
 | <a id="oci_image_layer-symlinks"></a>symlinks |  Dictionary of symlink -> target entries to place in the tarball   |  `None` |
+| <a id="oci_image_layer-compression_method"></a>compression_method |  A string, currently supports "gzip" and "zstd", defaults to "gzip"   |  `None` |
 | <a id="oci_image_layer-kwargs"></a>kwargs |  Additional arguments to pass to the rule, e.g. tags or visibility   |  none |
 
 

go.mod

@@ -10,6 +10,7 @@ require (
 	github.com/containerd/containerd v1.7.20
 	github.com/containerd/log v0.1.0
 	github.com/docker/docker-credential-helpers v0.8.1
+	github.com/klauspost/compress v1.17.8
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/opencontainers/go-digest v1.0.0
 	github.com/opencontainers/image-spec v1.1.0
@@ -41,7 +42,6 @@ require (
 	github.com/go-logr/logr v1.4.2 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/gorilla/mux v1.8.1 // indirect
-	github.com/klauspost/compress v1.17.8 // indirect
 	github.com/kr/text v0.2.0 // indirect
 	github.com/moby/locker v1.0.1 // indirect
 	github.com/pkg/errors v0.9.1 // indirect

go/cmd/ocitool/BUILD.bazel

@@ -31,6 +31,7 @@ go_library(
         "@com_github_containerd_containerd//images:go_default_library",
         "@com_github_containerd_containerd//platforms:go_default_library",
         "@com_github_containerd_log//:go_default_library",
+        "@com_github_klauspost_compress//zstd:go_default_library",
         "@com_github_opencontainers_go_digest//:go_default_library",
         "@com_github_opencontainers_image_spec//specs-go:go_default_library",
         "@com_github_opencontainers_image_spec//specs-go/v1:go_default_library",

go/cmd/ocitool/createlayer_cmd.go

@@ -17,6 +17,7 @@ import (
 	"github.com/DataDog/rules_oci/go/internal/tarutil"
 	"github.com/DataDog/rules_oci/go/pkg/layer"
 	"github.com/DataDog/rules_oci/go/pkg/ociutil"
+	"github.com/klauspost/compress/zstd"
 	"github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/urfave/cli/v2"
@@ -35,11 +36,32 @@ func CreateLayerCmd(c *cli.Context) error {
 
 	digester := digest.SHA256.Digester()
 	wc := ociutil.NewWriterCounter(io.MultiWriter(out, digester.Hash()))
-	gw := gzip.NewWriter(wc)
-	gw.Name = path.Base(out.Name())
-	defer gw.Close()
 
-	tw := tar.NewWriter(gw)
+	var compressWriter io.Writer
+	var compressCloser io.Closer
+	var mediaType string
+	switch config.CompressionMethod {
+	case "gzip":
+		gzipWriter := gzip.NewWriter(wc)
+		gzipWriter.Name = path.Base(out.Name())
+		compressWriter = gzipWriter
+		compressCloser = gzipWriter
+		mediaType = ocispec.MediaTypeImageLayerGzip
+	case "zstd":
+		zstdWriter, err := zstd.NewWriter(wc,
+			zstd.WithEncoderLevel(zstd.SpeedBestCompression))
+		if err != nil {
+			return err
+		}
+		compressWriter = zstdWriter
+		compressCloser = zstdWriter
+		mediaType = ocispec.MediaTypeImageLayerZstd
+	default:
+		return fmt.Errorf("uknown compress method %s", config.CompressionMethod)
+	}
+	defer compressCloser.Close()
+
+	tw := tar.NewWriter(compressWriter)
 	defer tw.Close()
 
 	slices.Sort(config.Files)
@@ -139,11 +161,12 @@ func CreateLayerCmd(c *cli.Context) error {
 	// Need to flush before we count bytes and digest, might as well close since
 	// it's not needed anymore.
 	tw.Close()
-	gw.Close()
+	compressCloser.Close()
+	out.Close()
 
 	desc := ocispec.Descriptor{
 		Digest:    digester.Digest(),
-		MediaType: ocispec.MediaTypeImageLayerGzip,
+		MediaType: mediaType,
 		Size:      int64(wc.Count()),
 	}
 
@@ -164,15 +187,16 @@ func CreateLayerCmd(c *cli.Context) error {
 }
 
 type createLayerConfig struct {
-	BazelLabel     string            `json:"bazel-label" toml:"bazel-label" yaml:"bazel-label"`
-	Descriptor     string            `json:"outd" toml:"outd" yaml:"outd"`
-	Directory      string            `json:"dir" toml:"dir" yaml:"dir"`
-	FileMapping    map[string]string `json:"file-map" toml:"file-map" yaml:"file-map"`
-	Files          []string          `json:"file" toml:"file" yaml:"file"`
-	ModeMapping    map[string]int64  `json:"mode-map" toml:"mode-map" yaml:"mode-map"`
-	OutputLayer    string            `json:"out" toml:"out" yaml:"out"`
-	OwnerMapping   map[string]string `json:"owner-map" toml:"owner-map" yaml:"owner-map"`
-	SymlinkMapping map[string]string `json:"symlink" toml:"symlink" yaml:"symlink"`
+	BazelLabel        string            `json:"bazel-label" toml:"bazel-label" yaml:"bazel-label"`
+	Descriptor        string            `json:"outd" toml:"outd" yaml:"outd"`
+	Directory         string            `json:"dir" toml:"dir" yaml:"dir"`
+	FileMapping       map[string]string `json:"file-map" toml:"file-map" yaml:"file-map"`
+	Files             []string          `json:"file" toml:"file" yaml:"file"`
+	ModeMapping       map[string]int64  `json:"mode-map" toml:"mode-map" yaml:"mode-map"`
+	OutputLayer       string            `json:"out" toml:"out" yaml:"out"`
+	OwnerMapping      map[string]string `json:"owner-map" toml:"owner-map" yaml:"owner-map"`
+	SymlinkMapping    map[string]string `json:"symlink" toml:"symlink" yaml:"symlink"`
+	CompressionMethod string            `json:"compression-method" toml:"compression-method" yaml:"compression-method"`
 }
 
 func newCreateLayerConfig(c *cli.Context) (*createLayerConfig, error) {
@@ -184,16 +208,23 @@ func newCreateLayerConfig(c *cli.Context) (*createLayerConfig, error) {
 		}
 		modeMapping[path] = mode
 	}
+
+	compressionMethod := c.String("compression-method")
+	if compressionMethod == "" {
+		compressionMethod = "gzip"
+	}
+
 	return &createLayerConfig{
-		BazelLabel:     c.String("bazel-label"),
-		Descriptor:     c.String("outd"),
-		Directory:      c.String("dir"),
-		FileMapping:    c.Generic("file-map").(*flagutil.KeyValueFlag).Map,
-		Files:          c.StringSlice("file"),
-		ModeMapping:    modeMapping,
-		OutputLayer:    c.String("out"),
-		OwnerMapping:   c.Generic("owner-map").(*flagutil.KeyValueFlag).Map,
-		SymlinkMapping: c.Generic("symlink").(*flagutil.KeyValueFlag).Map,
+		BazelLabel:        c.String("bazel-label"),
+		Descriptor:        c.String("outd"),
+		Directory:         c.String("dir"),
+		FileMapping:       c.Generic("file-map").(*flagutil.KeyValueFlag).Map,
+		Files:             c.StringSlice("file"),
+		ModeMapping:       modeMapping,
+		OutputLayer:       c.String("out"),
+		OwnerMapping:      c.Generic("owner-map").(*flagutil.KeyValueFlag).Map,
+		SymlinkMapping:    c.Generic("symlink").(*flagutil.KeyValueFlag).Map,
+		CompressionMethod: compressionMethod,
 	}, nil
 }
 

go/cmd/ocitool/main.go

@@ -78,6 +78,9 @@ var app = &cli.App{
 					Name:  "mode-map",
 					Value: &flagutil.KeyValueFlag{},
 				},
+				&cli.StringFlag{
+					Name: "compression-method",
+				},
 			},
 		},
 		{

go/pkg/ociutil/BUILD.bazel

@@ -35,6 +35,7 @@ go_library(
         "@com_github_containerd_containerd//reference/docker:go_default_library",
         "@com_github_containerd_containerd//remotes:go_default_library",
         "@com_github_containerd_containerd//remotes/docker:go_default_library",
+        "@com_github_klauspost_compress//zstd:go_default_library",
         "@com_github_opencontainers_go_digest//:go_default_library",
         "@com_github_opencontainers_image_spec//specs-go/v1:go_default_library",
         "@com_github_sethvargo_go_retry//:go_default_library",

go/pkg/ociutil/diff.go

@@ -4,8 +4,10 @@ import (
 	"compress/gzip"
 	"context"
 	"fmt"
+	"io"
 
 	"github.com/containerd/containerd/content"
+	"github.com/klauspost/compress/zstd"
 	"github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 )
@@ -14,23 +16,31 @@ import (
 // image config. If a layer is uncompressed, the diff ID is simply the digest but if the
 // layer is compressed, we must uncompress the file and acquire the digest.
 func GetLayerDiffID(ctx context.Context, store content.Store, desc ocispec.Descriptor) (digest.Digest, error) {
+	if desc.MediaType != ocispec.MediaTypeImageLayerGzip && desc.MediaType != ocispec.MediaTypeImageLayerZstd {
+		return desc.Digest, nil
+	}
+
+	r, err := store.ReaderAt(ctx, desc)
+	if err != nil {
+		return "", fmt.Errorf("failed to get reader for layer: %w", err)
+	}
+	defer r.Close()
+
+	var cr io.Reader
 	switch desc.MediaType {
 	case ocispec.MediaTypeImageLayerGzip:
-		r, err := store.ReaderAt(ctx, desc)
+		cr, err = gzip.NewReader(&readerAtReader{ReaderAt: r})
 		if err != nil {
-			return "", fmt.Errorf("failed to get reader for layer: %w", err)
+			return "", fmt.Errorf("failed to get gzip reader for layer: %w", err)
 		}
-		defer r.Close()
-
-		gr, err := gzip.NewReader(&readerAtReader{ReaderAt: r})
+	case ocispec.MediaTypeImageLayerZstd:
+		cr, err = zstd.NewReader(&readerAtReader{ReaderAt: r})
 		if err != nil {
-			return "", fmt.Errorf("failed to get gzip reader for layer: %w", err)
+			return "", fmt.Errorf("failed to get zstd reader for layer: %w", err)
 		}
-
-		return digest.SHA256.FromReader(gr)
-	default:
-		return desc.Digest, nil
 	}
+
+	return digest.SHA256.FromReader(cr)
 }
 
 type readerAtReader struct {

oci/layer.bzl

@@ -11,6 +11,7 @@ def oci_image_layer(
         mode_map = None,  # dict[str, int] | None,
         owner_map = None,  # dict[str, str] | None,
         symlinks = None,  # dict[str, str] | None,
+        compression_method = None,  # str | None
         **kwargs):
     """ Creates a tarball and an OCI descriptor for it
 
@@ -22,6 +23,7 @@ def oci_image_layer(
         mode_map: Dictionary of file location in tarball -> mode int (e.g. 0o755)
         owner_map: Dictionary of file location in tarball -> owner:group string (e.g. '501:501')
         symlinks: Dictionary of symlink -> target entries to place in the tarball
+        compression_method: A string, currently supports "gzip" and "zstd", defaults to "gzip"
         **kwargs: Additional arguments to pass to the rule, e.g. tags or visibility
     """
     mode_map = {k: str(v) for k, v in mode_map.items()} if mode_map else {}
@@ -34,22 +36,41 @@ def oci_image_layer(
         mode_map = mode_map,
         owner_map = owner_map,
         symlinks = symlinks,
+        compression_method = compression_method,
         **kwargs
     )
 
+def _get_output_file_name(name, compression_method):
+    file_extension = {
+        "gzip": "gz",
+        "zstd": "zst",
+    }.get(compression_method)
+
+    if not file_extension:
+        fail("Unknown compression_method: {}".format(compression_method))
+
+    return "{}-layer.tar.{}".format(name, file_extension)
+
 def _impl(ctx):
     toolchain = ctx.toolchains["@com_github_datadog_rules_oci//oci:toolchain"]
 
     descriptor_file = ctx.actions.declare_file("{}.descriptor.json".format(ctx.label.name))
 
+    compression_method = ctx.attr.compression_method
+    if not compression_method:
+        compression_method = "gzip"
+
+    output_file = ctx.actions.declare_file(_get_output_file_name(ctx.label.name, compression_method))
+
     ctx.actions.run(
         executable = toolchain.sdk.ocitool,
         arguments = [
                         "create-layer",
                         "--bazel-label={}".format(ctx.label),
                         "--dir={}".format(ctx.attr.directory),
-                        "--out={}".format(ctx.outputs.layer.path),
+                        "--out={}".format(output_file.path),
                         "--outd={}".format(descriptor_file.path),
+                        "--compression-method={}".format(compression_method),
                     ] +
                     ["--file-map={}={}".format(k.files.to_list()[0].path, v) for k, v in ctx.attr.file_map.items()] +
                     ["--file={}".format(f.path) for f in ctx.files.files] +
@@ -60,17 +81,24 @@ def _impl(ctx):
         mnemonic = "OCIImageCreateLayer",
         outputs = [
             descriptor_file,
-            ctx.outputs.layer,
+            output_file,
         ],
     )
 
     return [
         OCIDescriptor(
             descriptor_file = descriptor_file,
-            file = ctx.outputs.layer,
+            file = output_file,
         ),
     ]
 
+def _get_outputs(name, compression_method):
+    if not compression_method:
+        compression_method = "gzip"
+    return {
+        "layer": _get_output_file_name(name, compression_method),
+    }
+
 _oci_image_layer = rule(
     implementation = _impl,
     doc = "Create a tarball and an OCI descriptor for it",
@@ -81,9 +109,8 @@ _oci_image_layer = rule(
         "mode_map": attr.string_dict(),
         "owner_map": attr.string_dict(),
         "symlinks": attr.string_dict(),
+        "compression_method": attr.string(),
     },
+    outputs = _get_outputs,
     toolchains = ["@com_github_datadog_rules_oci//oci:toolchain"],
-    outputs = {
-        "layer": "%{name}-layer.tar.gz",
-    },
 )

tests/go-multiarch-image/BUILD.bazel

@@ -27,6 +27,18 @@ go_multiarch_image(
     visibility = ["//visibility:public"],
 )
 
+go_multiarch_image(
+    name = "zstd-image",
+    archs = [
+        "amd64",
+        "arm64",
+    ],
+    base = "@ubuntu_focal//image",
+    compression_method = "zstd",
+    embed = [":go_default_library"],
+    visibility = ["//visibility:public"],
+)
+
 oci_push(
     name = "push",
     manifest = ":image",