refactor(tvix/nar-bridge): drop reader package r/6693

Make the import function usable on any reader.

Change-Id: I84d2004cb73cdd7a11fe8efb0f2efb6335d5e6b0
Reviewed-on: https://cl.tvl.fyi/c/depot/+/9527
Reviewed-by: Connor Brewster <cbrewster@hey.com>
Tested-by: BuildkiteCI
Autosubmit: flokli <flokli@flokli.de>

6 files changed, 41 insertions, 56 deletions

diff --git a/tvix/nar-bridge/pkg/reader/hashers.go b/tvix/nar-bridge/pkg/hashers/hashers.go
index 7e98d2062d21..0c9e611799e3 100644
--- a/tvix/nar-bridge/pkg/reader/hashers.go
+++ b/tvix/nar-bridge/pkg/hashers/hashers.go
@@ -1,4 +1,4 @@
-package reader
+package hashers
 
 import (
 	"errors"
diff --git a/tvix/nar-bridge/pkg/reader/reader.go b/tvix/nar-bridge/pkg/importer/importer.go
index 08fdd39a7f06..de5456be0b3e 100644
--- a/tvix/nar-bridge/pkg/reader/reader.go
+++ b/tvix/nar-bridge/pkg/importer/importer.go
@@ -1,4 +1,4 @@
-package reader
+package importer
 
 import (
 	"context"
@@ -10,45 +10,36 @@ import (
 	"strings"
 
 	castorev1pb "code.tvl.fyi/tvix/castore/protos"
+	"code.tvl.fyi/tvix/nar-bridge/pkg/hashers"
 	storev1pb "code.tvl.fyi/tvix/store/protos"
 	"github.com/nix-community/go-nix/pkg/nar"
 	"lukechampine.com/blake3"
 )
 
-type Reader struct {
-	hrSha256 *Hasher
-}
-
 // An item on the directories stack
 type stackItem struct {
 	path      string
 	directory *castorev1pb.Directory
 }
 
-func New(r io.Reader) *Reader {
-	// Instead of using the underlying reader itself, wrap the reader
-	// with a hasher calculating sha256 and one calculating sha512,
-	// and feed that one into the NAR reader.
-	hrSha256 := NewHasher(r, sha256.New())
-
-	return &Reader{
-		hrSha256: hrSha256,
-	}
-}
-
-// Import reads from the internally-wrapped reader,
-// and calls the callback functions whenever regular file contents are
-// encountered, or a Directory node is about to be finished.
-func (r *Reader) Import(
+// Import reads NAR from a reader, and returns a (sparsely populated) PathInfo
+// object.
+func Import(
+	// a context, to support cancellation
 	ctx context.Context,
+	// The reader the data is read from
+	r io.Reader,
 	// callback function called with each regular file content
 	blobCb func(fileReader io.Reader) error,
 	// callback function called with each finalized directory node
 	directoryCb func(directory *castorev1pb.Directory) error,
 ) (*storev1pb.PathInfo, error) {
+	// wrap the passed reader in a reader that records the number of bytes read and
+	// their sha256 sum.
+	hr := hashers.NewHasher(r, sha256.New())
 
-	// construct a NAR reader, by reading through hrSha256
-	narReader, err := nar.NewReader(r.hrSha256)
+	// construct a NAR reader from the underlying data.
+	narReader, err := nar.NewReader(hr)
 	if err != nil {
 		return nil, fmt.Errorf("failed to instantiate nar reader: %w", err)
 	}
@@ -144,8 +135,8 @@ func (r *Reader) Import(
 					Node:       nil,
 					References: [][]byte{},
 					Narinfo: &storev1pb.NARInfo{
-						NarSize:        uint64(r.hrSha256.BytesWritten()),
-						NarSha256:      r.hrSha256.Sum(nil),
+						NarSize:        uint64(hr.BytesWritten()),
+						NarSha256:      hr.Sum(nil),
 						Signatures:     []*storev1pb.NARInfo_Signature{},
 						ReferenceNames: []string{},
 					},
@@ -215,7 +206,7 @@ func (r *Reader) Import(
 			}
 			if hdr.Type == nar.TypeRegular {
 				// wrap reader with a reader calculating the blake3 hash
-				fileReader := NewHasher(narReader, blake3.New(32, nil))
+				fileReader := hashers.NewHasher(narReader, blake3.New(32, nil))
 
 				err := blobCb(fileReader)
 				if err != nil {
diff --git a/tvix/nar-bridge/pkg/reader/reader_test.go b/tvix/nar-bridge/pkg/importer/importer_test.go
index 560b06ebdd4e..0557c1d6dd1f 100644
--- a/tvix/nar-bridge/pkg/reader/reader_test.go
+++ b/tvix/nar-bridge/pkg/importer/importer_test.go
@@ -1,4 +1,4 @@
-package reader_test
+package importer_test
 
 import (
 	"context"
@@ -8,7 +8,7 @@ import (
 	"testing"
 
 	castorev1pb "code.tvl.fyi/tvix/castore/protos"
-	"code.tvl.fyi/tvix/nar-bridge/pkg/reader"
+	"code.tvl.fyi/tvix/nar-bridge/pkg/importer"
 	storev1pb "code.tvl.fyi/tvix/store/protos"
 	"github.com/google/go-cmp/cmp"
 	"github.com/stretchr/testify/require"
@@ -33,10 +33,9 @@ func TestSymlink(t *testing.T) {
 	f, err := os.Open("../../testdata/symlink.nar")
 	require.NoError(t, err)
 
-	r := reader.New(f)
-
-	actualPathInfo, err := r.Import(
+	actualPathInfo, err := importer.Import(
 		context.Background(),
+		f,
 		func(fileReader io.Reader) error {
 			panic("no file contents expected!")
 		}, func(directory *castorev1pb.Directory) error {
@@ -72,10 +71,9 @@ func TestRegular(t *testing.T) {
 	f, err := os.Open("../../testdata/onebyteregular.nar")
 	require.NoError(t, err)
 
-	r := reader.New(f)
-
-	actualPathInfo, err := r.Import(
+	actualPathInfo, err := importer.Import(
 		context.Background(),
+		f,
 		func(fileReader io.Reader) error {
 			contents, err := io.ReadAll(fileReader)
 			require.NoError(t, err, "reading fileReader should not error")
@@ -123,15 +121,14 @@ func TestEmptyDirectory(t *testing.T) {
 	f, err := os.Open("../../testdata/emptydirectory.nar")
 	require.NoError(t, err)
 
-	r := reader.New(f)
-
 	expectedDirectory := &castorev1pb.Directory{
 		Directories: []*castorev1pb.DirectoryNode{},
 		Files:       []*castorev1pb.FileNode{},
 		Symlinks:    []*castorev1pb.SymlinkNode{},
 	}
-	actualPathInfo, err := r.Import(
+	actualPathInfo, err := importer.Import(
 		context.Background(),
+		f,
 		func(fileReader io.Reader) error {
 			panic("no file contents expected!")
 		}, func(directory *castorev1pb.Directory) error {
@@ -168,8 +165,6 @@ func TestFull(t *testing.T) {
 	f, err := os.Open("../../testdata/nar_1094wph9z4nwlgvsd53abfz8i117ykiv5dwnq9nnhz846s7xqd7d.nar")
 	require.NoError(t, err)
 
-	r := reader.New(f)
-
 	expectedDirectoryPaths := []string{
 		"/bin",
 		"/share/man/man1",
@@ -478,8 +473,9 @@ func TestFull(t *testing.T) {
 
 	numDirectoriesReceived := 0
 
-	actualPathInfo, err := r.Import(
+	actualPathInfo, err := importer.Import(
 		context.Background(),
+		f,
 		func(fileReader io.Reader) error {
 			// Don't really bother reading and comparing the contents here,
 			// We already verify the right digests are produced by comparing the
@@ -533,12 +529,11 @@ func TestCallbackErrors(t *testing.T) {
 		f, err := os.Open("../../testdata/onebyteregular.nar")
 		require.NoError(t, err)
 
-		r := reader.New(f)
-
 		targetErr := errors.New("expected error")
 
-		_, err = r.Import(
+		_, err = importer.Import(
 			context.Background(),
+			f,
 			func(fileReader io.Reader) error {
 				return targetErr
 			}, func(directory *castorev1pb.Directory) error {
@@ -552,12 +547,11 @@ func TestCallbackErrors(t *testing.T) {
 		f, err := os.Open("../../testdata/emptydirectory.nar")
 		require.NoError(t, err)
 
-		r := reader.New(f)
-
 		targetErr := errors.New("expected error")
 
-		_, err = r.Import(
+		_, err = importer.Import(
 			context.Background(),
+			f,
 			func(fileReader io.Reader) error {
 				panic("no file contents expected!")
 			}, func(directory *castorev1pb.Directory) error {
@@ -585,9 +579,9 @@ func TestPopDirectories(t *testing.T) {
 	require.NoError(t, err)
 	defer f.Close()
 
-	r := reader.New(f)
-	_, err = r.Import(
+	_, err = importer.Import(
 		context.Background(),
+		f,
 		func(fileReader io.Reader) error { return nil },
 		func(directory *castorev1pb.Directory) error {
 			return directory.Validate()
diff --git a/tvix/nar-bridge/pkg/server/blob_upload.go b/tvix/nar-bridge/pkg/server/blob_upload.go
index 87d3918efa87..5531335367df 100644
--- a/tvix/nar-bridge/pkg/server/blob_upload.go
+++ b/tvix/nar-bridge/pkg/server/blob_upload.go
@@ -16,7 +16,7 @@ import (
 const chunkSize = 1024 * 1024
 
 // this produces a callback function that can be used as blobCb for the
-// reader.Import function call
+// importer.Import function call.
 func genBlobServiceWriteCb(ctx context.Context, blobServiceClient castorev1pb.BlobServiceClient) func(io.Reader) error {
 	return func(blobReader io.Reader) error {
 		// Ensure the blobReader is buffered to at least the chunk size.
diff --git a/tvix/nar-bridge/pkg/server/nar_put.go b/tvix/nar-bridge/pkg/server/nar_put.go
index 68529d6c3a17..0cb0190b7c6f 100644
--- a/tvix/nar-bridge/pkg/server/nar_put.go
+++ b/tvix/nar-bridge/pkg/server/nar_put.go
@@ -7,7 +7,7 @@ import (
 	"net/http"
 
 	castorev1pb "code.tvl.fyi/tvix/castore/protos"
-	"code.tvl.fyi/tvix/nar-bridge/pkg/reader"
+	"code.tvl.fyi/tvix/nar-bridge/pkg/importer"
 	"github.com/go-chi/chi/v5"
 	nixhash "github.com/nix-community/go-nix/pkg/hash"
 	"github.com/nix-community/go-nix/pkg/nixbase32"
@@ -39,10 +39,10 @@ func registerNarPut(s *Server) {
 		directoriesUploader := NewDirectoriesUploader(ctx, s.directoryServiceClient)
 		defer directoriesUploader.Done() //nolint:errcheck
 
-		// buffer the body by 10MiB
-		rd := reader.New(bufio.NewReaderSize(r.Body, 10*1024*1024))
-		pathInfo, err := rd.Import(
+		pathInfo, err := importer.Import(
 			ctx,
+			// buffer the body by 10MiB
+			bufio.NewReaderSize(r.Body, 10*1024*1024),
 			genBlobServiceWriteCb(ctx, s.blobServiceClient),
 			func(directory *castorev1pb.Directory) error {
 				return directoriesUploader.Put(directory)
diff --git a/tvix/nar-bridge/pkg/writer/writer_test.go b/tvix/nar-bridge/pkg/writer/writer_test.go
index 83ac717aa3d6..1c63ba78355f 100644
--- a/tvix/nar-bridge/pkg/writer/writer_test.go
+++ b/tvix/nar-bridge/pkg/writer/writer_test.go
@@ -9,7 +9,7 @@ import (
 	"testing"
 
 	castorev1pb "code.tvl.fyi/tvix/castore/protos"
-	"code.tvl.fyi/tvix/nar-bridge/pkg/reader"
+	"code.tvl.fyi/tvix/nar-bridge/pkg/importer"
 	"code.tvl.fyi/tvix/nar-bridge/pkg/writer"
 	storev1pb "code.tvl.fyi/tvix/store/protos"
 	"github.com/stretchr/testify/require"
@@ -159,9 +159,9 @@ func TestFull(t *testing.T) {
 	filesMap := make(map[string][]byte, 0)
 	directoriesMap := make(map[string]*castorev1pb.Directory)
 
-	r := reader.New(bytes.NewBuffer(narContents))
-	pathInfo, err := r.Import(
+	pathInfo, err := importer.Import(
 		context.Background(),
+		bytes.NewBuffer(narContents),
 		func(fileReader io.Reader) error {
 			fileContents, err := io.ReadAll(fileReader)
 			require.NoError(t, err)