Check for fatal signals and reset watchdog during MM.mapASLocked().

PiperOrigin-RevId: 777697337

Author: nixprime

pkg/context/context.go

@@ -43,6 +43,9 @@ type Blocker interface {
 	// Interrupted notes whether this context is Interrupted.
 	Interrupted() bool
 
+	// Killed returns true if this context is interrupted by a fatal signal.
+	Killed() bool
+
 	// BlockOn blocks until one of the previously registered events occurs,
 	// or some external interrupt (cancellation).
 	//
@@ -94,6 +97,11 @@ func (nt *NoTask) Interrupted() bool {
 	return nt.cancel != nil && len(nt.cancel) > 0
 }
 
+// Killed implements Blocker.Killed.
+func (nt *NoTask) Killed() bool {
+	return false
+}
+
 // Block implements Blocker.Block.
 func (nt *NoTask) Block(C <-chan struct{}) error {
 	if nt.cancel == nil {

pkg/sentry/kernel/task_block.go

@@ -234,7 +234,7 @@ func (t *Task) Interrupted() bool {
 	}
 	// Indicate that t's task goroutine is still responsive (i.e. reset the
 	// watchdog timer).
-	t.accountTaskGoroutineRunning()
+	t.touchGostateTime()
 	return false
 }
 

pkg/sentry/kernel/task_exit.go

@@ -100,6 +100,17 @@ func (t *Task) killLocked() {
 	t.interrupt()
 }
 
+// Killed implements context.Blocker.Killed.
+func (t *Task) Killed() bool {
+	if t.killed() {
+		return true
+	}
+	// Indicate that t's task goroutine is still responsive (i.e. reset the
+	// watchdog timer).
+	t.touchGostateTime()
+	return false
+}
+
 // killed returns true if t has a SIGKILL pending. killed is analogous to
 // Linux's fatal_signal_pending().
 //

pkg/sentry/kernel/task_sched.go

@@ -116,14 +116,6 @@ func (t *Task) accountTaskGoroutineLeave(state TaskGoroutineState) {
 	t.gostateSeq.EndWrite()
 }
 
-// Preconditions: The caller must be running on the task goroutine.
-func (t *Task) accountTaskGoroutineRunning() {
-	if oldState := t.TaskGoroutineState(); oldState != TaskGoroutineRunningSys {
-		panic(fmt.Sprintf("Task goroutine in state %v (expected %v)", oldState, TaskGoroutineRunningSys))
-	}
-	t.touchGostateTime()
-}
-
 // Preconditions: The caller must be running on the task goroutine.
 func (t *Task) touchGostateTime() {
 	t.gostateTime.Store(t.k.cpuClock.Load())

pkg/sentry/mm/address_space.go

@@ -18,6 +18,7 @@ import (
 	"fmt"
 
 	"gvisor.dev/gvisor/pkg/context"
+	"gvisor.dev/gvisor/pkg/errors/linuxerr"
 	"gvisor.dev/gvisor/pkg/hostarch"
 	"gvisor.dev/gvisor/pkg/sentry/memmap"
 	"gvisor.dev/gvisor/pkg/sentry/pgalloc"
@@ -172,7 +173,7 @@ func (mm *MemoryManager) Deactivate() {
 //   - ar.Length() != 0.
 //   - ar must be page-aligned.
 //   - pseg == mm.pmas.LowerBoundSegment(ar.Start).
-func (mm *MemoryManager) mapASLocked(pseg pmaIterator, ar hostarch.AddrRange, platformEffect memmap.MMapPlatformEffect) error {
+func (mm *MemoryManager) mapASLocked(ctx context.Context, pseg pmaIterator, ar hostarch.AddrRange, platformEffect memmap.MMapPlatformEffect) error {
 	// By default, map entire pmas at a time, under the assumption that there
 	// is no cost to mapping more of a pma than necessary.
 	mapAR := hostarch.AddrRange{0, ^hostarch.Addr(hostarch.PageSize - 1)}
@@ -217,8 +218,28 @@ func (mm *MemoryManager) mapASLocked(pseg pmaIterator, ar hostarch.AddrRange, pl
 			perms.Write = false
 		}
 		if perms.Any() { // MapFile precondition
-			if err := mm.as.MapFile(pmaMapAR.Start, pma.file, pseg.fileRangeOf(pmaMapAR), perms, platformEffect == memmap.PlatformEffectCommit); err != nil {
-				return err
+			// If the length of the mapping exceeds singleMapThreshold, call
+			// AddressSpace.MapFile() on singleMapThreshold-aligned chunks so
+			// we can check ctx.Killed() reasonably frequently.
+			const singleMapThreshold = 1 << 30
+			if pmaMapAR.Length() <= singleMapThreshold {
+				if err := mm.as.MapFile(pmaMapAR.Start, pma.file, pseg.fileRangeOf(pmaMapAR), perms, platformEffect == memmap.PlatformEffectCommit); err != nil {
+					return err
+				}
+				if ctx.Killed() {
+					return linuxerr.EINTR
+				}
+			} else {
+				for windowStart := pmaMapAR.Start &^ (singleMapThreshold - 1); windowStart < pmaMapAR.End; windowStart += singleMapThreshold {
+					windowAR := hostarch.AddrRange{windowStart, windowStart + singleMapThreshold}
+					thisMapAR := pmaMapAR.Intersect(windowAR)
+					if err := mm.as.MapFile(thisMapAR.Start, pma.file, pseg.fileRangeOf(thisMapAR), perms, platformEffect == memmap.PlatformEffectCommit); err != nil {
+						return err
+					}
+					if ctx.Killed() {
+						return linuxerr.EINTR
+					}
+				}
 			}
 		}
 		pseg = pseg.NextSegment()

pkg/sentry/mm/io.go

@@ -514,7 +514,7 @@ func (mm *MemoryManager) handleASIOFault(ctx context.Context, addr hostarch.Addr
 	// anymore.
 	mm.activeMu.DowngradeLock()
 
-	err = mm.mapASLocked(pseg, ar, memmap.PlatformEffectDefault)
+	err = mm.mapASLocked(ctx, pseg, ar, memmap.PlatformEffectDefault)
 	mm.activeMu.RUnlock()
 	return translateIOError(ctx, err)
 }

pkg/sentry/mm/syscalls.go

@@ -66,7 +66,7 @@ func (mm *MemoryManager) HandleUserFault(ctx context.Context, addr hostarch.Addr
 	mm.activeMu.DowngradeLock()
 
 	// Map the faulted page into the active AddressSpace.
-	err = mm.mapASLocked(pseg, ar, memmap.PlatformEffectDefault)
+	err = mm.mapASLocked(ctx, pseg, ar, memmap.PlatformEffectDefault)
 	mm.activeMu.RUnlock()
 	return err
 }
@@ -201,7 +201,7 @@ func (mm *MemoryManager) populateVMA(ctx context.Context, vseg vmaIterator, ar h
 	// Downgrade to a read-lock on activeMu since we don't need to mutate pmas
 	// anymore.
 	mm.activeMu.DowngradeLock()
-	err = mm.mapASLocked(pseg, ar, platformEffect)
+	err = mm.mapASLocked(ctx, pseg, ar, platformEffect)
 	mm.activeMu.RUnlock()
 	return err
 }
@@ -250,7 +250,7 @@ func (mm *MemoryManager) populateVMAAndUnlock(ctx context.Context, vseg vmaItera
 
 	// As above, errors are silently ignored.
 	mm.activeMu.DowngradeLock()
-	mm.mapASLocked(pseg, ar, platformEffect)
+	mm.mapASLocked(ctx, pseg, ar, platformEffect)
 	mm.activeMu.RUnlock()
 }
 
@@ -930,7 +930,7 @@ func (mm *MemoryManager) MLock(ctx context.Context, addr hostarch.Addr, length u
 		mm.mappingMu.RUnlock()
 		if mm.as != nil {
 			mm.activeMu.DowngradeLock()
-			err := mm.mapASLocked(mm.pmas.LowerBoundSegment(ar.Start), ar, memmap.PlatformEffectCommit)
+			err := mm.mapASLocked(ctx, mm.pmas.LowerBoundSegment(ar.Start), ar, memmap.PlatformEffectCommit)
 			mm.activeMu.RUnlock()
 			if err != nil {
 				return err
@@ -1014,7 +1014,7 @@ func (mm *MemoryManager) MLockAll(ctx context.Context, opts MLockAllOpts) error
 		mm.mappingMu.RUnlock()
 		if mm.as != nil {
 			mm.activeMu.DowngradeLock()
-			mm.mapASLocked(mm.pmas.FirstSegment(), mm.applicationAddrRange(), memmap.PlatformEffectCommit)
+			mm.mapASLocked(ctx, mm.pmas.FirstSegment(), mm.applicationAddrRange(), memmap.PlatformEffectCommit)
 			mm.activeMu.RUnlock()
 		} else {
 			mm.activeMu.Unlock()