From d3c44ab03d14bc143e790426bede465feb871cd5 Mon Sep 17 00:00:00 2001 From: Florian Klink Date: Tue, 3 Oct 2023 14:29:24 +0300 Subject: refactor(tvix/nar-bridge): rename writer package to exporter Change-Id: I7cf6a3ed0dca4c450fadd9228dc0154ad6254f51 Reviewed-on: https://cl.tvl.fyi/c/depot/+/9530 Tested-by: BuildkiteCI Reviewed-by: Connor Brewster Autosubmit: flokli --- tvix/nar-bridge/pkg/exporter/export.go | 279 +++++++++++++++++++++ tvix/nar-bridge/pkg/exporter/export_test.go | 220 ++++++++++++++++ .../nar-bridge/pkg/exporter/pick_next_node_test.go | 51 ++++ tvix/nar-bridge/pkg/server/nar_get.go | 4 +- tvix/nar-bridge/pkg/writer/writer.go | 279 --------------------- .../pkg/writer/writer_pick_next_node_test.go | 51 ---- tvix/nar-bridge/pkg/writer/writer_test.go | 220 ---------------- 7 files changed, 552 insertions(+), 552 deletions(-) create mode 100644 tvix/nar-bridge/pkg/exporter/export.go create mode 100644 tvix/nar-bridge/pkg/exporter/export_test.go create mode 100644 tvix/nar-bridge/pkg/exporter/pick_next_node_test.go delete mode 100644 tvix/nar-bridge/pkg/writer/writer.go delete mode 100644 tvix/nar-bridge/pkg/writer/writer_pick_next_node_test.go delete mode 100644 tvix/nar-bridge/pkg/writer/writer_test.go diff --git a/tvix/nar-bridge/pkg/exporter/export.go b/tvix/nar-bridge/pkg/exporter/export.go new file mode 100644 index 000000000000..40ccc4b6380c --- /dev/null +++ b/tvix/nar-bridge/pkg/exporter/export.go @@ -0,0 +1,279 @@ +package exporter + +import ( + "fmt" + "io" + "path" + + castorev1pb "code.tvl.fyi/tvix/castore/protos" + storev1pb "code.tvl.fyi/tvix/store/protos" + "github.com/nix-community/go-nix/pkg/nar" +) + +type DirectoryLookupFn func([]byte) (*castorev1pb.Directory, error) +type BlobLookupFn func([]byte) (io.ReadCloser, error) + +// Export will traverse a given PathInfo structure, and write the contents +// in NAR format to the passed Writer. +// It uses directoryLookupFn and blobLookupFn to resolve references. +func Export( + w io.Writer, + pathInfo *storev1pb.PathInfo, + directoryLookupFn DirectoryLookupFn, + blobLookupFn BlobLookupFn, +) error { + // initialize a NAR writer + narWriter, err := nar.NewWriter(w) + if err != nil { + return fmt.Errorf("unable to initialize nar writer: %w", err) + } + defer narWriter.Close() + + // populate rootHeader + rootHeader := &nar.Header{ + Path: "/", + } + + // populate a stack + // we will push paths and directories to it when entering a directory, + // and emit individual elements to the NAR writer, draining the Directory object. + // once it's empty, we can pop it off the stack. + var stackPaths = []string{} + var stackDirectories = []*castorev1pb.Directory{} + + // peek at the pathInfo root and assemble the root node and write to writer + // in the case of a regular file, we retrieve and write the contents, close and exit + // in the case of a symlink, we write the symlink, close and exit + switch v := (pathInfo.GetNode().GetNode()).(type) { + case *castorev1pb.Node_File: + rootHeader.Type = nar.TypeRegular + rootHeader.Size = int64(v.File.GetSize()) + rootHeader.Executable = v.File.GetExecutable() + err := narWriter.WriteHeader(rootHeader) + if err != nil { + return fmt.Errorf("unable to write root header: %w", err) + } + + // if it's a regular file, retrieve and write the contents + blobReader, err := blobLookupFn(v.File.GetDigest()) + if err != nil { + return fmt.Errorf("unable to lookup blob: %w", err) + } + defer blobReader.Close() + + _, err = io.Copy(narWriter, blobReader) + if err != nil { + return fmt.Errorf("unable to read from blobReader: %w", err) + } + + err = blobReader.Close() + if err != nil { + return fmt.Errorf("unable to close content reader: %w", err) + } + + err = narWriter.Close() + if err != nil { + return fmt.Errorf("unable to close nar reader: %w", err) + } + + return nil + + case *castorev1pb.Node_Symlink: + rootHeader.Type = nar.TypeSymlink + rootHeader.LinkTarget = string(v.Symlink.GetTarget()) + err := narWriter.WriteHeader(rootHeader) + if err != nil { + return fmt.Errorf("unable to write root header: %w", err) + } + + err = narWriter.Close() + if err != nil { + return fmt.Errorf("unable to close nar reader: %w", err) + } + + return nil + case *castorev1pb.Node_Directory: + // We have a directory at the root, look it up and put in on the stack. + directory, err := directoryLookupFn(v.Directory.Digest) + if err != nil { + return fmt.Errorf("unable to lookup directory: %w", err) + } + stackDirectories = append(stackDirectories, directory) + stackPaths = append(stackPaths, "/") + + err = narWriter.WriteHeader(&nar.Header{ + Path: "/", + Type: nar.TypeDirectory, + }) + + if err != nil { + return fmt.Errorf("error writing header: %w", err) + } + } + + // as long as the stack is not empty, we keep running. + for { + if len(stackDirectories) == 0 { + return nil + } + + // Peek at the current top of the stack. + topOfStack := stackDirectories[len(stackDirectories)-1] + topOfStackPath := stackPaths[len(stackPaths)-1] + + // get the next element that's lexicographically smallest, and drain it from + // the current directory on top of the stack. + nextNode := drainNextNode(topOfStack) + + // If nextNode returns nil, there's nothing left in the directory node, so we + // can emit it from the stack. + // Contrary to the import case, we don't emit the node popping from the stack, but when pushing. + if nextNode == nil { + // pop off stack + stackDirectories = stackDirectories[:len(stackDirectories)-1] + stackPaths = stackPaths[:len(stackPaths)-1] + + continue + } + + switch n := (nextNode).(type) { + case *castorev1pb.DirectoryNode: + err := narWriter.WriteHeader(&nar.Header{ + Path: path.Join(topOfStackPath, string(n.GetName())), + Type: nar.TypeDirectory, + }) + if err != nil { + return fmt.Errorf("unable to write nar header: %w", err) + } + + d, err := directoryLookupFn(n.GetDigest()) + if err != nil { + return fmt.Errorf("unable to lookup directory: %w", err) + } + + // add to stack + stackDirectories = append(stackDirectories, d) + stackPaths = append(stackPaths, path.Join(topOfStackPath, string(n.GetName()))) + case *castorev1pb.FileNode: + err := narWriter.WriteHeader(&nar.Header{ + Path: path.Join(topOfStackPath, string(n.GetName())), + Type: nar.TypeRegular, + Size: int64(n.GetSize()), + Executable: n.GetExecutable(), + }) + if err != nil { + return fmt.Errorf("unable to write nar header: %w", err) + } + + // copy file contents + contentReader, err := blobLookupFn(n.GetDigest()) + if err != nil { + return fmt.Errorf("unable to get blob: %w", err) + } + defer contentReader.Close() + + _, err = io.Copy(narWriter, contentReader) + if err != nil { + return fmt.Errorf("unable to copy contents from contentReader: %w", err) + } + + err = contentReader.Close() + if err != nil { + return fmt.Errorf("unable to close content reader: %w", err) + } + case *castorev1pb.SymlinkNode: + err := narWriter.WriteHeader(&nar.Header{ + Path: path.Join(topOfStackPath, string(n.GetName())), + Type: nar.TypeSymlink, + LinkTarget: string(n.GetTarget()), + }) + if err != nil { + return fmt.Errorf("unable to write nar header: %w", err) + } + } + } +} + +// TODO: add validation functions to Directory in both rust and golang, to +// validate the keys in directories, files and symlinks are sorted. + +// drainNextNode will drain a directory message with one of its child nodes, +// whichever comes first alphabetically. +func drainNextNode(d *castorev1pb.Directory) interface{} { + switch v := (smallestNode(d)).(type) { + case *castorev1pb.DirectoryNode: + d.Directories = d.Directories[1:] + return v + case *castorev1pb.FileNode: + d.Files = d.Files[1:] + return v + case *castorev1pb.SymlinkNode: + d.Symlinks = d.Symlinks[1:] + return v + case nil: + return nil + default: + panic("invalid type encountered") + } +} + +// smallestNode will return the node from a directory message, +// whichever comes first alphabetically. +func smallestNode(d *castorev1pb.Directory) interface{} { + childDirectories := d.GetDirectories() + childFiles := d.GetFiles() + childSymlinks := d.GetSymlinks() + + if len(childDirectories) > 0 { + if len(childFiles) > 0 { + if len(childSymlinks) > 0 { + // directories,files,symlinks + return smallerNode(smallerNode(childDirectories[0], childFiles[0]), childSymlinks[0]) + } else { + // directories,files,!symlinks + return smallerNode(childDirectories[0], childFiles[0]) + } + } else { + // directories,!files + if len(childSymlinks) > 0 { + // directories,!files,symlinks + return smallerNode(childDirectories[0], childSymlinks[0]) + } else { + // directories,!files,!symlinks + return childDirectories[0] + } + } + } else { + // !directories + if len(childFiles) > 0 { + // !directories,files + if len(childSymlinks) > 0 { + // !directories,files,symlinks + return smallerNode(childFiles[0], childSymlinks[0]) + } else { + // !directories,files,!symlinks + return childFiles[0] + } + } else { + //!directories,!files + if len(childSymlinks) > 0 { + //!directories,!files,symlinks + return childSymlinks[0] + } else { + //!directories,!files,!symlinks + return nil + } + } + } +} + +// smallerNode compares two nodes by their name, +// and returns the one with the smaller name. +// both nodes may not be nil, we do check for these cases in smallestNode. +func smallerNode(a interface{ GetName() []byte }, b interface{ GetName() []byte }) interface{ GetName() []byte } { + if string(a.GetName()) < string(b.GetName()) { + return a + } else { + return b + } +} diff --git a/tvix/nar-bridge/pkg/exporter/export_test.go b/tvix/nar-bridge/pkg/exporter/export_test.go new file mode 100644 index 000000000000..ffee62f29183 --- /dev/null +++ b/tvix/nar-bridge/pkg/exporter/export_test.go @@ -0,0 +1,220 @@ +package exporter_test + +import ( + "bytes" + "context" + "encoding/base64" + "fmt" + "io" + "os" + "testing" + + castorev1pb "code.tvl.fyi/tvix/castore/protos" + "code.tvl.fyi/tvix/nar-bridge/pkg/exporter" + "code.tvl.fyi/tvix/nar-bridge/pkg/importer" + storev1pb "code.tvl.fyi/tvix/store/protos" + "github.com/stretchr/testify/require" + "lukechampine.com/blake3" +) + +func mustDirectoryDigest(d *castorev1pb.Directory) []byte { + dgst, err := d.Digest() + if err != nil { + panic(err) + } + return dgst +} + +func mustBlobDigest(r io.Reader) []byte { + hasher := blake3.New(32, nil) + _, err := io.Copy(hasher, r) + if err != nil { + panic(err) + } + return hasher.Sum([]byte{}) +} + +func TestSymlink(t *testing.T) { + pathInfo := &storev1pb.PathInfo{ + + Node: &castorev1pb.Node{ + Node: &castorev1pb.Node_Symlink{ + Symlink: &castorev1pb.SymlinkNode{ + Name: []byte("doesntmatter"), + Target: []byte("/nix/store/somewhereelse"), + }, + }, + }, + } + + var buf bytes.Buffer + + err := exporter.Export(&buf, pathInfo, func([]byte) (*castorev1pb.Directory, error) { + panic("no directories expected") + }, func([]byte) (io.ReadCloser, error) { + panic("no files expected") + }) + require.NoError(t, err, "exporter shouldn't fail") + + f, err := os.Open("../../testdata/symlink.nar") + require.NoError(t, err) + + bytesExpected, err := io.ReadAll(f) + if err != nil { + panic(err) + } + + require.Equal(t, bytesExpected, buf.Bytes(), "expected nar contents to match") +} + +func TestRegular(t *testing.T) { + // The blake3 digest of the 0x01 byte. + BLAKE3_DIGEST_0X01 := []byte{ + 0x48, 0xfc, 0x72, 0x1f, 0xbb, 0xc1, 0x72, 0xe0, 0x92, 0x5f, 0xa2, 0x7a, 0xf1, 0x67, 0x1d, + 0xe2, 0x25, 0xba, 0x92, 0x71, 0x34, 0x80, 0x29, 0x98, 0xb1, 0x0a, 0x15, 0x68, 0xa1, 0x88, + 0x65, 0x2b, + } + + pathInfo := &storev1pb.PathInfo{ + Node: &castorev1pb.Node{ + Node: &castorev1pb.Node_File{ + File: &castorev1pb.FileNode{ + Name: []byte("doesntmatter"), + Digest: BLAKE3_DIGEST_0X01, + Size: 1, + Executable: false, + }, + }, + }, + } + + var buf bytes.Buffer + + err := exporter.Export(&buf, pathInfo, func([]byte) (*castorev1pb.Directory, error) { + panic("no directories expected") + }, func(blobRef []byte) (io.ReadCloser, error) { + if !bytes.Equal(blobRef, BLAKE3_DIGEST_0X01) { + panic("unexpected blobref") + } + return io.NopCloser(bytes.NewBuffer([]byte{0x01})), nil + }) + require.NoError(t, err, "exporter shouldn't fail") + + f, err := os.Open("../../testdata/onebyteregular.nar") + require.NoError(t, err) + + bytesExpected, err := io.ReadAll(f) + if err != nil { + panic(err) + } + + require.Equal(t, bytesExpected, buf.Bytes(), "expected nar contents to match") +} + +func TestEmptyDirectory(t *testing.T) { + // construct empty directory node this refers to + emptyDirectory := &castorev1pb.Directory{ + Directories: []*castorev1pb.DirectoryNode{}, + Files: []*castorev1pb.FileNode{}, + Symlinks: []*castorev1pb.SymlinkNode{}, + } + emptyDirectoryDigest := mustDirectoryDigest(emptyDirectory) + + pathInfo := &storev1pb.PathInfo{ + Node: &castorev1pb.Node{ + Node: &castorev1pb.Node_Directory{ + Directory: &castorev1pb.DirectoryNode{ + Name: []byte("doesntmatter"), + Digest: emptyDirectoryDigest, + Size: 0, + }, + }, + }, + } + + var buf bytes.Buffer + + err := exporter.Export(&buf, pathInfo, func(directoryRef []byte) (*castorev1pb.Directory, error) { + if !bytes.Equal(directoryRef, emptyDirectoryDigest) { + panic("unexpected directoryRef") + } + return emptyDirectory, nil + }, func([]byte) (io.ReadCloser, error) { + panic("no files expected") + }) + require.NoError(t, err, "exporter shouldn't fail") + + f, err := os.Open("../../testdata/emptydirectory.nar") + require.NoError(t, err) + + bytesExpected, err := io.ReadAll(f) + if err != nil { + panic(err) + } + + require.Equal(t, bytesExpected, buf.Bytes(), "expected nar contents to match") +} + +func TestFull(t *testing.T) { + // We pipe nar_1094wph9z4nwlgvsd53abfz8i117ykiv5dwnq9nnhz846s7xqd7d.nar to the exporter, + // and store all the file contents and directory objects received in two hashmaps. + // We then feed it to the writer, and test we come up with the same NAR file. + + f, err := os.Open("../../testdata/nar_1094wph9z4nwlgvsd53abfz8i117ykiv5dwnq9nnhz846s7xqd7d.nar") + require.NoError(t, err) + + narContents, err := io.ReadAll(f) + require.NoError(t, err) + + blobsMap := make(map[string][]byte, 0) + directoriesMap := make(map[string]*castorev1pb.Directory) + + pathInfo, err := importer.Import( + context.Background(), + bytes.NewBuffer(narContents), + func(blobReader io.Reader) ([]byte, error) { + // read in contents, we need to put it into filesMap later. + contents, err := io.ReadAll(blobReader) + require.NoError(t, err) + + dgst := mustBlobDigest(bytes.NewReader(contents)) + + // put it in filesMap + blobsMap[base64.StdEncoding.EncodeToString(dgst)] = contents + + return dgst, nil + }, + func(directory *castorev1pb.Directory) ([]byte, error) { + dgst := mustDirectoryDigest(directory) + + directoriesMap[base64.StdEncoding.EncodeToString(dgst)] = directory + return dgst, nil + }, + ) + + require.NoError(t, err) + + // done populating everything, now actually test the export :-) + var buf bytes.Buffer + err = exporter.Export( + &buf, + pathInfo, + func(directoryDgst []byte) (*castorev1pb.Directory, error) { + d, found := directoriesMap[base64.StdEncoding.EncodeToString(directoryDgst)] + if !found { + panic(fmt.Sprintf("directory %v not found", base64.StdEncoding.EncodeToString(directoryDgst))) + } + return d, nil + }, + func(blobDgst []byte) (io.ReadCloser, error) { + blobContents, found := blobsMap[base64.StdEncoding.EncodeToString(blobDgst)] + if !found { + panic(fmt.Sprintf("blob %v not found", base64.StdEncoding.EncodeToString(blobDgst))) + } + return io.NopCloser(bytes.NewReader(blobContents)), nil + }, + ) + + require.NoError(t, err, "exporter shouldn't fail") + require.Equal(t, narContents, buf.Bytes()) +} diff --git a/tvix/nar-bridge/pkg/exporter/pick_next_node_test.go b/tvix/nar-bridge/pkg/exporter/pick_next_node_test.go new file mode 100644 index 000000000000..c80261ae6b6a --- /dev/null +++ b/tvix/nar-bridge/pkg/exporter/pick_next_node_test.go @@ -0,0 +1,51 @@ +package exporter + +import ( + "testing" + + castorev1pb "code.tvl.fyi/tvix/castore/protos" + "github.com/google/go-cmp/cmp" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/testing/protocmp" +) + +func requireProtoEq(t *testing.T, expected interface{}, actual interface{}) { + if diff := cmp.Diff(expected, actual, protocmp.Transform()); diff != "" { + t.Errorf("unexpected difference:\n%v", diff) + } +} + +func TestPopNextNode(t *testing.T) { + t.Run("empty directory", func(t *testing.T) { + d := &castorev1pb.Directory{ + Directories: []*castorev1pb.DirectoryNode{}, + Files: []*castorev1pb.FileNode{}, + Symlinks: []*castorev1pb.SymlinkNode{}, + } + + n := drainNextNode(d) + require.Equal(t, nil, n) + }) + t.Run("only directories", func(t *testing.T) { + ds := &castorev1pb.Directory{ + Directories: []*castorev1pb.DirectoryNode{{ + Name: []byte("a"), + Digest: []byte{}, + Size: 0, + }, { + Name: []byte("b"), + Digest: []byte{}, + Size: 0, + }}, + Files: []*castorev1pb.FileNode{}, + Symlinks: []*castorev1pb.SymlinkNode{}, + } + + n := drainNextNode(ds) + requireProtoEq(t, &castorev1pb.DirectoryNode{ + Name: []byte("a"), + Digest: []byte{}, + Size: 0, + }, n) + }) +} diff --git a/tvix/nar-bridge/pkg/server/nar_get.go b/tvix/nar-bridge/pkg/server/nar_get.go index 7397992611fc..3ccb8d658d72 100644 --- a/tvix/nar-bridge/pkg/server/nar_get.go +++ b/tvix/nar-bridge/pkg/server/nar_get.go @@ -13,7 +13,7 @@ import ( "sync" castorev1pb "code.tvl.fyi/tvix/castore/protos" - "code.tvl.fyi/tvix/nar-bridge/pkg/writer" + "code.tvl.fyi/tvix/nar-bridge/pkg/exporter" storev1pb "code.tvl.fyi/tvix/store/protos" "github.com/go-chi/chi/v5" nixhash "github.com/nix-community/go-nix/pkg/hash" @@ -94,7 +94,7 @@ func renderNar( } // render the NAR file - err := writer.Export( + err := exporter.Export( w, pathInfo, func(directoryDigest []byte) (*castorev1pb.Directory, error) { diff --git a/tvix/nar-bridge/pkg/writer/writer.go b/tvix/nar-bridge/pkg/writer/writer.go deleted file mode 100644 index fd154e061bf0..000000000000 --- a/tvix/nar-bridge/pkg/writer/writer.go +++ /dev/null @@ -1,279 +0,0 @@ -package writer - -import ( - "fmt" - "io" - "path" - - castorev1pb "code.tvl.fyi/tvix/castore/protos" - storev1pb "code.tvl.fyi/tvix/store/protos" - "github.com/nix-community/go-nix/pkg/nar" -) - -type DirectoryLookupFn func([]byte) (*castorev1pb.Directory, error) -type BlobLookupFn func([]byte) (io.ReadCloser, error) - -// Export will traverse a given pathInfo structure, and write the contents -// in NAR format to the passed Writer. -// It uses directoryLookupFn and blobLookupFn to resolve references. -func Export( - w io.Writer, - pathInfo *storev1pb.PathInfo, - directoryLookupFn DirectoryLookupFn, - blobLookupFn BlobLookupFn, -) error { - // initialize a NAR writer - narWriter, err := nar.NewWriter(w) - if err != nil { - return fmt.Errorf("unable to initialize nar writer: %w", err) - } - defer narWriter.Close() - - // populate rootHeader - rootHeader := &nar.Header{ - Path: "/", - } - - // populate a stack - // we will push paths and directories to it when entering a directory, - // and emit individual elements to the NAR writer, draining the Directory object. - // once it's empty, we can pop it off the stack. - var stackPaths = []string{} - var stackDirectories = []*castorev1pb.Directory{} - - // peek at the pathInfo root and assemble the root node and write to writer - // in the case of a regular file, we retrieve and write the contents, close and exit - // in the case of a symlink, we write the symlink, close and exit - switch v := (pathInfo.GetNode().GetNode()).(type) { - case *castorev1pb.Node_File: - rootHeader.Type = nar.TypeRegular - rootHeader.Size = int64(v.File.GetSize()) - rootHeader.Executable = v.File.GetExecutable() - err := narWriter.WriteHeader(rootHeader) - if err != nil { - return fmt.Errorf("unable to write root header: %w", err) - } - - // if it's a regular file, retrieve and write the contents - blobReader, err := blobLookupFn(v.File.GetDigest()) - if err != nil { - return fmt.Errorf("unable to lookup blob: %w", err) - } - defer blobReader.Close() - - _, err = io.Copy(narWriter, blobReader) - if err != nil { - return fmt.Errorf("unable to read from blobReader: %w", err) - } - - err = blobReader.Close() - if err != nil { - return fmt.Errorf("unable to close content reader: %w", err) - } - - err = narWriter.Close() - if err != nil { - return fmt.Errorf("unable to close nar reader: %w", err) - } - - return nil - - case *castorev1pb.Node_Symlink: - rootHeader.Type = nar.TypeSymlink - rootHeader.LinkTarget = string(v.Symlink.GetTarget()) - err := narWriter.WriteHeader(rootHeader) - if err != nil { - return fmt.Errorf("unable to write root header: %w", err) - } - - err = narWriter.Close() - if err != nil { - return fmt.Errorf("unable to close nar reader: %w", err) - } - - return nil - case *castorev1pb.Node_Directory: - // We have a directory at the root, look it up and put in on the stack. - directory, err := directoryLookupFn(v.Directory.Digest) - if err != nil { - return fmt.Errorf("unable to lookup directory: %w", err) - } - stackDirectories = append(stackDirectories, directory) - stackPaths = append(stackPaths, "/") - - err = narWriter.WriteHeader(&nar.Header{ - Path: "/", - Type: nar.TypeDirectory, - }) - - if err != nil { - return fmt.Errorf("error writing header: %w", err) - } - } - - // as long as the stack is not empty, we keep running. - for { - if len(stackDirectories) == 0 { - return nil - } - - // Peek at the current top of the stack. - topOfStack := stackDirectories[len(stackDirectories)-1] - topOfStackPath := stackPaths[len(stackPaths)-1] - - // get the next element that's lexicographically smallest, and drain it from - // the current directory on top of the stack. - nextNode := drainNextNode(topOfStack) - - // If nextNode returns nil, there's nothing left in the directory node, so we - // can emit it from the stack. - // Contrary to the import case, we don't emit the node popping from the stack, but when pushing. - if nextNode == nil { - // pop off stack - stackDirectories = stackDirectories[:len(stackDirectories)-1] - stackPaths = stackPaths[:len(stackPaths)-1] - - continue - } - - switch n := (nextNode).(type) { - case *castorev1pb.DirectoryNode: - err := narWriter.WriteHeader(&nar.Header{ - Path: path.Join(topOfStackPath, string(n.GetName())), - Type: nar.TypeDirectory, - }) - if err != nil { - return fmt.Errorf("unable to write nar header: %w", err) - } - - d, err := directoryLookupFn(n.GetDigest()) - if err != nil { - return fmt.Errorf("unable to lookup directory: %w", err) - } - - // add to stack - stackDirectories = append(stackDirectories, d) - stackPaths = append(stackPaths, path.Join(topOfStackPath, string(n.GetName()))) - case *castorev1pb.FileNode: - err := narWriter.WriteHeader(&nar.Header{ - Path: path.Join(topOfStackPath, string(n.GetName())), - Type: nar.TypeRegular, - Size: int64(n.GetSize()), - Executable: n.GetExecutable(), - }) - if err != nil { - return fmt.Errorf("unable to write nar header: %w", err) - } - - // copy file contents - contentReader, err := blobLookupFn(n.GetDigest()) - if err != nil { - return fmt.Errorf("unable to get blob: %w", err) - } - defer contentReader.Close() - - _, err = io.Copy(narWriter, contentReader) - if err != nil { - return fmt.Errorf("unable to copy contents from contentReader: %w", err) - } - - err = contentReader.Close() - if err != nil { - return fmt.Errorf("unable to close content reader: %w", err) - } - case *castorev1pb.SymlinkNode: - err := narWriter.WriteHeader(&nar.Header{ - Path: path.Join(topOfStackPath, string(n.GetName())), - Type: nar.TypeSymlink, - LinkTarget: string(n.GetTarget()), - }) - if err != nil { - return fmt.Errorf("unable to write nar header: %w", err) - } - } - } -} - -// TODO: add validation functions to Directory in both rust and golang, to -// validate the keys in directories, files and symlinks are sorted. - -// drainNextNode will drain a directory message with one of its child nodes, -// whichever comes first alphabetically. -func drainNextNode(d *castorev1pb.Directory) interface{} { - switch v := (smallestNode(d)).(type) { - case *castorev1pb.DirectoryNode: - d.Directories = d.Directories[1:] - return v - case *castorev1pb.FileNode: - d.Files = d.Files[1:] - return v - case *castorev1pb.SymlinkNode: - d.Symlinks = d.Symlinks[1:] - return v - case nil: - return nil - default: - panic("invalid type encountered") - } -} - -// smallestNode will return the node from a directory message, -// whichever comes first alphabetically. -func smallestNode(d *castorev1pb.Directory) interface{} { - childDirectories := d.GetDirectories() - childFiles := d.GetFiles() - childSymlinks := d.GetSymlinks() - - if len(childDirectories) > 0 { - if len(childFiles) > 0 { - if len(childSymlinks) > 0 { - // directories,files,symlinks - return smallerNode(smallerNode(childDirectories[0], childFiles[0]), childSymlinks[0]) - } else { - // directories,files,!symlinks - return smallerNode(childDirectories[0], childFiles[0]) - } - } else { - // directories,!files - if len(childSymlinks) > 0 { - // directories,!files,symlinks - return smallerNode(childDirectories[0], childSymlinks[0]) - } else { - // directories,!files,!symlinks - return childDirectories[0] - } - } - } else { - // !directories - if len(childFiles) > 0 { - // !directories,files - if len(childSymlinks) > 0 { - // !directories,files,symlinks - return smallerNode(childFiles[0], childSymlinks[0]) - } else { - // !directories,files,!symlinks - return childFiles[0] - } - } else { - //!directories,!files - if len(childSymlinks) > 0 { - //!directories,!files,symlinks - return childSymlinks[0] - } else { - //!directories,!files,!symlinks - return nil - } - } - } -} - -// smallerNode compares two nodes by their name, -// and returns the one with the smaller name. -// both nodes may not be nil, we do check for these cases in smallestNode. -func smallerNode(a interface{ GetName() []byte }, b interface{ GetName() []byte }) interface{ GetName() []byte } { - if string(a.GetName()) < string(b.GetName()) { - return a - } else { - return b - } -} diff --git a/tvix/nar-bridge/pkg/writer/writer_pick_next_node_test.go b/tvix/nar-bridge/pkg/writer/writer_pick_next_node_test.go deleted file mode 100644 index d69adca89cb7..000000000000 --- a/tvix/nar-bridge/pkg/writer/writer_pick_next_node_test.go +++ /dev/null @@ -1,51 +0,0 @@ -package writer - -import ( - "testing" - - castorev1pb "code.tvl.fyi/tvix/castore/protos" - "github.com/google/go-cmp/cmp" - "github.com/stretchr/testify/require" - "google.golang.org/protobuf/testing/protocmp" -) - -func requireProtoEq(t *testing.T, expected interface{}, actual interface{}) { - if diff := cmp.Diff(expected, actual, protocmp.Transform()); diff != "" { - t.Errorf("unexpected difference:\n%v", diff) - } -} - -func TestPopNextNode(t *testing.T) { - t.Run("empty directory", func(t *testing.T) { - d := &castorev1pb.Directory{ - Directories: []*castorev1pb.DirectoryNode{}, - Files: []*castorev1pb.FileNode{}, - Symlinks: []*castorev1pb.SymlinkNode{}, - } - - n := drainNextNode(d) - require.Equal(t, nil, n) - }) - t.Run("only directories", func(t *testing.T) { - ds := &castorev1pb.Directory{ - Directories: []*castorev1pb.DirectoryNode{{ - Name: []byte("a"), - Digest: []byte{}, - Size: 0, - }, { - Name: []byte("b"), - Digest: []byte{}, - Size: 0, - }}, - Files: []*castorev1pb.FileNode{}, - Symlinks: []*castorev1pb.SymlinkNode{}, - } - - n := drainNextNode(ds) - requireProtoEq(t, &castorev1pb.DirectoryNode{ - Name: []byte("a"), - Digest: []byte{}, - Size: 0, - }, n) - }) -} diff --git a/tvix/nar-bridge/pkg/writer/writer_test.go b/tvix/nar-bridge/pkg/writer/writer_test.go deleted file mode 100644 index 881045b173d5..000000000000 --- a/tvix/nar-bridge/pkg/writer/writer_test.go +++ /dev/null @@ -1,220 +0,0 @@ -package writer_test - -import ( - "bytes" - "context" - "encoding/base64" - "fmt" - "io" - "os" - "testing" - - castorev1pb "code.tvl.fyi/tvix/castore/protos" - "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" - "lukechampine.com/blake3" -) - -func mustDirectoryDigest(d *castorev1pb.Directory) []byte { - dgst, err := d.Digest() - if err != nil { - panic(err) - } - return dgst -} - -func mustBlobDigest(r io.Reader) []byte { - hasher := blake3.New(32, nil) - _, err := io.Copy(hasher, r) - if err != nil { - panic(err) - } - return hasher.Sum([]byte{}) -} - -func TestSymlink(t *testing.T) { - pathInfo := &storev1pb.PathInfo{ - - Node: &castorev1pb.Node{ - Node: &castorev1pb.Node_Symlink{ - Symlink: &castorev1pb.SymlinkNode{ - Name: []byte("doesntmatter"), - Target: []byte("/nix/store/somewhereelse"), - }, - }, - }, - } - - var buf bytes.Buffer - - err := writer.Export(&buf, pathInfo, func([]byte) (*castorev1pb.Directory, error) { - panic("no directories expected") - }, func([]byte) (io.ReadCloser, error) { - panic("no files expected") - }) - require.NoError(t, err, "exporter shouldn't fail") - - f, err := os.Open("../../testdata/symlink.nar") - require.NoError(t, err) - - bytesExpected, err := io.ReadAll(f) - if err != nil { - panic(err) - } - - require.Equal(t, bytesExpected, buf.Bytes(), "expected nar contents to match") -} - -func TestRegular(t *testing.T) { - // The blake3 digest of the 0x01 byte. - BLAKE3_DIGEST_0X01 := []byte{ - 0x48, 0xfc, 0x72, 0x1f, 0xbb, 0xc1, 0x72, 0xe0, 0x92, 0x5f, 0xa2, 0x7a, 0xf1, 0x67, 0x1d, - 0xe2, 0x25, 0xba, 0x92, 0x71, 0x34, 0x80, 0x29, 0x98, 0xb1, 0x0a, 0x15, 0x68, 0xa1, 0x88, - 0x65, 0x2b, - } - - pathInfo := &storev1pb.PathInfo{ - Node: &castorev1pb.Node{ - Node: &castorev1pb.Node_File{ - File: &castorev1pb.FileNode{ - Name: []byte("doesntmatter"), - Digest: BLAKE3_DIGEST_0X01, - Size: 1, - Executable: false, - }, - }, - }, - } - - var buf bytes.Buffer - - err := writer.Export(&buf, pathInfo, func([]byte) (*castorev1pb.Directory, error) { - panic("no directories expected") - }, func(blobRef []byte) (io.ReadCloser, error) { - if !bytes.Equal(blobRef, BLAKE3_DIGEST_0X01) { - panic("unexpected blobref") - } - return io.NopCloser(bytes.NewBuffer([]byte{0x01})), nil - }) - require.NoError(t, err, "exporter shouldn't fail") - - f, err := os.Open("../../testdata/onebyteregular.nar") - require.NoError(t, err) - - bytesExpected, err := io.ReadAll(f) - if err != nil { - panic(err) - } - - require.Equal(t, bytesExpected, buf.Bytes(), "expected nar contents to match") -} - -func TestEmptyDirectory(t *testing.T) { - // construct empty directory node this refers to - emptyDirectory := &castorev1pb.Directory{ - Directories: []*castorev1pb.DirectoryNode{}, - Files: []*castorev1pb.FileNode{}, - Symlinks: []*castorev1pb.SymlinkNode{}, - } - emptyDirectoryDigest := mustDirectoryDigest(emptyDirectory) - - pathInfo := &storev1pb.PathInfo{ - Node: &castorev1pb.Node{ - Node: &castorev1pb.Node_Directory{ - Directory: &castorev1pb.DirectoryNode{ - Name: []byte("doesntmatter"), - Digest: emptyDirectoryDigest, - Size: 0, - }, - }, - }, - } - - var buf bytes.Buffer - - err := writer.Export(&buf, pathInfo, func(directoryRef []byte) (*castorev1pb.Directory, error) { - if !bytes.Equal(directoryRef, emptyDirectoryDigest) { - panic("unexpected directoryRef") - } - return emptyDirectory, nil - }, func([]byte) (io.ReadCloser, error) { - panic("no files expected") - }) - require.NoError(t, err, "exporter shouldn't fail") - - f, err := os.Open("../../testdata/emptydirectory.nar") - require.NoError(t, err) - - bytesExpected, err := io.ReadAll(f) - if err != nil { - panic(err) - } - - require.Equal(t, bytesExpected, buf.Bytes(), "expected nar contents to match") -} - -func TestFull(t *testing.T) { - // We pipe nar_1094wph9z4nwlgvsd53abfz8i117ykiv5dwnq9nnhz846s7xqd7d.nar to the exporter, - // and store all the file contents and directory objects received in two hashmaps. - // We then feed it to the writer, and test we come up with the same NAR file. - - f, err := os.Open("../../testdata/nar_1094wph9z4nwlgvsd53abfz8i117ykiv5dwnq9nnhz846s7xqd7d.nar") - require.NoError(t, err) - - narContents, err := io.ReadAll(f) - require.NoError(t, err) - - blobsMap := make(map[string][]byte, 0) - directoriesMap := make(map[string]*castorev1pb.Directory) - - pathInfo, err := importer.Import( - context.Background(), - bytes.NewBuffer(narContents), - func(blobReader io.Reader) ([]byte, error) { - // read in contents, we need to put it into filesMap later. - contents, err := io.ReadAll(blobReader) - require.NoError(t, err) - - dgst := mustBlobDigest(bytes.NewReader(contents)) - - // put it in filesMap - blobsMap[base64.StdEncoding.EncodeToString(dgst)] = contents - - return dgst, nil - }, - func(directory *castorev1pb.Directory) ([]byte, error) { - dgst := mustDirectoryDigest(directory) - - directoriesMap[base64.StdEncoding.EncodeToString(dgst)] = directory - return dgst, nil - }, - ) - - require.NoError(t, err) - - // done populating everything, now actually test the export :-) - var buf bytes.Buffer - err = writer.Export( - &buf, - pathInfo, - func(directoryDgst []byte) (*castorev1pb.Directory, error) { - d, found := directoriesMap[base64.StdEncoding.EncodeToString(directoryDgst)] - if !found { - panic(fmt.Sprintf("directory %v not found", base64.StdEncoding.EncodeToString(directoryDgst))) - } - return d, nil - }, - func(blobDgst []byte) (io.ReadCloser, error) { - blobContents, found := blobsMap[base64.StdEncoding.EncodeToString(blobDgst)] - if !found { - panic(fmt.Sprintf("blob %v not found", base64.StdEncoding.EncodeToString(blobDgst))) - } - return io.NopCloser(bytes.NewReader(blobContents)), nil - }, - ) - - require.NoError(t, err, "exporter shouldn't fail") - require.Equal(t, narContents, buf.Bytes()) -} -- cgit 1.4.1