restore: track expected filenames in a folder

internal/restorer/restorer.go

@@ -37,6 +37,7 @@ type Options struct {
 	Sparse    bool
 	Progress  *restoreui.Progress
 	Overwrite OverwriteBehavior
+	Delete    bool
 }
 
 type OverwriteBehavior int
@@ -107,22 +108,61 @@ func NewRestorer(repo restic.Repository, sn *restic.Snapshot, opts Options) *Res
 type treeVisitor struct {
 	enterDir  func(node *restic.Node, target, location string) error
 	visitNode func(node *restic.Node, target, location string) error
-	leaveDir  func(node *restic.Node, target, location string) error
+	// 'entries' contains all files the snapshot contains for this node. This also includes files
+	// ignored by the SelectFilter.
+	leaveDir func(node *restic.Node, target, location string, entries []string) error
 }
 
 // traverseTree traverses a tree from the repo and calls treeVisitor.
 // target is the path in the file system, location within the snapshot.
-func (res *Restorer) traverseTree(ctx context.Context, target, location string, treeID restic.ID, visitor treeVisitor) (hasRestored bool, err error) {
+func (res *Restorer) traverseTree(ctx context.Context, target string, treeID restic.ID, visitor treeVisitor) error {
+	location := string(filepath.Separator)
+	sanitizeError := func(err error) error {
+		switch err {
+		case nil, context.Canceled, context.DeadlineExceeded:
+			// Context errors are permanent.
+			return err
+		default:
+			return res.Error(location, err)
+		}
+	}
+
+	if visitor.enterDir != nil {
+		err := sanitizeError(visitor.enterDir(nil, target, location))
+		if err != nil {
+			return err
+		}
+	}
+	childFilenames, hasRestored, err := res.traverseTreeInner(ctx, target, location, treeID, visitor)
+	if err != nil {
+		return err
+	}
+	if hasRestored && visitor.leaveDir != nil {
+		err = sanitizeError(visitor.leaveDir(nil, target, location, childFilenames))
+	}
+
+	return err
+}
+
+func (res *Restorer) traverseTreeInner(ctx context.Context, target, location string, treeID restic.ID, visitor treeVisitor) (filenames []string, hasRestored bool, err error) {
 	debug.Log("%v %v %v", target, location, treeID)
 	tree, err := restic.LoadTree(ctx, res.repo, treeID)
 	if err != nil {
 		debug.Log("error loading tree %v: %v", treeID, err)
-		return hasRestored, res.Error(location, err)
+		return nil, hasRestored, res.Error(location, err)
 	}
 
+	if res.opts.Delete {
+		filenames = make([]string, 0, len(tree.Nodes))
+	}
 	for i, node := range tree.Nodes {
 		// allow GC of tree node
 		tree.Nodes[i] = nil
+		if res.opts.Delete {
+			// just track all files included in the tree node to simplify the control flow.
+			// tracking too many files does not matter except for a slightly elevated memory usage
+			filenames = append(filenames, node.Name)
+		}
 
 		// ensure that the node name does not contain anything that refers to a
 		// top-level directory.
@@ -131,8 +171,10 @@ func (res *Restorer) traverseTree(ctx context.Context, target, location string,
 			debug.Log("node %q has invalid name %q", node.Name, nodeName)
 			err := res.Error(location, errors.Errorf("invalid child node name %s", node.Name))
 			if err != nil {
-				return hasRestored, err
+				return nil, hasRestored, err
 			}
+			// force disable deletion to prevent unexpected behavior
+			res.opts.Delete = false
 			continue
 		}
 
@@ -144,8 +186,10 @@ func (res *Restorer) traverseTree(ctx context.Context, target, location string,
 			debug.Log("node %q has invalid target path %q", node.Name, nodeTarget)
 			err := res.Error(nodeLocation, errors.New("node has invalid path"))
 			if err != nil {
-				return hasRestored, err
+				return nil, hasRestored, err
 			}
+			// force disable deletion to prevent unexpected behavior
+			res.opts.Delete = false
 			continue
 		}
 
@@ -173,25 +217,26 @@ func (res *Restorer) traverseTree(ctx context.Context, target, location string,
 
 		if node.Type == "dir" {
 			if node.Subtree == nil {
-				return hasRestored, errors.Errorf("Dir without subtree in tree %v", treeID.Str())
+				return nil, hasRestored, errors.Errorf("Dir without subtree in tree %v", treeID.Str())
 			}
 
 			if selectedForRestore && visitor.enterDir != nil {
 				err = sanitizeError(visitor.enterDir(node, nodeTarget, nodeLocation))
 				if err != nil {
-					return hasRestored, err
+					return nil, hasRestored, err
 				}
 			}
 
 			// keep track of restored child status
 			// so metadata of the current directory are restored on leaveDir
 			childHasRestored := false
+			var childFilenames []string
 
 			if childMayBeSelected {
-				childHasRestored, err = res.traverseTree(ctx, nodeTarget, nodeLocation, *node.Subtree, visitor)
+				childFilenames, childHasRestored, err = res.traverseTreeInner(ctx, nodeTarget, nodeLocation, *node.Subtree, visitor)
 				err = sanitizeError(err)
 				if err != nil {
-					return hasRestored, err
+					return nil, hasRestored, err
 				}
 				// inform the parent directory to restore parent metadata on leaveDir if needed
 				if childHasRestored {
@@ -202,9 +247,9 @@ func (res *Restorer) traverseTree(ctx context.Context, target, location string,
 			// metadata need to be restore when leaving the directory in both cases
 			// selected for restore or any child of any subtree have been restored
 			if (selectedForRestore || childHasRestored) && visitor.leaveDir != nil {
-				err = sanitizeError(visitor.leaveDir(node, nodeTarget, nodeLocation))
+				err = sanitizeError(visitor.leaveDir(node, nodeTarget, nodeLocation, childFilenames))
 				if err != nil {
-					return hasRestored, err
+					return nil, hasRestored, err
 				}
 			}
 
@@ -214,12 +259,12 @@ func (res *Restorer) traverseTree(ctx context.Context, target, location string,
 		if selectedForRestore {
 			err = sanitizeError(visitor.visitNode(node, nodeTarget, nodeLocation))
 			if err != nil {
-				return hasRestored, err
+				return nil, hasRestored, err
 			}
 		}
 	}
 
-	return hasRestored, nil
+	return filenames, hasRestored, nil
 }
 
 func (res *Restorer) restoreNodeTo(ctx context.Context, node *restic.Node, target, location string) error {
@@ -310,10 +355,12 @@ func (res *Restorer) RestoreTo(ctx context.Context, dst string) error {
 	var buf []byte
 
 	// first tree pass: create directories and collect all files to restore
-	_, err = res.traverseTree(ctx, dst, string(filepath.Separator), *res.sn.Tree, treeVisitor{
+	err = res.traverseTree(ctx, dst, *res.sn.Tree, treeVisitor{
 		enterDir: func(_ *restic.Node, target, location string) error {
 			debug.Log("first pass, enterDir: mkdir %q, leaveDir should restore metadata", location)
+			if location != "/" {
 				res.opts.Progress.AddFile(0)
+			}
 			return res.ensureDir(target)
 		},
 
@@ -373,7 +420,7 @@ func (res *Restorer) RestoreTo(ctx context.Context, dst string) error {
 	debug.Log("second pass for %q", dst)
 
 	// second tree pass: restore special files and filesystem metadata
-	_, err = res.traverseTree(ctx, dst, string(filepath.Separator), *res.sn.Tree, treeVisitor{
+	err = res.traverseTree(ctx, dst, *res.sn.Tree, treeVisitor{
 		visitNode: func(node *restic.Node, target, location string) error {
 			debug.Log("second pass, visitNode: restore node %q", location)
 			if node.Type != "file" {
@@ -396,7 +443,11 @@ func (res *Restorer) RestoreTo(ctx context.Context, dst string) error {
 			// don't touch skipped files
 			return nil
 		},
-		leaveDir: func(node *restic.Node, target, location string) error {
+		leaveDir: func(node *restic.Node, target, location string, expectedFilenames []string) error {
+			if node == nil {
+				return nil
+			}
+
 			err := res.restoreNodeMetadataTo(node, target, location)
 			if err == nil {
 				res.opts.Progress.AddProgress(location, restoreui.ActionDirRestored, 0, 0)
@@ -493,7 +544,7 @@ func (res *Restorer) VerifyFiles(ctx context.Context, dst string) (int, error) {
 	g.Go(func() error {
 		defer close(work)
 
-		_, err := res.traverseTree(ctx, dst, string(filepath.Separator), *res.sn.Tree, treeVisitor{
+		err := res.traverseTree(ctx, dst, *res.sn.Tree, treeVisitor{
 			visitNode: func(node *restic.Node, target, location string) error {
 				if node.Type != "file" {
 					return nil

internal/restorer/restorer_test.go

@@ -8,6 +8,7 @@ import (
 	"math"
 	"os"
 	"path/filepath"
+	"reflect"
 	"runtime"
 	"strings"
 	"syscall"
@@ -529,14 +530,15 @@ type TraverseTreeCheck func(testing.TB) treeVisitor
 type TreeVisit struct {
 	funcName string   // name of the function
 	location string   // location passed to the function
+	files    []string // file list passed to the function
 }
 
 func checkVisitOrder(list []TreeVisit) TraverseTreeCheck {
 	var pos int
 
 	return func(t testing.TB) treeVisitor {
-		check := func(funcName string) func(*restic.Node, string, string) error {
+		check := func(funcName string) func(*restic.Node, string, string, []string) error {
-			return func(node *restic.Node, target, location string) error {
+			return func(node *restic.Node, target, location string, expectedFilenames []string) error {
 				if pos >= len(list) {
 					t.Errorf("step %v, %v(%v): expected no more than %d function calls", pos, funcName, location, len(list))
 					pos++
@@ -554,14 +556,24 @@ func checkVisitOrder(list []TreeVisit) TraverseTreeCheck {
 					t.Errorf("step %v: want location %v, got %v", pos, list[pos].location, location)
 				}
 
+				if !reflect.DeepEqual(expectedFilenames, v.files) {
+					t.Errorf("step %v: want files %v, got %v", pos, list[pos].files, expectedFilenames)
+				}
+
 				pos++
 				return nil
 			}
 		}
+		checkNoFilename := func(funcName string) func(*restic.Node, string, string) error {
+			f := check(funcName)
+			return func(node *restic.Node, target, location string) error {
+				return f(node, target, location, nil)
+			}
+		}
 
 		return treeVisitor{
-			enterDir:  check("enterDir"),
+			enterDir:  checkNoFilename("enterDir"),
-			visitNode: check("visitNode"),
+			visitNode: checkNoFilename("visitNode"),
 			leaveDir:  check("leaveDir"),
 		}
 	}
@@ -590,13 +602,15 @@ func TestRestorerTraverseTree(t *testing.T) {
 				return true, true
 			},
 			Visitor: checkVisitOrder([]TreeVisit{
-				{"enterDir", "/dir"},
+				{"enterDir", "/", nil},
-				{"visitNode", "/dir/otherfile"},
+				{"enterDir", "/dir", nil},
-				{"enterDir", "/dir/subdir"},
+				{"visitNode", "/dir/otherfile", nil},
-				{"visitNode", "/dir/subdir/file"},
+				{"enterDir", "/dir/subdir", nil},
-				{"leaveDir", "/dir/subdir"},
+				{"visitNode", "/dir/subdir/file", nil},
-				{"leaveDir", "/dir"},
+				{"leaveDir", "/dir/subdir", []string{"file"}},
-				{"visitNode", "/foo"},
+				{"leaveDir", "/dir", []string{"otherfile", "subdir"}},
+				{"visitNode", "/foo", nil},
+				{"leaveDir", "/", []string{"dir", "foo"}},
 			}),
 		},
 
@@ -620,7 +634,9 @@ func TestRestorerTraverseTree(t *testing.T) {
 				return false, false
 			},
 			Visitor: checkVisitOrder([]TreeVisit{
-				{"visitNode", "/foo"},
+				{"enterDir", "/", nil},
+				{"visitNode", "/foo", nil},
+				{"leaveDir", "/", []string{"dir", "foo"}},
 			}),
 		},
 		{
@@ -642,7 +658,9 @@ func TestRestorerTraverseTree(t *testing.T) {
 				return false, false
 			},
 			Visitor: checkVisitOrder([]TreeVisit{
-				{"visitNode", "/aaa"},
+				{"enterDir", "/", nil},
+				{"visitNode", "/aaa", nil},
+				{"leaveDir", "/", []string{"aaa", "dir"}},
 			}),
 		},
 
@@ -666,12 +684,14 @@ func TestRestorerTraverseTree(t *testing.T) {
 				return false, false
 			},
 			Visitor: checkVisitOrder([]TreeVisit{
-				{"enterDir", "/dir"},
+				{"enterDir", "/", nil},
-				{"visitNode", "/dir/otherfile"},
+				{"enterDir", "/dir", nil},
-				{"enterDir", "/dir/subdir"},
+				{"visitNode", "/dir/otherfile", nil},
-				{"visitNode", "/dir/subdir/file"},
+				{"enterDir", "/dir/subdir", nil},
-				{"leaveDir", "/dir/subdir"},
+				{"visitNode", "/dir/subdir/file", nil},
-				{"leaveDir", "/dir"},
+				{"leaveDir", "/dir/subdir", []string{"file"}},
+				{"leaveDir", "/dir", []string{"otherfile", "subdir"}},
+				{"leaveDir", "/", []string{"dir", "foo"}},
 			}),
 		},
 
@@ -699,8 +719,10 @@ func TestRestorerTraverseTree(t *testing.T) {
 				}
 			},
 			Visitor: checkVisitOrder([]TreeVisit{
-				{"visitNode", "/dir/otherfile"},
+				{"enterDir", "/", nil},
-				{"leaveDir", "/dir"},
+				{"visitNode", "/dir/otherfile", nil},
+				{"leaveDir", "/dir", []string{"otherfile", "subdir"}},
+				{"leaveDir", "/", []string{"dir", "foo"}},
 			}),
 		},
 	}
@@ -710,7 +732,8 @@ func TestRestorerTraverseTree(t *testing.T) {
 			repo := repository.TestRepository(t)
 			sn, _ := saveSnapshot(t, repo, test.Snapshot, noopGetGenericAttributes)
 
-			res := NewRestorer(repo, sn, Options{})
+			// set Delete option to enable tracking filenames in a directory
+			res := NewRestorer(repo, sn, Options{Delete: true})
 
 			res.SelectFilter = test.Select
 
@@ -721,7 +744,7 @@ func TestRestorerTraverseTree(t *testing.T) {
 			// make sure we're creating a new subdir of the tempdir
 			target := filepath.Join(tempdir, "target")
 
-			_, err := res.traverseTree(ctx, target, string(filepath.Separator), *sn.Tree, test.Visitor(t))
+			err := res.traverseTree(ctx, target, *sn.Tree, test.Visitor(t))
 			if err != nil {
 				t.Fatal(err)
 			}