From fd244bfc23084adf31edc44aecee8863fcbe95dd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Motiejus=20Jak=C5=A1tys?= <motiejus@jakstys.lt>
Date: Mon, 24 May 2021 00:11:58 +0300
Subject: [PATCH] read gzip archives

---
 internal/tartest/tartest.go      | 16 ++++++++++++
 internal/tartest/tartest_test.go | 17 +++++++++++-
 rootfs/rootfs.go                 | 44 +++++++++++++++++++++++++++++---
 rootfs/rootfs_test.go            | 12 +++++++++
 4 files changed, 84 insertions(+), 5 deletions(-)

diff --git a/internal/tartest/tartest.go b/internal/tartest/tartest.go
index a4fc4ae..5620aca 100644
--- a/internal/tartest/tartest.go
+++ b/internal/tartest/tartest.go
@@ -3,6 +3,8 @@ package tartest
 import (
 	"archive/tar"
 	"bytes"
+	"compress/gzip"
+	"fmt"
 	"io"
 	"testing"
 
@@ -53,6 +55,20 @@ func (tb Tarball) Buffer() *bytes.Buffer {
 	return &buf
 }
 
+// Gzip returns a gzipped buffer
+func (tb Tarball) Gzip() *bytes.Buffer {
+	var buf bytes.Buffer
+	w := gzip.NewWriter(&buf)
+	_, err := io.Copy(w, tb.Buffer())
+	if err != nil {
+		panic(fmt.Errorf("Gzip(): %w", err))
+	}
+	if err := w.Close(); err != nil {
+		panic(fmt.Errorf("gzip.Close(): %w", err))
+	}
+	return &buf
+}
+
 // Tar tars the Dir
 func (d Dir) Tar(tw *tar.Writer) error {
 	hdr := &tar.Header{
diff --git a/internal/tartest/tartest_test.go b/internal/tartest/tartest_test.go
index 5596ef8..2f58e4c 100644
--- a/internal/tartest/tartest_test.go
+++ b/internal/tartest/tartest_test.go
@@ -2,9 +2,12 @@ package tartest
 
 import (
 	"bytes"
+	"compress/gzip"
+	"io"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestTarball(t *testing.T) {
@@ -23,6 +26,18 @@ func TestTarball(t *testing.T) {
 		Dir{Name: "bin"},
 		Hardlink{Name: "entrypoint2"},
 	}
-
 	assert.Equal(t, want, got)
 }
+
+func TestGzip(t *testing.T) {
+	tb := Tarball{File{Name: "entrypoint.sh", Contents: bytes.NewBufferString("hello")}}
+	tbuf := tb.Buffer()
+
+	tgz, err := gzip.NewReader(tb.Gzip())
+	require.NoError(t, err)
+	var uncompressed bytes.Buffer
+	_, err = io.Copy(&uncompressed, tgz)
+	require.NoError(t, err)
+	assert.Equal(t, tbuf.Bytes(), uncompressed.Bytes())
+
+}
diff --git a/rootfs/rootfs.go b/rootfs/rootfs.go
index 14f0550..38ab8f3 100644
--- a/rootfs/rootfs.go
+++ b/rootfs/rootfs.go
@@ -2,6 +2,8 @@ package rootfs
 
 import (
 	"archive/tar"
+	"bytes"
+	"compress/gzip"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -9,7 +11,7 @@ import (
 	"path/filepath"
 	"strings"
 
-     "github.com/motiejus/code/undocker/internal/bytecounter"
+	"github.com/motiejus/code/undocker/internal/bytecounter"
 	"go.uber.org/multierr"
 )
 
@@ -43,6 +45,7 @@ func (r *RootFS) WriteTo(w io.Writer) (n int64, err error) {
 	wr := bytecounter.New(w)
 	tr := tar.NewReader(r.rd)
 	tw := tar.NewWriter(wr)
+	var closer func() error
 	defer func() {
 		err = multierr.Append(err, tw.Close())
 		n = wr.N
@@ -110,7 +113,7 @@ func (r *RootFS) WriteTo(w io.Writer) (n int64, err error) {
 		if _, err := r.rd.Seek(no.offset, io.SeekStart); err != nil {
 			return n, err
 		}
-		tr = tar.NewReader(r.rd)
+		tr, closer = readTar(r.rd)
 		for {
 			hdr, err := tr.Next()
 			if err == io.EOF {
@@ -138,9 +141,11 @@ func (r *RootFS) WriteTo(w io.Writer) (n int64, err error) {
 					continue
 				}
 			}
-
 			file2layer[hdr.Name] = i
 		}
+		if err := closer(); err != nil {
+			return n, err
+		}
 	}
 
 	// construct directories to whiteout, for each layer.
@@ -151,7 +156,7 @@ func (r *RootFS) WriteTo(w io.Writer) (n int64, err error) {
 		if _, err := r.rd.Seek(no.offset, io.SeekStart); err != nil {
 			return n, err
 		}
-		tr = tar.NewReader(r.rd)
+		tr, closer = readTar(r.rd)
 		for {
 			hdr, err := tr.Next()
 			if err == io.EOF {
@@ -173,6 +178,9 @@ func (r *RootFS) WriteTo(w io.Writer) (n int64, err error) {
 				return n, err
 			}
 		}
+		if err := closer(); err != nil {
+			return n, err
+		}
 	}
 	return n, nil
 }
@@ -223,3 +231,31 @@ func whiteoutDirs(whreaddir map[string]int, nlayers int) []*tree {
 	}
 	return ret
 }
+
+// readTar creates a tar reader from a targzip or tar
+func readTar(r io.Reader) (*tar.Reader, func() error) {
+	var buf bytes.Buffer
+	w := &discarder{w: &buf}
+	r2 := io.TeeReader(r, w)
+	gz, err := gzip.NewReader(r2)
+	if err == nil {
+		w.discard = true
+		buf.Reset()
+		return tar.NewReader(gz), gz.Close
+	}
+	return tar.NewReader(io.MultiReader(&buf, r)), func() error { return nil }
+}
+
+// discarder is a pass-through writer until asked to 'discard' its writes.
+// useful for proxying writes from a TeeReader until a certain point.
+type discarder struct {
+	w       io.Writer
+	discard bool
+}
+
+func (d *discarder) Write(p []byte) (int, error) {
+	if d.discard {
+		return len(p), nil
+	}
+	return d.w.Write(p)
+}
diff --git a/rootfs/rootfs_test.go b/rootfs/rootfs_test.go
index 14d17e1..3612dea 100644
--- a/rootfs/rootfs_test.go
+++ b/rootfs/rootfs_test.go
@@ -167,6 +167,18 @@ func TestRootFS(t *testing.T) {
 				file{Name: "a/fileb"},
 			},
 		},
+		{
+			name: "archived layer",
+			image: tarball{
+				file{Name: "layer1/layer.tar", Contents: layer1.Gzip()},
+				file{Name: "layer0/layer.tar", Contents: layer0.Gzip()},
+				manifest{"layer0/layer.tar", "layer1/layer.tar"},
+			},
+			want: []extractable{
+				dir{Name: "/", UID: 0},
+				file{Name: "/file", UID: 1, Contents: bytes.NewBufferString("from 1")},
+			},
+		},
 	}
 
 	for _, tt := range tests {