From 7f9d150a331d3244122b2c4f1369e96548c2a22c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Motiejus=20Jak=C5=A1tys?= Date: Mon, 24 May 2021 00:11:58 +0300 Subject: [PATCH] wip xz/commander --- internal/cmdrootfs/BUILD | 9 ++++++- internal/cmdrootfs/cmdrootfs.go | 25 ++++++++++++++++-- internal/cmdrootfs/cmdrootfs_test.go | 7 +++++ rootfs/rootfs.go | 38 ++++++++++++++++++++-------- rootfs/rootfs_test.go | 6 +++-- 5 files changed, 69 insertions(+), 16 deletions(-) create mode 100644 internal/cmdrootfs/cmdrootfs_test.go diff --git a/internal/cmdrootfs/BUILD b/internal/cmdrootfs/BUILD index c6f077f..6bf183a 100644 --- a/internal/cmdrootfs/BUILD +++ b/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"], +) diff --git a/internal/cmdrootfs/cmdrootfs.go b/internal/cmdrootfs/cmdrootfs.go index b53add4..5c560f2 100644 --- a/internal/cmdrootfs/cmdrootfs.go +++ b/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 } diff --git a/internal/cmdrootfs/cmdrootfs_test.go b/internal/cmdrootfs/cmdrootfs_test.go new file mode 100644 index 0000000..bcebe32 --- /dev/null +++ b/internal/cmdrootfs/cmdrootfs_test.go @@ -0,0 +1,7 @@ +package cmdrootfs + +import "testing" + +func TestExecute(t *testing.T) { + +} diff --git a/rootfs/rootfs.go b/rootfs/rootfs.go index 316d6e5..75382bb 100644 --- a/rootfs/rootfs.go +++ b/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) +} diff --git a/rootfs/rootfs_test.go b/rootfs/rootfs_test.go index f445289..14d17e1 100644 --- a/rootfs/rootfs_test.go +++ b/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) }) } }