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 +++++++++++++++++++++++++++++++++ 1 file changed, 279 insertions(+) create mode 100644 tvix/nar-bridge/pkg/exporter/export.go (limited to 'tvix/nar-bridge/pkg/exporter/export.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 + } +} -- cgit 1.4.1