maurice/snix
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

chore(tvix/nar-bridge): move to Export from storev1pb

Browse source 
This removes the Export method in nar-bridge, and updates all users to
the version now in storev1pb.

It moves the roundtrip test to the importer crate, and some of the
utility functions into a separate util_test.go file.

Change-Id: I81d9e0b35dfd78ef1042bed307281eecd2aaa2a8
Reviewed-on: https://cl.tvl.fyi/c/depot/+/9603
Reviewed-by: Brian McGee <brian@bmcgee.ie>
Tested-by: BuildkiteCI
This commit is contained in:
Florian Klink 2023-10-10 20:52:03 +02:00 • committed by flokli
parent e6ba84ea50
commit fe963ae0a3
8 changed files with 49 additions and 335 deletions
Download patch file Download diff file Expand all files Collapse all files

276
tvix/nar-bridge/pkg/exporter/export.go
View file

@ -1,276 +0,0 @@
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)
}
}
}
}
// 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
}
}

3
tvix/nar-bridge/pkg/http/nar_get.go
View file

@ -13,7 +13,6 @@ import (
"sync"
castorev1pb "code.tvl.fyi/tvix/castore/protos"
"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 +93,7 @@ func renderNar(
}
// render the NAR file
err := exporter.Export(
err := storev1pb.Export(
w,
pathInfo,
func(directoryDigest []byte) (*castorev1pb.Directory, error) {

26
tvix/nar-bridge/pkg/importer/importer_test.go
View file

@ -11,35 +11,9 @@ import (
castorev1pb "code.tvl.fyi/tvix/castore/protos"
"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"
"google.golang.org/protobuf/testing/protocmp"
"lukechampine.com/blake3"
)
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 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) {
f, err := os.Open("../../testdata/symlink.nar")
require.NoError(t, err)

31
tvix/nar-bridge/pkg/exporter/full_test.go → tvix/nar-bridge/pkg/importer/roundtrip_test.go
View file

@ -1,4 +1,4 @@
package exporter_test
package importer_test
import (
"bytes"
@ -10,15 +10,15 @@ import (
"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 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.
func TestRoundtrip(t *testing.T) {
// We pipe nar_1094wph9z4nwlgvsd53abfz8i117ykiv5dwnq9nnhz846s7xqd7d.nar to
// storev1pb.Export, 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")
@ -57,7 +57,7 @@ func TestFull(t *testing.T) {
// done populating everything, now actually test the export :-)
var buf bytes.Buffer
err = exporter.Export(
err = storev1pb.Export(
&buf,
pathInfo,
func(directoryDgst []byte) (*castorev1pb.Directory, error) {
@ -79,20 +79,3 @@ func TestFull(t *testing.T) {
require.NoError(t, err, "exporter shouldn't fail")
require.Equal(t, narContents, buf.Bytes())
}
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{})
}

34
tvix/nar-bridge/pkg/importer/util_test.go Normal file
View file

@ -0,0 +1,34 @@
package importer_test
import (
"io"
"testing"
castorev1pb "code.tvl.fyi/tvix/castore/protos"
"github.com/google/go-cmp/cmp"
"google.golang.org/protobuf/testing/protocmp"
"lukechampine.com/blake3"
)
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 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{})
}