Error when too few bytes are read

.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,20 @@
+package consumer_test
+
+import (
+	"math/rand"
+)
+
+const (
+	kB int64 = 1024
+)
+
+// 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,22 @@
+package consumer_test
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/replicate/pget/pkg/consumer"
+)
+
+func TestNullWriter_Consume(t *testing.T) {
+	r := require.New(t)
+	buf := generateTestContent(kB)
+	reader := bytes.NewReader(buf)
+
+	nullConsumer := consumer.NullWriter{}
+	r.NoError(nullConsumer.Consume(reader, "", kB))
+
+	_, _ = reader.Seek(0, 0)
+	r.Error(nullConsumer.Consume(reader, "", kB-100))
+}

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,166 @@
+package consumer_test
+
+import (
+	"archive/tar"
+	"bytes"
+	"os"
+	"path"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"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) {
+	r := require.New(t)
+
+	tarFileBytes, err := createTarFileBytesBuffer()
+	r.NoError(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-")
+	r.NoError(err)
+
+	t.Cleanup(func() { os.RemoveAll(tmpDir) })
+
+	tarConsumer := consumer.TarExtractor{}
+	targetDir := path.Join(tmpDir, "extract")
+	r.NoError(tarConsumer.Consume(reader, targetDir, int64(len(tarFileBytes))))
+
+	// Check if the extraction was successful
+	checkTarExtraction(t, targetDir)
+
+	// Test with incorrect expectedBytes
+	_, _ = reader.Seek(0, 0)
+	targetDir = path.Join(tmpDir, "extract-fail")
+	r.Error(tarConsumer.Consume(reader, targetDir, int64(len(tarFileBytes)-1)))
+}
+
+func checkTarExtraction(t *testing.T, targetDir string) {
+	r := require.New(t)
+
+	// Verify that file1.txt is correctly extracted
+	fqFile1Path := path.Join(targetDir, file1Path)
+	content, err := os.ReadFile(fqFile1Path)
+	r.NoError(err)
+	r.Equal(file1Content, string(content))
+
+	// Verify that file2.txt is correctly extracted
+	fqFile2Path := path.Join(targetDir, file2Path)
+	content, err = os.ReadFile(fqFile2Path)
+	r.NoError(err)
+	r.Equal(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)
+	r.NoError(err)
+	r.Equal(file1Path, linkTarget)
+	r.Equal(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)
+	r.NoError(err)
+	file2Stat, err := os.Stat(fqFile2Path)
+	r.NoError(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/require"
+
+	"github.com/replicate/pget/pkg/consumer"
+)
+
+func TestFileWriter_Consume(t *testing.T) {
+	r := require.New(t)
+
+	buf := generateTestContent(kB)
+	reader := bytes.NewReader(buf)
+
+	writeFileConsumer := consumer.FileWriter{}
+	tmpFile, _ := os.CreateTemp("", "fileWriterTest-")
+
+	t.Cleanup(func() {
+		tmpFile.Close()
+		os.Remove(tmpFile.Name())
+	})
+
+	r.NoError(writeFileConsumer.Consume(reader, tmpFile.Name(), kB))
+
+	// Check the file content is correct
+	fileContent, _ := os.ReadFile(tmpFile.Name())
+	r.Equal(buf, fileContent)
+
+	_, _ = reader.Seek(0, 0)
+	r.Error(writeFileConsumer.Consume(reader, "", kB-100))
+
+	// test overwrite
+	// overwrite the file
+	f, err := os.OpenFile(tmpFile.Name(), os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0755)
+	r.NoError(err)
+	_, _ = f.Write([]byte("different content"))
+	f.Close()
+
+	// consume the reader
+	_, _ = reader.Seek(0, 0)
+	writeFileConsumer.Overwrite = true
+	r.NoError(writeFileConsumer.Consume(reader, tmpFile.Name(), kB))
+
+	// check the file content is correct
+	fileContent, _ = os.ReadFile(tmpFile.Name())
+	r.Equal(buf, fileContent)
+}

pkg/download/buffer_slow_test.go

@@ -17,5 +17,5 @@ func BenchmarkDownload10GH2(b *testing.B) {
 }
 
 func BenchmarkDownloadDollyTensors(b *testing.B) {
-	benchmarkDownloadURL(defaultOpts, "https://storage.googleapis.com/replicate-weights/dolly-v2-12b-fp16.tensors", b)
+	benchmarkDownloadURL(defaultOpts, "https://weights.replicate.delivery/default/dolly-v2-12b-fp16.tensors", b)
 }

pkg/download/buffer_test.go

@@ -31,14 +31,14 @@ func benchmarkDownloadURL(opts download.Options, url string, b *testing.B) {
 }
 
 func BenchmarkDownloadBertH1(b *testing.B) {
-	benchmarkDownloadURL(defaultOpts, "https://storage.googleapis.com/replicate-weights/bert-base-uncased-hf-cache.tar", b)
+	benchmarkDownloadURL(defaultOpts, "https://weights.replicate.delivery/default/bert-base-uncased-hf-cache.tar", b)
 }
 func BenchmarkDownloadBertH2(b *testing.B) {
-	benchmarkDownloadURL(http2Opts, "https://storage.googleapis.com/replicate-weights/bert-base-uncased-hf-cache.tar", b)
+	benchmarkDownloadURL(http2Opts, "https://weights.replicate.delivery/default/bert-base-uncased-hf-cache.tar", b)
 }
 func BenchmarkDownloadLlama7bChatH1(b *testing.B) {
-	benchmarkDownloadURL(defaultOpts, "https://storage.googleapis.com/replicate-weights/Llama-2-7b-Chat-GPTQ/gptq_model-4bit-32g.safetensors", b)
+	benchmarkDownloadURL(defaultOpts, "https://weights.replicate.delivery/default/Llama-2-7b-Chat-GPTQ/gptq_model-4bit-32g.safetensors", b)
 }
 func BenchmarkDownloadLlama7bChatH2(b *testing.B) {
-	benchmarkDownloadURL(http2Opts, "https://storage.googleapis.com/replicate-weights/Llama-2-7b-Chat-GPTQ/gptq_model-4bit-32g.safetensors", b)
+	benchmarkDownloadURL(http2Opts, "https://weights.replicate.delivery/default/Llama-2-7b-Chat-GPTQ/gptq_model-4bit-32g.safetensors", b)
 }

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)
 	}