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