tvl-depot/tvix/nar-bridge/pkg/reader/reader.go

package reader

import (
	"context"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	"path"
	"strings"

	castorev1pb "code.tvl.fyi/tvix/castore/protos"
	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(
	ctx context.Context,
	// 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) {

	// construct a NAR reader, by reading through hrSha256
	narReader, err := nar.NewReader(r.hrSha256)
	if err != nil {
		return nil, fmt.Errorf("failed to instantiate nar reader: %w", err)
	}
	defer narReader.Close()

	// If we store a symlink or regular file at the root, these are not nil.
	// If they are nil, we instead have a stackDirectory.
	var rootSymlink *castorev1pb.SymlinkNode
	var rootFile *castorev1pb.FileNode
	var stackDirectory *castorev1pb.Directory

	var stack = []stackItem{}

	// popFromStack is used when we transition to a different directory or
	// drain the stack when we reach the end of the NAR.
	// It adds the popped element to the element underneath if any,
	// and passes it to the directoryCb callback.
	// This function may only be called if the stack is not already empty.
	popFromStack := func() error {
		// Keep the top item, and "resize" the stack slice.
		// This will only make the last element unaccessible, but chances are high
		// we're re-using that space anyways.
		toPop := stack[len(stack)-1]
		stack = stack[:len(stack)-1]

		// if there's still a parent left on the stack, refer to it from there.
		if len(stack) > 0 {
			dgst, err := toPop.directory.Digest()
			if err != nil {
				return fmt.Errorf("unable to calculate directory digest: %w", err)
			}

			topOfStack := stack[len(stack)-1].directory
			topOfStack.Directories = append(topOfStack.Directories, &castorev1pb.DirectoryNode{
				Name:   []byte(path.Base(toPop.path)),
				Digest: dgst,
				Size:   toPop.directory.Size(),
			})
		}
		// call the directoryCb
		if err := directoryCb(toPop.directory); err != nil {
			return fmt.Errorf("failed calling directoryCb: %w", err)
		}
		// Keep track that we have encounter at least one directory
		stackDirectory = toPop.directory
		return nil
	}

	getBasename := func(p string) string {
		// extract the basename. In case of "/", replace with empty string.
		basename := path.Base(p)
		if basename == "/" {
			basename = ""
		}
		return basename
	}

	for {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
			// call narReader.Next() to get the next element
			hdr, err := narReader.Next()

			// If this returns an error, it's either EOF (when we're done reading from the NAR),
			// or another error
			if err != nil {
				// if this returns no EOF, bail out
				if !errors.Is(err, io.EOF) {
					return nil, fmt.Errorf("failed getting next nar element: %w", err)
				}

				// The NAR has been read all the way to the end…
				// Make sure we close the nar reader, which might read some final trailers.
				if err := narReader.Close(); err != nil {
					return nil, fmt.Errorf("unable to close nar reader: %w", err)
				}

				// Check the stack. While it's not empty, we need to pop things off the stack.
				for len(stack) > 0 {
					err := popFromStack()
					if err != nil {
						return nil, fmt.Errorf("unable to pop from stack: %w", err)
					}

				}

				// Stack is empty. We now either have a regular or symlink root node,
				// or we encountered at least one directory assemble pathInfo with these and
				// return.
				pi := &storev1pb.PathInfo{
					Node:       nil,
					References: [][]byte{},
					Narinfo: &storev1pb.NARInfo{
						NarSize:        uint64(r.hrSha256.BytesWritten()),
						NarSha256:      r.hrSha256.Sum(nil),
						Signatures:     []*storev1pb.NARInfo_Signature{},
						ReferenceNames: []string{},
					},
				}

				if rootFile != nil {
					pi.Node = &castorev1pb.Node{
						Node: &castorev1pb.Node_File{
							File: rootFile,
						},
					}
				}
				if rootSymlink != nil {
					pi.Node = &castorev1pb.Node{
						Node: &castorev1pb.Node_Symlink{
							Symlink: rootSymlink,
						},
					}
				}
				if stackDirectory != nil {
					// calculate directory digest (i.e. after we received all its contents)
					dgst, err := stackDirectory.Digest()
					if err != nil {
						return nil, fmt.Errorf("unable to calculate root directory digest: %w", err)
					}

					pi.Node = &castorev1pb.Node{
						Node: &castorev1pb.Node_Directory{
							Directory: &castorev1pb.DirectoryNode{
								Name:   []byte{},
								Digest: dgst,
								Size:   stackDirectory.Size(),
							},
						},
					}
				}
				return pi, nil
			}

			// Check for valid path transitions, pop from stack if needed
			// The nar reader already gives us some guarantees about ordering and illegal transitions,
			// So we really only need to check if the top-of-stack path is a prefix of the path,
			// and if it's not, pop from the stack. We do this repeatedly until the top of the stack is
			// the subdirectory the new entry is in, or we hit the root directory.

			// We don't need to worry about the root node case, because we can only finish the root "/"
			// If we're at the end of the NAR reader (covered by the EOF check)
			for len(stack) > 1 && !strings.HasPrefix(hdr.Path, stack[len(stack)-1].path+"/") {
				err := popFromStack()
				if err != nil {
					return nil, fmt.Errorf("unable to pop from stack: %w", err)
				}
			}

			if hdr.Type == nar.TypeSymlink {
				symlinkNode := &castorev1pb.SymlinkNode{
					Name:   []byte(getBasename(hdr.Path)),
					Target: []byte(hdr.LinkTarget),
				}
				if len(stack) > 0 {
					topOfStack := stack[len(stack)-1].directory
					topOfStack.Symlinks = append(topOfStack.Symlinks, symlinkNode)
				} else {
					rootSymlink = symlinkNode
				}

			}
			if hdr.Type == nar.TypeRegular {
				// wrap reader with a reader calculating the blake3 hash
				fileReader := NewHasher(narReader, blake3.New(32, nil))

				err := blobCb(fileReader)
				if err != nil {
					return nil, fmt.Errorf("failure from blobCb: %w", err)
				}

				// drive the file reader to the end, in case the CB function doesn't read
				// all the way to the end on its own
				if fileReader.BytesWritten() != uint32(hdr.Size) {
					_, err := io.ReadAll(fileReader)
					if err != nil {
						return nil, fmt.Errorf("unable to read until the end of the file content: %w", err)
					}
				}

				// read the blake3 hash
				dgst := fileReader.Sum(nil)

				fileNode := &castorev1pb.FileNode{
					Name:       []byte(getBasename(hdr.Path)),
					Digest:     dgst,
					Size:       uint32(hdr.Size),
					Executable: hdr.Executable,
				}
				if len(stack) > 0 {
					topOfStack := stack[len(stack)-1].directory
					topOfStack.Files = append(topOfStack.Files, fileNode)
				} else {
					rootFile = fileNode
				}
			}
			if hdr.Type == nar.TypeDirectory {
				directory := &castorev1pb.Directory{
					Directories: []*castorev1pb.DirectoryNode{},
					Files:       []*castorev1pb.FileNode{},
					Symlinks:    []*castorev1pb.SymlinkNode{},
				}
				stack = append(stack, stackItem{
					directory: directory,
					path:      hdr.Path,
				})
			}
		}
	}
}