From cf93d156901b040a50c1baea2cc32b68089bd3e0 Mon Sep 17 00:00:00 2001
From: Andrei Vagin <avagin@google.com>
Date: Wed, 4 Oct 2023 12:56:25 -0700
Subject: [PATCH] mount: update root and cwd from cloneMountTree

In this case, we don't need to enumerate all tree mounts and look up pair
mounts in the cloned tree.

PiperOrigin-RevId: 570778945
---
 pkg/sentry/vfs/mount.go     | 14 +++++++++++++-
 pkg/sentry/vfs/namespace.go | 10 ++++------
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/pkg/sentry/vfs/mount.go b/pkg/sentry/vfs/mount.go
index 3f804b65b..8ad994636 100644
--- a/pkg/sentry/vfs/mount.go
+++ b/pkg/sentry/vfs/mount.go
@@ -377,10 +377,19 @@ type cloneTreeNode struct {
 
 // cloneMountTree creates a copy of mnt's tree with the specified root
 // dentry at root. The new descendents are added to mnt's pending mount list.
+// `cloneFunc` is a callback that is executed for each cloned mount.
 //
 // +checklocks:vfs.mountMu
-func (vfs *VirtualFilesystem) cloneMountTree(ctx context.Context, mnt *Mount, root *Dentry) (*Mount, error) {
+func (vfs *VirtualFilesystem) cloneMountTree(
+	ctx context.Context,
+	mnt *Mount,
+	root *Dentry,
+	cloneFunc func(ctx context.Context, oldmnt, newMnt *Mount),
+) (*Mount, error) {
 	clone := vfs.cloneMount(mnt, root, nil)
+	if cloneFunc != nil {
+		cloneFunc(ctx, mnt, clone)
+	}
 	queue := []cloneTreeNode{{mnt, clone}}
 	for len(queue) != 0 {
 		p := queue[len(queue)-1]
@@ -397,6 +406,9 @@ func (vfs *VirtualFilesystem) cloneMountTree(ctx context.Context, mnt *Mount, ro
 			if len(c.children) != 0 {
 				queue = append(queue, cloneTreeNode{c, m})
 			}
+			if cloneFunc != nil {
+				cloneFunc(ctx, c, m)
+			}
 		}
 	}
 	return clone, nil
diff --git a/pkg/sentry/vfs/namespace.go b/pkg/sentry/vfs/namespace.go
index 4cb39a89b..8e275e870 100644
--- a/pkg/sentry/vfs/namespace.go
+++ b/pkg/sentry/vfs/namespace.go
@@ -142,10 +142,6 @@ func (vfs *VirtualFilesystem) updateRootAndCWD(ctx context.Context, root *Virtua
 		cwd.mount = dst
 		cwd.mount.IncRef()
 	}
-	for srcChild := range src.children {
-		dstChild := vfs.mounts.Lookup(dst, srcChild.point())
-		vfs.updateRootAndCWD(ctx, root, cwd, srcChild, dstChild)
-	}
 }
 
 // NamespaceInodeGetter is an interface that provides the GetNamespaceInode method.
@@ -174,7 +170,10 @@ func (vfs *VirtualFilesystem) CloneMountNamespace(
 	vfs.lockMounts()
 	defer vfs.unlockMounts(ctx)
 
-	newRoot, err := vfs.cloneMountTree(ctx, ns.root, ns.root.root)
+	newRoot, err := vfs.cloneMountTree(ctx, ns.root, ns.root.root,
+		func(ctx context.Context, src, dst *Mount) {
+			vfs.updateRootAndCWD(ctx, root, cwd, src, dst) // +checklocksforce: vfs.mountMu is locked.
+		})
 	if err != nil {
 		newns.DecRef(ctx)
 		vfs.abortTree(ctx, newRoot)
@@ -183,7 +182,6 @@ func (vfs *VirtualFilesystem) CloneMountNamespace(
 	newns.root = newRoot
 	newns.root.ns = newns
 	vfs.commitPendingTree(ctx, newRoot)
-	vfs.updateRootAndCWD(ctx, root, cwd, ns.root, newns.root)
 	return newns, nil
 }