Add MemoryManager.Pin.

PiperOrigin-RevId: 204162313
Change-Id: Ib0593dde88ac33e222c12d0dca6733ef1f1035dc

2 files changed, 92 insertions, 2 deletions

diff --git a/pkg/sentry/mm/pma.go b/pkg/sentry/mm/pma.go
index 35e873762..9febb25ac 100644
--- a/pkg/sentry/mm/pma.go
+++ b/pkg/sentry/mm/pma.go
@@ -578,6 +578,96 @@ func (mm *MemoryManager) invalidateLocked(ar usermem.AddrRange, invalidatePrivat
 	}
 }
 
+// Pin returns the platform.File ranges currently mapped by addresses in ar in
+// mm, acquiring a reference on the returned ranges which the caller must
+// release by calling Unpin. If not all addresses are mapped, Pin returns a
+// non-nil error. Note that Pin may return both a non-empty slice of
+// PinnedRanges and a non-nil error.
+//
+// Pin does not prevent mapped ranges from changing, making it unsuitable for
+// most I/O. It should only be used in contexts that would use get_user_pages()
+// in the Linux kernel.
+//
+// Preconditions: ar.Length() != 0. ar must be page-aligned.
+func (mm *MemoryManager) Pin(ctx context.Context, ar usermem.AddrRange, at usermem.AccessType, ignorePermissions bool) ([]PinnedRange, error) {
+	if checkInvariants {
+		if !ar.WellFormed() || ar.Length() <= 0 || !ar.IsPageAligned() {
+			panic(fmt.Sprintf("invalid ar: %v", ar))
+		}
+	}
+
+	// Ensure that we have usable vmas.
+	mm.mappingMu.RLock()
+	vseg, vend, verr := mm.getVMAsLocked(ctx, ar, at, ignorePermissions)
+	if vendaddr := vend.Start(); vendaddr < ar.End {
+		if vendaddr <= ar.Start {
+			mm.mappingMu.RUnlock()
+			return nil, verr
+		}
+		ar.End = vendaddr
+	}
+
+	// Ensure that we have usable pmas.
+	mm.activeMu.Lock()
+	pseg, pend, perr := mm.getPMAsLocked(ctx, vseg, ar, pmaOpts{
+		breakCOW: at.Write,
+	})
+	mm.mappingMu.RUnlock()
+	if pendaddr := pend.Start(); pendaddr < ar.End {
+		if pendaddr <= ar.Start {
+			mm.activeMu.Unlock()
+			return nil, perr
+		}
+		ar.End = pendaddr
+	}
+
+	// Gather pmas.
+	var prs []PinnedRange
+	for pseg.Ok() && pseg.Start() < ar.End {
+		psar := pseg.Range().Intersect(ar)
+		f := pseg.ValuePtr().file
+		fr := pseg.fileRangeOf(psar)
+		f.IncRef(fr)
+		prs = append(prs, PinnedRange{
+			Source: psar,
+			File:   f,
+			Offset: fr.Start,
+		})
+		pseg = pseg.NextSegment()
+	}
+	mm.activeMu.Unlock()
+
+	// Return the first error in order of progress through ar.
+	if perr != nil {
+		return prs, perr
+	}
+	return prs, verr
+}
+
+// PinnedRanges are returned by MemoryManager.Pin.
+type PinnedRange struct {
+	// Source is the corresponding range of addresses.
+	Source usermem.AddrRange
+
+	// File is the mapped file.
+	File platform.File
+
+	// Offset is the offset into File at which this PinnedRange begins.
+	Offset uint64
+}
+
+// FileRange returns the platform.File offsets mapped by pr.
+func (pr PinnedRange) FileRange() platform.FileRange {
+	return platform.FileRange{pr.Offset, pr.Offset + uint64(pr.Source.Length())}
+}
+
+// Unpin releases the reference held by prs.
+func Unpin(prs []PinnedRange) {
+	for i := range prs {
+		prs[i].File.DecRef(prs[i].FileRange())
+	}
+}
+
 // movePMAsLocked moves all pmas in oldAR to newAR.
 //
 // Preconditions: mm.activeMu must be locked for writing. oldAR.Length() != 0.
diff --git a/pkg/sentry/platform/platform.go b/pkg/sentry/platform/platform.go
index 1c385bc5a..f2fe163e8 100644
--- a/pkg/sentry/platform/platform.go
+++ b/pkg/sentry/platform/platform.go
@@ -305,8 +305,8 @@ type File interface {
 	MapInto(as AddressSpace, addr usermem.Addr, fr FileRange, at usermem.AccessType, precommit bool) error
 
 	// MapInternal returns a mapping of the given file offsets in the invoking
-	// process' address space for reading and writing. The lifetime of the
-	// returned mapping is implementation-defined.
+	// process' address space for reading and writing. The returned mapping is
+	// valid as long as a reference is held on the mapped range.
 	//
 	// Note that fr.Start and fr.End need not be page-aligned.
 	//