Use a limit writer to capture actual tar size

coder · Emyrk · Jun 5, 2023 · Jun 1, 2023 · Jun 1, 2023 · Jun 1, 2023
commit 0cd070ca085e61a2885ed8902b051f275bb06b7c
diff --git a/coderd/util/xio/limitwriter.go b/coderd/util/xio/limitwriter.go
@@ -0,0 +1,43 @@
+package xio
+
+import (
+	"io"
+
+	"golang.org/x/xerrors"
+)
+
+var ErrLimitReached = xerrors.Errorf("i/o limit reached")
+
+// LimitWriter will only write bytes to the underlying writer until the limit is reached.
+type LimitWriter struct {
+	Limit int64
+	N     int64
+	W     io.Writer
+}
+
+func NewLimitWriter(w io.Writer, n int64) *LimitWriter {
+	// If anyone tries this, just make a 0 writer.
+	if n < 0 {
+		n = 0
+	}
+	return &LimitWriter{
+		Limit: n,
+		N:     0,
+		W:     w,
+	}
+}
+
+func (l *LimitWriter) Write(p []byte) (int, error) {
+	if l.N >= l.Limit {
+		return 0, ErrLimitReached
+	}
+
+	// Write 0 bytes if the limit is to be exceeded.
+	if int64(len(p)) > l.Limit-l.N {
+		return 0, ErrLimitReached
+	}
+
+	n, err := l.W.Write(p)
+	l.N += int64(n)
+	return n, err
+}
diff --git a/coderd/util/xio/limitwriter_test.go b/coderd/util/xio/limitwriter_test.go
@@ -0,0 +1,136 @@
+package xio_test
+
+import (
+	"bytes"
+	cryptorand "crypto/rand"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/coderd/util/xio"
+)
+
+func TestLimitWriter(t *testing.T) {
+	type writeCase struct {
+		N    int
+		ExpN int
+		Err  bool
+	}
+
+	// testCases will do multiple writes to the same limit writer and check the output.
+	testCases := []struct {
+		Name   string
+		L      int64
+		Writes []writeCase
+		N      int
+		ExpN   int
+	}{
+		{
+			Name: "Empty",
+			L:    1000,
+			Writes: []writeCase{
+				// A few empty writes
+				{N: 0, ExpN: 0}, {N: 0, ExpN: 0}, {N: 0, ExpN: 0},
+			},
+		},
+		{
+			Name: "NotFull",
+			L:    1000,
+			Writes: []writeCase{
+				{N: 250, ExpN: 250},
+				{N: 250, ExpN: 250},
+				{N: 250, ExpN: 250},
+			},
+		},
+		{
+			Name: "Short",
+			L:    1000,
+			Writes: []writeCase{
+				{N: 250, ExpN: 250},
+				{N: 250, ExpN: 250},
+				{N: 250, ExpN: 250},
+				{N: 250, ExpN: 250},
+				{N: 250, ExpN: 0, Err: true},
+			},
+		},
+		{
+			Name: "Exact",
+			L:    1000,
+			Writes: []writeCase{
+				{
+					N:    1000,
+					ExpN: 1000,
+				},
+				{
+					N:   1000,
+					Err: true,
+				},
+			},
+		},
+		{
+			Name: "Over",
+			L:    1000,
+			Writes: []writeCase{
+				{
+					N:    5000,
+					ExpN: 0,
+					Err:  true,
+				},
+				{
+					N:   5000,
+					Err: true,
+				},
+				{
+					N:   5000,
+					Err: true,
+				},
+			},
+		},
+		{
+			Name: "Strange",
+			L:    -1,
+			Writes: []writeCase{
+				{
+					N:    5,
+					ExpN: 0,
+					Err:  true,
+				},
+				{
+					N:    0,
+					ExpN: 0,
+					Err:  true,
+				},
+			},
+		},
+	}
+
+	for _, c := range testCases {
+		t.Run(c.Name, func(t *testing.T) {
+			buf := bytes.NewBuffer([]byte{})
+			allBuff := bytes.NewBuffer([]byte{})
+			w := xio.NewLimitWriter(buf, c.L)
+
+			for _, wc := range c.Writes {
+				data := make([]byte, wc.N)
+
+				n, err := cryptorand.Read(data)
+				require.NoError(t, err, "crand read")
+				require.Equal(t, wc.N, n, "correct bytes read")
+				max := data[:wc.ExpN]
+				n, err = w.Write(data)
+				if wc.Err {
+					require.Error(t, err, "exp error")
+				} else {
+					require.NoError(t, err, "write")
+				}
+
+				// Need to use this to compare across multiple writes.
+				// Each write appends to the expected output.
+				allBuff.Write(max)
+
+				require.Equal(t, wc.ExpN, n, "correct bytes written")
+				require.Equal(t, allBuff.Bytes(), buf.Bytes(), "expected data")
+			}
+		})
+	}
+}
diff --git a/provisionersdk/archive.go b/provisionersdk/archive.go
@@ -8,6 +8,8 @@ import (
 	"strings"
 
 	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/coderd/util/xio"
 )
 
 const (
@@ -32,8 +34,9 @@ func dirHasExt(dir string, ext string) (bool, error) {
 
 // Tar archives a Terraform directory.
 func Tar(w io.Writer, directory string, limit int64) error {
+	// The total bytes written must be under the limit.
+	w = xio.NewLimitWriter(w, limit)
 	tarWriter := tar.NewWriter(w)
-	totalSize := int64(0)
 
 	const tfExt = ".tf"
 	hasTf, err := dirHasExt(directory, tfExt)
@@ -54,7 +57,6 @@ func Tar(w io.Writer, directory string, limit int64) error {
 		)
 	}
 
-	fileTooBigError := xerrors.Errorf("Archive too big. Must be <= %d bytes", limit)
 	err = filepath.Walk(directory, func(file string, fileInfo os.FileInfo, err error) error {
 		if err != nil {
 			return err
@@ -96,23 +98,20 @@ func Tar(w io.Writer, directory string, limit int64) error {
 		if !fileInfo.Mode().IsRegular() {
 			return nil
 		}
-		// Before we even open the file, check if it is going to exceed our limit.
-		if fileInfo.Size()+totalSize > limit {
-			return fileTooBigError
-		}
+
 		data, err := os.Open(file)
 		if err != nil {
 			return err
 		}
 		defer data.Close()
-		wrote, err := io.Copy(tarWriter, data)
+		_, err = io.Copy(tarWriter, data)
 		if err != nil {
+			if xerrors.Is(err, xio.ErrLimitReached) {
+				return xerrors.Errorf("Archive too big. Must be <= %d bytes", limit)
+			}
 			return err
 		}
-		totalSize += wrote
-		if limit != 0 && totalSize > limit {
-			return fileTooBigError
-		}
+
 		return data.Close()
 	})
 	if err != nil {