wip xz/commander

internal/cmdrootfs/BUILD

@@ -1,4 +1,4 @@
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "go_default_library",
@@ -8,6 +8,13 @@ go_library(
     deps = [
         "//src/undocker/rootfs:go_default_library",
         "@com_github_jessevdk_go_flags//:go_default_library",
+        "@com_github_ulikunitz_xz//:go_default_library",
         "@org_uber_go_multierr//:go_default_library",
     ],
 )
+
+go_test(
+    name = "go_default_test",
+    srcs = ["cmdrootfs_test.go"],
+    embed = [":go_default_library"],
+)

internal/cmdrootfs/cmdrootfs.go

@@ -2,10 +2,12 @@ package cmdrootfs
 
 import (
 	"errors"
+	"io"
 	"os"
 
 	goflags "github.com/jessevdk/go-flags"
 	"github.com/motiejus/code/undocker/rootfs"
+	"github.com/ulikunitz/xz"
 	"go.uber.org/multierr"
 )
 
@@ -15,6 +17,10 @@ type Command struct {
 		Infile  goflags.Filename `long:"infile" description:"Input tarball"`
 		Outfile string           `long:"outfile" description:"Output path, stdout is '-'"`
 	} `positional-args:"yes" required:"yes"`
+
+	Xz bool `short:"J" long:"xz" description:"create XZ archive"`
+
+	rootfsNew func(io.ReadSeeker) io.WriterTo
 }
 
 // Execute executes rootfs Command
@@ -22,6 +28,11 @@ func (c *Command) Execute(args []string) (err error) {
 	if len(args) != 0 {
 		return errors.New("too many args")
 	}
+	if c.rootfsNew == nil {
+		c.rootfsNew = func(r io.ReadSeeker) io.WriterTo {
+			return rootfs.New(r)
+		}
+	}
 
 	rd, err := os.Open(string(c.PositionalArgs.Infile))
 	if err != nil {
@@ -29,7 +40,7 @@ func (c *Command) Execute(args []string) (err error) {
 	}
 	defer func() { err = multierr.Append(err, rd.Close()) }()
 
-	var out *os.File
+	var out io.WriteCloser
 	outf := string(c.PositionalArgs.Outfile)
 	if outf == "-" {
 		out = os.Stdout
@@ -41,5 +52,15 @@ func (c *Command) Execute(args []string) (err error) {
 	}
 	defer func() { err = multierr.Append(err, out.Close()) }()
 
-	return rootfs.New(rd).WriteTo(out)
+	if c.Xz {
+		if out, err = xz.NewWriter(out); err != nil {
+			return err
+		}
+		defer func() { err = multierr.Append(err, out.Close()) }()
+	}
+
+	if _, err := c.rootfsNew(rd).WriteTo(out); err != nil {
+		return err
+	}
+	return nil
 }

internal/cmdrootfs/cmdrootfs_test.go

@@ -0,0 +1,7 @@
+package cmdrootfs
+
+import "testing"
+
+func TestExecute(t *testing.T) {
+
+}

rootfs/rootfs.go

@@ -38,10 +38,14 @@ func New(rd io.ReadSeeker) *RootFS {
 }
 
 // WriteTo writes a docker image to an open tarball.
-func (r *RootFS) WriteTo(wr io.Writer) (err error) {
+func (r *RootFS) WriteTo(w io.Writer) (n int64, err error) {
+	wr := &byteCounter{rw: w}
 	tr := tar.NewReader(r.rd)
 	tw := tar.NewWriter(wr)
-	defer func() { err = multierr.Append(err, tw.Close()) }()
+	defer func() {
+		err = multierr.Append(err, tw.Close())
+		n = wr.n
+	}()
 
 	// layerOffsets maps a layer name (a9b123c0daa/layer.tar) to it's offset
 	layerOffsets := map[string]int64{}
@@ -62,19 +66,19 @@ func (r *RootFS) WriteTo(wr io.Writer) (err error) {
 		case filepath.Clean(hdr.Name) == _manifestJSON:
 			dec := json.NewDecoder(tr)
 			if err := dec.Decode(&manifest); err != nil {
-				return fmt.Errorf("decode %s: %w", _manifestJSON, err)
+				return n, fmt.Errorf("decode %s: %w", _manifestJSON, err)
 			}
 		case strings.HasSuffix(hdr.Name, _layerSuffix):
 			here, err := r.rd.Seek(0, io.SeekCurrent)
 			if err != nil {
-				return err
+				return n, err
 			}
 			layerOffsets[hdr.Name] = here
 		}
 	}
 
 	if len(manifest) == 0 || len(layerOffsets) != len(manifest[0].Layers) {
-		return errBadManifest
+		return n, errBadManifest
 	}
 
 	// enumerate layers the way they would be laid down in the image
@@ -96,7 +100,7 @@ func (r *RootFS) WriteTo(wr io.Writer) (err error) {
 	// iterate over all files, construct `file2layer`, `whreaddir`, `wh`
 	for i, offset := range layers {
 		if _, err := r.rd.Seek(offset, io.SeekStart); err != nil {
-			return err
+			return n, err
 		}
 		tr = tar.NewReader(r.rd)
 		for {
@@ -105,7 +109,7 @@ func (r *RootFS) WriteTo(wr io.Writer) (err error) {
 				break
 			}
 			if err != nil {
-				return err
+				return n, err
 			}
 			if hdr.Typeflag == tar.TypeDir {
 				continue
@@ -137,7 +141,7 @@ func (r *RootFS) WriteTo(wr io.Writer) (err error) {
 	// iterate through all layers, all files, and write files.
 	for i, offset := range layers {
 		if _, err := r.rd.Seek(offset, io.SeekStart); err != nil {
-			return err
+			return n, err
 		}
 		tr = tar.NewReader(r.rd)
 		for {
@@ -146,7 +150,7 @@ func (r *RootFS) WriteTo(wr io.Writer) (err error) {
 				break
 			}
 			if err != nil {
-				return err
+				return n, err
 			}
 			if layer, ok := wh[hdr.Name]; ok && layer >= i {
 				continue
@@ -158,11 +162,11 @@ func (r *RootFS) WriteTo(wr io.Writer) (err error) {
 				continue
 			}
 			if err := writeFile(tr, tw, hdr); err != nil {
-				return err
+				return n, err
 			}
 		}
 	}
-	return nil
+	return n, nil
 }
 
 func writeFile(tr *tar.Reader, tw *tar.Writer, hdr *tar.Header) error {
@@ -211,3 +215,15 @@ func whiteoutDirs(whreaddir map[string]int, nlayers int) []*tree {
 	}
 	return ret
 }
+
+// byteCounter is an io.Writer that counts bytes written to it
+type byteCounter struct {
+	rw io.Writer
+	n  int64
+}
+
+// Write writes to the underlying io.Writer and counts total written bytes
+func (b *byteCounter) Write(data []byte) (n int, err error) {
+	defer func() { b.n += int64(n) }()
+	return b.rw.Write(data)
+}

rootfs/rootfs_test.go

@@ -174,14 +174,16 @@ func TestRootFS(t *testing.T) {
 			in := bytes.NewReader(tt.image.Buffer().Bytes())
 			out := bytes.Buffer{}
 
-			err := New(in).WriteTo(&out)
+			n, err := New(in).WriteTo(&out)
 			if tt.wantErr != "" {
 				assert.EqualError(t, err, tt.wantErr)
 				return
 			}
+			outb := out.Bytes()
 			require.NoError(t, err)
-			got := tartest.Extract(t, &out)
+			got := tartest.Extract(t, bytes.NewReader(outb))
 			assert.Equal(t, tt.want, got)
+			assert.Equal(t, int64(len(outb)), n)
 		})
 	}
 }