From b483177ee2089fd8284258840e2526370c15eab3 Mon Sep 17 00:00:00 2001
From: ktock <ktokunaga.mail@gmail.com>
Date: Wed, 14 Jul 2021 15:52:04 +0900
Subject: [PATCH] Support custom compressor for walking differ

Signed-off-by: Kohei Tokunaga <ktokunaga.mail@gmail.com>
---
 diff/diff.go           | 15 +++++++++++++++
 diff/walking/differ.go | 38 ++++++++++++++++++++++++++------------
 2 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/diff/diff.go b/diff/diff.go
index 17aab616e..235d6377c 100644
--- a/diff/diff.go
+++ b/diff/diff.go
@@ -18,6 +18,7 @@ package diff
 
 import (
 	"context"
+	"io"
 
 	"github.com/containerd/containerd/mount"
 	"github.com/gogo/protobuf/types"
@@ -37,6 +38,12 @@ type Config struct {
 
 	// Labels are the labels to apply to the generated content
 	Labels map[string]string
+
+	// Compressor is a function to compress the diff stream
+	// instead of the default gzip compressor. Differ passes
+	// the MediaType of the target diff content to the compressor.
+	// When using this config, MediaType must be specified as well.
+	Compressor func(dest io.Writer, mediaType string) (io.WriteCloser, error)
 }
 
 // Opt is used to configure a diff operation
@@ -71,6 +78,14 @@ type Applier interface {
 	Apply(ctx context.Context, desc ocispec.Descriptor, mount []mount.Mount, opts ...ApplyOpt) (ocispec.Descriptor, error)
 }
 
+// WithCompressor sets the function to be used to compress the diff stream.
+func WithCompressor(f func(dest io.Writer, mediaType string) (io.WriteCloser, error)) Opt {
+	return func(c *Config) error {
+		c.Compressor = f
+		return nil
+	}
+}
+
 // WithMediaType sets the media type to use for creating the diff, without
 // specifying the differ will choose a default.
 func WithMediaType(m string) Opt {
diff --git a/diff/walking/differ.go b/diff/walking/differ.go
index 8f6362992..c5d142bbd 100644
--- a/diff/walking/differ.go
+++ b/diff/walking/differ.go
@@ -65,17 +65,24 @@ func (s *walkingDiff) Compare(ctx context.Context, lower, upper []mount.Mount, o
 		}
 	}
 
-	if config.MediaType == "" {
-		config.MediaType = ocispec.MediaTypeImageLayerGzip
-	}
-
 	var isCompressed bool
-	switch config.MediaType {
-	case ocispec.MediaTypeImageLayer:
-	case ocispec.MediaTypeImageLayerGzip:
+	if config.Compressor != nil {
+		if config.MediaType == "" {
+			return emptyDesc, errors.New("media type must be explicitly specified when using custom compressor")
+		}
 		isCompressed = true
-	default:
-		return emptyDesc, errors.Wrapf(errdefs.ErrNotImplemented, "unsupported diff media type: %v", config.MediaType)
+	} else {
+		if config.MediaType == "" {
+			config.MediaType = ocispec.MediaTypeImageLayerGzip
+		}
+
+		switch config.MediaType {
+		case ocispec.MediaTypeImageLayer:
+		case ocispec.MediaTypeImageLayerGzip:
+			isCompressed = true
+		default:
+			return emptyDesc, errors.Wrapf(errdefs.ErrNotImplemented, "unsupported diff media type: %v", config.MediaType)
+		}
 	}
 
 	var ocidesc ocispec.Descriptor
@@ -118,9 +125,16 @@ func (s *walkingDiff) Compare(ctx context.Context, lower, upper []mount.Mount, o
 			if isCompressed {
 				dgstr := digest.SHA256.Digester()
 				var compressed io.WriteCloser
-				compressed, errOpen = compression.CompressStream(cw, compression.Gzip)
-				if errOpen != nil {
-					return errors.Wrap(errOpen, "failed to get compressed stream")
+				if config.Compressor != nil {
+					compressed, errOpen = config.Compressor(cw, config.MediaType)
+					if errOpen != nil {
+						return errors.Wrap(errOpen, "failed to get compressed stream")
+					}
+				} else {
+					compressed, errOpen = compression.CompressStream(cw, compression.Gzip)
+					if errOpen != nil {
+						return errors.Wrap(errOpen, "failed to get compressed stream")
+					}
 				}
 				errOpen = archive.WriteDiff(ctx, io.MultiWriter(compressed, dgstr.Hash()), lowerRoot, upperRoot)
 				compressed.Close()