exporter: workaround containerd client issue with grpc size limits

The grpc client for containerd won't properly break up the writes for a
file when used in combination with `bytes.NewReader` and
`content.WriteBlob`.

To workaround the issue, this wraps the container image exporter for
containerd with a bytes reader that will chunk the `Write` calls instead
of relying on the client to do the chunking.

See containerd/containerd#11440 for the
upstream containerd issue.

Signed-off-by: Jonathan A. Sternberg <[email protected]>

Author: jsternberg

exporter/containerimage/writer.go

@@ -1,10 +1,10 @@
 package containerimage
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"reflect"
 	"strconv"
 	"strings"
@@ -342,7 +342,7 @@ func (ic *ImageWriter) Commit(ctx context.Context, inp *exporter.Source, session
 	}
 	idxDone := progress.OneOff(ctx, "exporting manifest list "+idxDigest.String())
 
-	if err := content.WriteBlob(ctx, ic.opt.ContentStore, idxDigest.String(), bytes.NewReader(idxBytes), idxDesc, content.WithLabels(labels)); err != nil {
+	if err := content.WriteBlob(ctx, ic.opt.ContentStore, idxDigest.String(), newBytesReader(idxBytes), idxDesc, content.WithLabels(labels)); err != nil {
 		return nil, idxDone(errors.Wrapf(err, "error writing manifest list blob %s", idxDigest))
 	}
 	idxDone(nil)
@@ -530,7 +530,7 @@ func (ic *ImageWriter) commitDistributionManifest(ctx context.Context, opts *Ima
 	}
 	mfstDone := progress.OneOff(ctx, "exporting manifest "+mfstDigest.String())
 
-	if err := content.WriteBlob(ctx, ic.opt.ContentStore, mfstDigest.String(), bytes.NewReader(mfstJSON), mfstDesc, content.WithLabels((labels))); err != nil {
+	if err := content.WriteBlob(ctx, ic.opt.ContentStore, mfstDigest.String(), newBytesReader(mfstJSON), mfstDesc, content.WithLabels((labels))); err != nil {
 		return nil, nil, mfstDone(errors.Wrapf(err, "error writing manifest blob %s", mfstDigest))
 	}
 	mfstDone(nil)
@@ -542,7 +542,7 @@ func (ic *ImageWriter) commitDistributionManifest(ctx context.Context, opts *Ima
 	}
 	configDone := progress.OneOff(ctx, "exporting config "+configDigest.String())
 
-	if err := content.WriteBlob(ctx, ic.opt.ContentStore, configDigest.String(), bytes.NewReader(config), configDesc); err != nil {
+	if err := content.WriteBlob(ctx, ic.opt.ContentStore, configDigest.String(), newBytesReader(config), configDesc); err != nil {
 		return nil, nil, configDone(errors.Wrap(err, "error writing config blob"))
 	}
 	configDone(nil)
@@ -584,7 +584,7 @@ func (ic *ImageWriter) commitAttestationsManifest(ctx context.Context, opts *Ima
 			},
 		}
 
-		if err := content.WriteBlob(ctx, ic.opt.ContentStore, digest.String(), bytes.NewReader(data), desc); err != nil {
+		if err := content.WriteBlob(ctx, ic.opt.ContentStore, digest.String(), newBytesReader(data), desc); err != nil {
 			return nil, errors.Wrapf(err, "error writing data blob %s", digest)
 		}
 		layers[i] = desc
@@ -641,10 +641,10 @@ func (ic *ImageWriter) commitAttestationsManifest(ctx context.Context, opts *Ima
 	}
 
 	done := progress.OneOff(ctx, "exporting attestation manifest "+mfstDigest.String())
-	if err := content.WriteBlob(ctx, ic.opt.ContentStore, mfstDigest.String(), bytes.NewReader(mfstJSON), mfstDesc, content.WithLabels((labels))); err != nil {
+	if err := content.WriteBlob(ctx, ic.opt.ContentStore, mfstDigest.String(), newBytesReader(mfstJSON), mfstDesc, content.WithLabels((labels))); err != nil {
 		return nil, done(errors.Wrapf(err, "error writing manifest blob %s", mfstDigest))
 	}
-	if err := content.WriteBlob(ctx, ic.opt.ContentStore, configDesc.Digest.String(), bytes.NewReader(config), configDesc); err != nil {
+	if err := content.WriteBlob(ctx, ic.opt.ContentStore, configDesc.Digest.String(), newBytesReader(config), configDesc); err != nil {
 		return nil, done(errors.Wrap(err, "error writing config blob"))
 	}
 	done(nil)
@@ -947,3 +947,41 @@ func getRefMetadata(ref cache.ImmutableRef, limit int) []refMetadata {
 	}
 	return metas
 }
+
+// bytesReader is an alternative version of the bytes.Reader that implements
+// io.WriterTo by sending chunks instead of the entire buffer at once.
+//
+// This is needed to workaround https://github.com/containerd/containerd/issues/11440.
+// When that is fixed, this can be removed.
+type bytesReader struct {
+	p []byte
+}
+
+func newBytesReader(p []byte) *bytesReader {
+	return &bytesReader{p}
+}
+
+func (r *bytesReader) Read(p []byte) (n int, err error) {
+	if len(r.p) == 0 {
+		return 0, io.EOF
+	}
+	n = copy(p, r.p)
+	r.p = r.p[n:]
+	return n, nil
+}
+
+func (r *bytesReader) WriteTo(w io.Writer) (n int64, err error) {
+	const maxBufferSize = 8 << 20
+	for len(r.p) > 0 && err == nil {
+		sz := len(r.p)
+		if sz > maxBufferSize {
+			sz = maxBufferSize
+		}
+
+		var n0 int
+		n0, err = w.Write(r.p[:sz])
+		n += int64(n0)
+		r.p = r.p[n0:]
+	}
+	return n, err
+}