From 303e951585c76403b69008c4eb58027a5a31f9c1 Mon Sep 17 00:00:00 2001
From: Jon Johnson <jonjohnson@google.com>
Date: Mon, 13 Jan 2020 16:21:50 -0800
Subject: [PATCH] Use CompressedSize in legacy/tarball

---
 pkg/legacy/tarball/write.go      | 32 +++++++++++++------
 pkg/legacy/tarball/write_test.go | 55 ++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 9 deletions(-)

diff --git a/pkg/legacy/tarball/write.go b/pkg/legacy/tarball/write.go
index e5555fb47c..48b3c0391b 100644
--- a/pkg/legacy/tarball/write.go
+++ b/pkg/legacy/tarball/write.go
@@ -249,20 +249,34 @@ func MultiWrite(refToImage map[name.Reference]v1.Image, w io.Writer) error {
 				return err
 			}
 			layerFiles[i] = fmt.Sprintf("%s/layer.tar", cur.config.ID)
+
 			u, err := l.Uncompressed()
 			if err != nil {
 				return err
 			}
 			defer u.Close()
-			// Reads the entire uncompressed blob into memory! Can be avoided
-			// for some layer implementations where the uncompressed blob is
-			// stored on disk and the layer can just stat the file.
-			uncompressedBlob, err := ioutil.ReadAll(u)
-			if err != nil {
-				return err
-			}
-			if err := writeTarEntry(tf, layerFiles[i], bytes.NewReader(uncompressedBlob), int64(len(uncompressedBlob))); err != nil {
-				return err
+
+			if partial.HasUncompressedSize(l) {
+				// If the v1.Layer implements UncompressedSize efficiently, use that
+				// for the tar header.
+				size, err := partial.UncompressedSize(l)
+				if err != nil {
+					return err
+				}
+				if err := writeTarEntry(tf, layerFiles[i], u, size); err != nil {
+					return err
+				}
+			} else {
+				// Reads the entire uncompressed blob into memory! Can be avoided
+				// for some layer implementations where the uncompressed blob is
+				// stored on disk and the layer can just stat the file.
+				uncompressedBlob, err := ioutil.ReadAll(u)
+				if err != nil {
+					return err
+				}
+				if err := writeTarEntry(tf, layerFiles[i], bytes.NewReader(uncompressedBlob), int64(len(uncompressedBlob))); err != nil {
+					return err
+				}
 			}
 			j, err := cur.json()
 			if err != nil {
diff --git a/pkg/legacy/tarball/write_test.go b/pkg/legacy/tarball/write_test.go
index 3b9ce495d4..717ee55806 100644
--- a/pkg/legacy/tarball/write_test.go
+++ b/pkg/legacy/tarball/write_test.go
@@ -24,6 +24,7 @@ import (
 	"github.com/google/go-containerregistry/pkg/name"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/mutate"
+	"github.com/google/go-containerregistry/pkg/v1/partial"
 	"github.com/google/go-containerregistry/pkg/v1/random"
 	"github.com/google/go-containerregistry/pkg/v1/tarball"
 	"github.com/google/go-containerregistry/pkg/v1/types"
@@ -401,3 +402,57 @@ func TestMultiWriteMismatchedHistory(t *testing.T) {
 		t.Errorf("Got unexpected error when writing image with mismatched history & layer, got %v, want substring %q", err, want)
 	}
 }
+
+type fastSizeLayer struct {
+	v1.Layer
+	size   int64
+	called bool
+}
+
+func (l *fastSizeLayer) UncompressedSize() (int64, error) {
+	l.called = true
+	return l.size, nil
+}
+
+func TestUncompressedSize(t *testing.T) {
+	// Make a random image
+	img, err := random.Image(256, 8)
+	if err != nil {
+		t.Fatalf("Error creating random image: %v", err)
+	}
+
+	rand, err := random.Layer(1000, types.DockerLayer)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	size, err := partial.UncompressedSize(rand)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	l := &fastSizeLayer{Layer: rand, size: size}
+
+	img, err = mutate.AppendLayers(img, l)
+	if err != nil {
+		t.Fatal(err)
+	}
+	tag, err := name.NewTag("gcr.io/foo/bar:latest", name.StrictValidation)
+	if err != nil {
+		t.Fatalf("Error creating test tag: %v", err)
+	}
+	// Make a tempfile for tarball writes.
+	fp, err := ioutil.TempFile("", "")
+	if err != nil {
+		t.Fatalf("Error creating temp file: %v", err)
+	}
+	t.Log(fp.Name())
+	defer fp.Close()
+	defer os.Remove(fp.Name())
+	if err := Write(tag, img, fp); err != nil {
+		t.Fatalf("Write(): %v", err)
+	}
+	if !l.called {
+		t.Errorf("expected UncompressedSize to be called, but it wasn't")
+	}
+}