Error when too few bytes are read

In the case where we have a scenario that results in fewer than expected
bytes being read from the tar compression or fewer than expected bytes
written from the file writer consumer, propagate an error to the main
program and exit with non-success status.

Observed behavior shows that within the tar-extraction code path in
certain circumstances a clean EOF is received by the tarReader and the
archive is partially extracted. This generally presents as files missing
in the destination directory.

New code paths expect the correct number of bytes to be read or it will
propagate a non-zero exit.

.goreleaser.yaml

@@ -1,4 +1,5 @@
 project_name: pget
+version: 2
 before:
   hooks:
     - go mod tidy

pkg/consumer/consumer.go

@@ -3,5 +3,5 @@ package consumer
 import "io"
 
 type Consumer interface {
-	Consume(reader io.Reader, destPath string) error
+	Consume(reader io.Reader, destPath string, expectedBytes int64) error
 }

pkg/consumer/consumer_test.go

@@ -0,0 +1,16 @@
+package consumer_test
+
+import (
+	"math/rand"
+)
+
+// generateTestContent generates a byte slice of a random size > 1KiB
+func generateTestContent(size int64) []byte {
+	content := make([]byte, size)
+	// Generate random bytes and write them to the content slice
+	for i := range content {
+		content[i] = byte(rand.Intn(256))
+	}
+	return content
+
+}

pkg/consumer/null.go

@@ -1,15 +1,19 @@
 package consumer
 
 import (
+	"fmt"
 	"io"
 )
 
 type NullWriter struct{}
 
 var _ Consumer = &NullWriter{}
 
-func (NullWriter) Consume(reader io.Reader, destPath string) error {
+func (NullWriter) Consume(reader io.Reader, destPath string, expectedBytes int64) error {
 	// io.Discard is explicitly designed to always succeed, ignore errors.
-	_, _ = io.Copy(io.Discard, reader)
+	bytesRead, _ := io.Copy(io.Discard, reader)
+	if bytesRead != expectedBytes {
+		return fmt.Errorf("expected %d bytes, read %d", expectedBytes, bytesRead)
+	}
 	return nil
 }

pkg/consumer/null_test.go

@@ -0,0 +1,23 @@
+package consumer_test
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/replicate/pget/pkg/consumer"
+)
+
+func TestNullWriter_Consume(t *testing.T) {
+	buf := generateTestContent(1024)
+	reader := bytes.NewReader(buf)
+
+	nullConsumer := consumer.NullWriter{}
+	err := nullConsumer.Consume(reader, "", 1024)
+	assert.NoError(t, err)
+
+	_, _ = reader.Seek(0, 0)
+	err = nullConsumer.Consume(reader, "", 100)
+	assert.Error(t, err)
+}

pkg/consumer/tar_extractor.go

@@ -14,10 +14,27 @@ type TarExtractor struct {
 
 var _ Consumer = &TarExtractor{}
 
-func (f *TarExtractor) Consume(reader io.Reader, destPath string) error {
-	err := extract.TarFile(bufio.NewReader(reader), destPath, f.Overwrite)
+var _ io.Reader = &byteTrackingReader{}
+
+type byteTrackingReader struct {
+	bytesRead int64
+	r         io.Reader
+}
+
+func (b *byteTrackingReader) Read(p []byte) (n int, err error) {
+	n, err = b.r.Read(p)
+	b.bytesRead += int64(n)
+	return
+}
+
+func (f *TarExtractor) Consume(reader io.Reader, destPath string, expectedBytes int64) error {
+	btReader := &byteTrackingReader{r: reader}
+	err := extract.TarFile(bufio.NewReader(btReader), destPath, f.Overwrite)
 	if err != nil {
 		return fmt.Errorf("error extracting file: %w", err)
 	}
+	if btReader.bytesRead != expectedBytes {
+		return fmt.Errorf("expected %d bytes, read %d from archive", expectedBytes, btReader.bytesRead)
+	}
 	return nil
 }

pkg/consumer/tar_extractor_test.go

@@ -0,0 +1,167 @@
+package consumer_test
+
+import (
+	"archive/tar"
+	"bytes"
+	"os"
+	"path"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/replicate/pget/pkg/consumer"
+)
+
+const (
+	file1Content     = "This is the content of file1."
+	file2Content     = "This is the content of file2."
+	file1Path        = "file1.txt"
+	file2Path        = "file2.txt"
+	fileSymLinkPath  = "link_to_file1.txt"
+	fileHardLinkPath = "subdir/hard_link_to_file2.txt"
+)
+
+func createTarFileBytesBuffer() ([]byte, error) {
+	// Create an in-memory representation of a tar file dynamically. This will be used to test the TarExtractor
+
+	var buf bytes.Buffer
+	tw := tar.NewWriter(&buf)
+
+	// Create first file
+	content1 := []byte(file1Content)
+	hdr := &tar.Header{
+		Name:    file1Path,
+		Mode:    0600,
+		Size:    int64(len(content1)),
+		ModTime: time.Now(),
+	}
+	if err := tw.WriteHeader(hdr); err != nil {
+		return nil, err
+	}
+	if _, err := tw.Write(content1); err != nil {
+		return nil, err
+	}
+
+	// Create second file
+	content2 := []byte(file2Content)
+	hdr = &tar.Header{
+		Name:    file2Path,
+		Mode:    0600,
+		Size:    int64(len(content2)),
+		ModTime: time.Now(),
+	}
+	if err := tw.WriteHeader(hdr); err != nil {
+		return nil, err
+	}
+	if _, err := tw.Write(content2); err != nil {
+		return nil, err
+	}
+
+	// Create a symlink to file1
+	hdr = &tar.Header{
+		Name:     fileSymLinkPath,
+		Mode:     0777,
+		Size:     0,
+		Linkname: file1Path,
+		Typeflag: tar.TypeSymlink,
+		ModTime:  time.Now(),
+	}
+	if err := tw.WriteHeader(hdr); err != nil {
+		return nil, err
+	}
+
+	// Create a subdirectory or path for the hardlink
+	hdr = &tar.Header{
+		Name:     "subdir/",
+		Mode:     0755,
+		Typeflag: tar.TypeDir,
+		ModTime:  time.Now(),
+	}
+	if err := tw.WriteHeader(hdr); err != nil {
+		return nil, err
+	}
+
+	// Create a hardlink to file2 in the subdirectory
+	hdr = &tar.Header{
+		Name:     fileHardLinkPath,
+		Mode:     0600,
+		Size:     0,
+		Linkname: file2Path,
+		Typeflag: tar.TypeLink,
+		ModTime:  time.Now(),
+	}
+	if err := tw.WriteHeader(hdr); err != nil {
+		return nil, err
+	}
+
+	// Close the tar writer to flush the data
+	if err := tw.Close(); err != nil {
+		return nil, err
+	}
+
+	return buf.Bytes(), nil
+}
+
+func TestTarExtractor_Consume(t *testing.T) {
+	tarFileBytes, err := createTarFileBytesBuffer()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Create a reader from the tar file bytes
+	reader := bytes.NewReader(tarFileBytes)
+
+	// Create a temporary directory to extract the tar file
+	tmpDir, err := os.MkdirTemp("", "tarExtractorTest-")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(tmpDir)
+
+	tarConsumer := consumer.TarExtractor{}
+	targetDir := path.Join(tmpDir, "extract")
+	err = tarConsumer.Consume(reader, targetDir, int64(len(tarFileBytes)))
+	assert.NoError(t, err)
+
+	// Check if the extraction was successful
+	checkTarExtraction(t, targetDir)
+
+	// Test with incorrect expectedBytes
+	_, _ = reader.Seek(0, 0)
+	targetDir = path.Join(tmpDir, "extract-fail")
+	err = tarConsumer.Consume(reader, targetDir, int64(len(tarFileBytes)-1))
+	assert.Error(t, err)
+}
+
+func checkTarExtraction(t *testing.T, targetDir string) {
+	// Verify that file1.txt is correctly extracted
+	fqFile1Path := path.Join(targetDir, file1Path)
+	content, err := os.ReadFile(fqFile1Path)
+	assert.NoError(t, err)
+	assert.Equal(t, file1Content, string(content))
+
+	// Verify that file2.txt is correctly extracted
+	fqFile2Path := path.Join(targetDir, file2Path)
+	content, err = os.ReadFile(fqFile2Path)
+	assert.NoError(t, err)
+	assert.Equal(t, file2Content, string(content))
+
+	// Verify that link_to_file1.txt is a symlink pointing to file1.txt
+	linkToFile1Path := path.Join(targetDir, fileSymLinkPath)
+	linkTarget, err := os.Readlink(linkToFile1Path)
+	assert.NoError(t, err)
+	assert.Equal(t, file1Path, linkTarget)
+	assert.Equal(t, os.ModeSymlink, os.ModeSymlink&os.ModeType)
+
+	// Verify that subdir/hard_link_to_file2.txt is a hard link to file2.txt
+	hardLinkToFile2Path := path.Join(targetDir, fileHardLinkPath)
+	hardLinkStat, err := os.Stat(hardLinkToFile2Path)
+	assert.NoError(t, err)
+	file2Stat, err := os.Stat(fqFile2Path)
+	assert.NoError(t, err)
+
+	if !os.SameFile(hardLinkStat, file2Stat) {
+		t.Errorf("hard link does not match file2.txt")
+	}
+}

pkg/consumer/write_file.go

@@ -13,7 +13,7 @@ type FileWriter struct {
 
 var _ Consumer = &FileWriter{}
 
-func (f *FileWriter) Consume(reader io.Reader, destPath string) error {
+func (f *FileWriter) Consume(reader io.Reader, destPath string, expectedBytes int64) error {
 	openFlags := os.O_WRONLY | os.O_CREATE
 	targetDir := filepath.Dir(destPath)
 	if err := os.MkdirAll(targetDir, 0755); err != nil {
@@ -28,9 +28,13 @@ func (f *FileWriter) Consume(reader io.Reader, destPath string) error {
 	}
 	defer out.Close()
 
-	_, err = io.Copy(out, reader)
+	written, err := io.Copy(out, reader)
 	if err != nil {
 		return fmt.Errorf("error writing file: %w", err)
 	}
+
+	if written != expectedBytes {
+		return fmt.Errorf("expected %d bytes, wrote %d", expectedBytes, written)
+	}
 	return nil
 }

pkg/consumer/write_file_test.go

@@ -0,0 +1,51 @@
+package consumer_test
+
+import (
+	"bytes"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/replicate/pget/pkg/consumer"
+)
+
+func TestFileWriter_Consume(t *testing.T) {
+
+	buf := generateTestContent(1024)
+	reader := bytes.NewReader(buf)
+
+	writeFileConsumer := consumer.FileWriter{}
+	tmpFile, _ := os.CreateTemp("", "fileWriterTest-")
+
+	defer tmpFile.Close()
+	defer os.Remove(tmpFile.Name())
+
+	err := writeFileConsumer.Consume(reader, tmpFile.Name(), 1024)
+	assert.NoError(t, err)
+
+	// Check the file content is correct
+	fileContent, _ := os.ReadFile(tmpFile.Name())
+	assert.Equal(t, buf, fileContent)
+
+	_, _ = reader.Seek(0, 0)
+	err = writeFileConsumer.Consume(reader, "", 100)
+	assert.Error(t, err)
+
+	// test overwrite
+	// overwrite the file
+	f, err := os.OpenFile(tmpFile.Name(), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0755)
+	assert.NoError(t, err)
+	_, _ = f.Write([]byte("different content"))
+	f.Close()
+
+	// consume the reader
+	_, _ = reader.Seek(0, 0)
+	writeFileConsumer.Overwrite = true
+	err = writeFileConsumer.Consume(reader, tmpFile.Name(), 1024)
+	assert.NoError(t, err)
+
+	// check the file content is correct
+	fileContent, _ = os.ReadFile(tmpFile.Name())
+	assert.Equal(t, buf, fileContent)
+}

pkg/pget.go

@@ -49,7 +49,7 @@ func (g *Getter) DownloadFile(ctx context.Context, url string, dest string) (int
 	// downloadElapsed := time.Since(downloadStartTime)
 	// writeStartTime := time.Now()
 
-	err = g.Consumer.Consume(buffer, dest)
+	err = g.Consumer.Consume(buffer, dest, fileSize)
 	if err != nil {
 		return fileSize, 0, fmt.Errorf("error writing file: %w", err)
 	}