1 files changed, 45 insertions, 8 deletions
diff --git a/pkg/sentry/mm/lifecycle.go b/pkg/sentry/mm/lifecycle.go
index 7646d5ab2..4e9ca1de6 100644
--- a/pkg/sentry/mm/lifecycle.go
+++ b/pkg/sentry/mm/lifecycle.go
@@ -18,14 +18,14 @@ import (
 	"fmt"
 	"sync/atomic"
 
-	"gvisor.googlesource.com/gvisor/pkg/atomicbitops"
-	"gvisor.googlesource.com/gvisor/pkg/sentry/arch"
-	"gvisor.googlesource.com/gvisor/pkg/sentry/context"
-	"gvisor.googlesource.com/gvisor/pkg/sentry/limits"
-	"gvisor.googlesource.com/gvisor/pkg/sentry/memmap"
-	"gvisor.googlesource.com/gvisor/pkg/sentry/pgalloc"
-	"gvisor.googlesource.com/gvisor/pkg/sentry/platform"
-	"gvisor.googlesource.com/gvisor/pkg/sentry/usermem"
+	"gvisor.dev/gvisor/pkg/atomicbitops"
+	"gvisor.dev/gvisor/pkg/sentry/arch"
+	"gvisor.dev/gvisor/pkg/sentry/context"
+	"gvisor.dev/gvisor/pkg/sentry/limits"
+	"gvisor.dev/gvisor/pkg/sentry/memmap"
+	"gvisor.dev/gvisor/pkg/sentry/pgalloc"
+	"gvisor.dev/gvisor/pkg/sentry/platform"
+	"gvisor.dev/gvisor/pkg/sentry/usermem"
 )
 
 // NewMemoryManager returns a new MemoryManager with no mappings and 1 user.
@@ -86,10 +86,22 @@ func (mm *MemoryManager) Fork(ctx context.Context) (*MemoryManager, error) {
 	}
 
 	// Copy vmas.
+	dontforks := false
 	dstvgap := mm2.vmas.FirstGap()
 	for srcvseg := mm.vmas.FirstSegment(); srcvseg.Ok(); srcvseg = srcvseg.NextSegment() {
 		vma := srcvseg.Value() // makes a copy of the vma
 		vmaAR := srcvseg.Range()
+
+		if vma.dontfork {
+			length := uint64(vmaAR.Length())
+			mm2.usageAS -= length
+			if vma.isPrivateDataLocked() {
+				mm2.dataAS -= length
+			}
+			dontforks = true
+			continue
+		}
+
 		// Inform the Mappable, if any, of the new mapping.
 		if vma.mappable != nil {
 			if err := vma.mappable.AddMapping(ctx, mm2, vmaAR, vma.off, vma.canWriteMappableLocked()); err != nil {
@@ -118,6 +130,10 @@ func (mm *MemoryManager) Fork(ctx context.Context) (*MemoryManager, error) {
 	defer mm2.activeMu.Unlock()
 	mm.activeMu.Lock()
 	defer mm.activeMu.Unlock()
+	if dontforks {
+		defer mm.pmas.MergeRange(mm.applicationAddrRange())
+	}
+	srcvseg := mm.vmas.FirstSegment()
 	dstpgap := mm2.pmas.FirstGap()
 	var unmapAR usermem.AddrRange
 	for srcpseg := mm.pmas.FirstSegment(); srcpseg.Ok(); srcpseg = srcpseg.NextSegment() {
@@ -125,6 +141,27 @@ func (mm *MemoryManager) Fork(ctx context.Context) (*MemoryManager, error) {
 		if !pma.private {
 			continue
 		}
+
+		if dontforks {
+			// Find the 'vma' that contains the starting address
+			// associated with the 'pma' (there must be one).
+			srcvseg = srcvseg.seekNextLowerBound(srcpseg.Start())
+			if checkInvariants {
+				if !srcvseg.Ok() {
+					panic(fmt.Sprintf("no vma covers pma range %v", srcpseg.Range()))
+				}
+				if srcpseg.Start() < srcvseg.Start() {
+					panic(fmt.Sprintf("vma %v ran ahead of pma %v", srcvseg.Range(), srcpseg.Range()))
+				}
+			}
+
+			srcpseg = mm.pmas.Isolate(srcpseg, srcvseg.Range())
+			if srcvseg.ValuePtr().dontfork {
+				continue
+			}
+			pma = srcpseg.ValuePtr()
+		}
+
 		if !pma.needCOW {
 			pma.needCOW = true
 			if pma.effectivePerms.Write {