From 98fd78bb5002906eaf66a4b20bf2b480f2cb2b39 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] bytecounter

---
 internal/bytecounter/BUILD               | 18 ++++++++++++++++++
 internal/bytecounter/bytecounter.go      | 20 ++++++++++++++++++++
 internal/bytecounter/bytecounter_test.go | 24 ++++++++++++++++++++++++
 internal/cmdlxcconfig/BUILD              |  1 +
 internal/cmdlxcconfig/cmdlxcconfig.go    | 17 +++++++++++------
 rootfs/rootfs.go                         | 12 ------------
 6 files changed, 74 insertions(+), 18 deletions(-)
 create mode 100644 internal/bytecounter/BUILD
 create mode 100644 internal/bytecounter/bytecounter.go
 create mode 100644 internal/bytecounter/bytecounter_test.go

diff --git a/internal/bytecounter/BUILD b/internal/bytecounter/BUILD
new file mode 100644
index 0000000..4a30583
--- /dev/null
+++ b/internal/bytecounter/BUILD
@@ -0,0 +1,18 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "go_default_library",
+    srcs = ["bytecounter.go"],
+    importpath = "github.com/motiejus/code/undocker/internal/bytecounter",
+    visibility = ["//src/undocker:__subpackages__"],
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = ["bytecounter_test.go"],
+    embed = [":go_default_library"],
+    deps = [
+        "@com_github_stretchr_testify//assert:go_default_library",
+        "@com_github_stretchr_testify//require:go_default_library",
+    ],
+)
diff --git a/internal/bytecounter/bytecounter.go b/internal/bytecounter/bytecounter.go
new file mode 100644
index 0000000..530269c
--- /dev/null
+++ b/internal/bytecounter/bytecounter.go
@@ -0,0 +1,20 @@
+package bytecounter
+
+import "io"
+
+// ByteCounter is an io.Writer that counts bytes written to it
+type ByteCounter struct {
+	N int64
+	w io.Writer
+}
+
+// New returns a new ByteCounter
+func New(w io.Writer) *ByteCounter {
+	return &ByteCounter{w: w}
+}
+
+// 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.w.Write(data)
+}
diff --git a/internal/bytecounter/bytecounter_test.go b/internal/bytecounter/bytecounter_test.go
new file mode 100644
index 0000000..17f5c81
--- /dev/null
+++ b/internal/bytecounter/bytecounter_test.go
@@ -0,0 +1,24 @@
+package bytecounter
+
+import (
+	"bytes"
+	"io"
+	"testing"
+	"testing/iotest"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestByteCounter(t *testing.T) {
+	r := bytes.NewBufferString("0123456789")
+	w := bytes.NewBuffer(nil)
+
+	tw := iotest.TruncateWriter(w, 4)
+	bc := New(tw)
+
+	_, err := io.Copy(bc, r)
+	require.NoError(t, err)
+	assert.Len(t, w.Bytes(), 4)
+	assert.Equal(t, 4, bc.N)
+}
diff --git a/internal/cmdlxcconfig/BUILD b/internal/cmdlxcconfig/BUILD
index c738ba3..94facd9 100644
--- a/internal/cmdlxcconfig/BUILD
+++ b/internal/cmdlxcconfig/BUILD
@@ -6,6 +6,7 @@ go_library(
     importpath = "github.com/motiejus/code/undocker/internal/cmdlxcconfig",
     visibility = ["//src/undocker:__subpackages__"],
     deps = [
+        "//src/undocker/internal/cmd:go_default_library",
         "//src/undocker/lxcconfig:go_default_library",
         "@com_github_jessevdk_go_flags//:go_default_library",
         "@org_uber_go_multierr//:go_default_library",
diff --git a/internal/cmdlxcconfig/cmdlxcconfig.go b/internal/cmdlxcconfig/cmdlxcconfig.go
index 9feb357..cd80670 100644
--- a/internal/cmdlxcconfig/cmdlxcconfig.go
+++ b/internal/cmdlxcconfig/cmdlxcconfig.go
@@ -5,17 +5,22 @@ import (
 	"os"
 
 	goflags "github.com/jessevdk/go-flags"
+	"github.com/motiejus/code/undocker/internal/cmd"
 	"github.com/motiejus/code/undocker/lxcconfig"
 	"go.uber.org/multierr"
 )
 
 // Command is "lxcconfig" command
-type Command struct {
-	PositionalArgs struct {
-		Infile  goflags.Filename `long:"infile" description:"Input tarball"`
-		Outfile string           `long:"outfile" description:"Output path, stdout is '-'"`
-	} `positional-args:"yes" required:"yes"`
-}
+type (
+	Command struct {
+		PositionalArgs struct {
+			cmd.BaseCommand
+
+			Infile  goflags.Filename `long:"infile" description:"Input tarball"`
+			Outfile string           `long:"outfile" description:"Output path, stdout is '-'"`
+		} `positional-args:"yes" required:"yes"`
+	}
+)
 
 // Execute executes lxcconfig Command
 func (c *Command) Execute(args []string) (err error) {
diff --git a/rootfs/rootfs.go b/rootfs/rootfs.go
index 02c666a..8fd2e5d 100644
--- a/rootfs/rootfs.go
+++ b/rootfs/rootfs.go
@@ -222,15 +222,3 @@ 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)
-}