Deflake kvm_test.

PiperOrigin-RevId: 200439846
Change-Id: I9970fe0716cb02f0f41b754891d55db7e0729f56

3 files changed, 120 insertions, 71 deletions

diff --git a/pkg/sentry/platform/kvm/BUILD b/pkg/sentry/platform/kvm/BUILD
index 89d98c5c7..135861368 100644
--- a/pkg/sentry/platform/kvm/BUILD
+++ b/pkg/sentry/platform/kvm/BUILD
@@ -71,7 +71,6 @@ go_library(
 
 go_test(
     name = "kvm_test",
-    size = "small",
     srcs = [
         "kvm_test.go",
         "virtual_map_test.go",
diff --git a/pkg/sentry/platform/kvm/kvm_test.go b/pkg/sentry/platform/kvm/kvm_test.go
index 71c5c856e..180bf7bb0 100644
--- a/pkg/sentry/platform/kvm/kvm_test.go
+++ b/pkg/sentry/platform/kvm/kvm_test.go
@@ -157,7 +157,9 @@ func TestApplicationSyscall(t *testing.T) {
 			FloatingPointState: dummyFPState,
 			PageTables:         pt,
 			FullRestore:        true,
-		}); err != nil {
+		}); err == platform.ErrContextInterrupt {
+			return true // Retry.
+		} else if err != nil {
 			t.Errorf("application syscall with full restore failed: %v", err)
 		}
 		return false
@@ -167,7 +169,9 @@ func TestApplicationSyscall(t *testing.T) {
 			Registers:          regs,
 			FloatingPointState: dummyFPState,
 			PageTables:         pt,
-		}); err != nil {
+		}); err == platform.ErrContextInterrupt {
+			return true // Retry.
+		} else if err != nil {
 			t.Errorf("application syscall with partial restore failed: %v", err)
 		}
 		return false
@@ -182,7 +186,9 @@ func TestApplicationFault(t *testing.T) {
 			FloatingPointState: dummyFPState,
 			PageTables:         pt,
 			FullRestore:        true,
-		}); err != platform.ErrContextSignal || (si != nil && si.Signo != int32(syscall.SIGSEGV)) {
+		}); err == platform.ErrContextInterrupt {
+			return true // Retry.
+		} else if err != platform.ErrContextSignal || (si != nil && si.Signo != int32(syscall.SIGSEGV)) {
 			t.Errorf("application fault with full restore got (%v, %v), expected (%v, SIGSEGV)", err, si, platform.ErrContextSignal)
 		}
 		return false
@@ -193,7 +199,9 @@ func TestApplicationFault(t *testing.T) {
 			Registers:          regs,
 			FloatingPointState: dummyFPState,
 			PageTables:         pt,
-		}); err != platform.ErrContextSignal || (si != nil && si.Signo != int32(syscall.SIGSEGV)) {
+		}); err == platform.ErrContextInterrupt {
+			return true // Retry.
+		} else if err != platform.ErrContextSignal || (si != nil && si.Signo != int32(syscall.SIGSEGV)) {
 			t.Errorf("application fault with partial restore got (%v, %v), expected (%v, SIGSEGV)", err, si, platform.ErrContextSignal)
 		}
 		return false
@@ -203,15 +211,20 @@ func TestApplicationFault(t *testing.T) {
 func TestRegistersSyscall(t *testing.T) {
 	applicationTest(t, true, testutil.TwiddleRegsSyscall, func(c *vCPU, regs *syscall.PtraceRegs, pt *pagetables.PageTables) bool {
 		testutil.SetTestRegs(regs) // Fill values for all registers.
-		if _, _, err := c.SwitchToUser(ring0.SwitchOpts{
-			Registers:          regs,
-			FloatingPointState: dummyFPState,
-			PageTables:         pt,
-		}); err != nil {
-			t.Errorf("application register check with partial restore got unexpected error: %v", err)
-		}
-		if err := testutil.CheckTestRegs(regs, false); err != nil {
-			t.Errorf("application register check with partial restore failed: %v", err)
+		for {
+			if _, _, err := c.SwitchToUser(ring0.SwitchOpts{
+				Registers:          regs,
+				FloatingPointState: dummyFPState,
+				PageTables:         pt,
+			}); err == platform.ErrContextInterrupt {
+				continue // Retry.
+			} else if err != nil {
+				t.Errorf("application register check with partial restore got unexpected error: %v", err)
+			}
+			if err := testutil.CheckTestRegs(regs, false); err != nil {
+				t.Errorf("application register check with partial restore failed: %v", err)
+			}
+			break // Done.
 		}
 		return false
 	})
@@ -220,16 +233,21 @@ func TestRegistersSyscall(t *testing.T) {
 func TestRegistersFault(t *testing.T) {
 	applicationTest(t, true, testutil.TwiddleRegsFault, func(c *vCPU, regs *syscall.PtraceRegs, pt *pagetables.PageTables) bool {
 		testutil.SetTestRegs(regs) // Fill values for all registers.
-		if si, _, err := c.SwitchToUser(ring0.SwitchOpts{
-			Registers:          regs,
-			FloatingPointState: dummyFPState,
-			PageTables:         pt,
-			FullRestore:        true,
-		}); err != platform.ErrContextSignal || si.Signo != int32(syscall.SIGSEGV) {
-			t.Errorf("application register check with full restore got unexpected error: %v", err)
-		}
-		if err := testutil.CheckTestRegs(regs, true); err != nil {
-			t.Errorf("application register check with full restore failed: %v", err)
+		for {
+			if si, _, err := c.SwitchToUser(ring0.SwitchOpts{
+				Registers:          regs,
+				FloatingPointState: dummyFPState,
+				PageTables:         pt,
+				FullRestore:        true,
+			}); err == platform.ErrContextInterrupt {
+				continue // Retry.
+			} else if err != platform.ErrContextSignal || si.Signo != int32(syscall.SIGSEGV) {
+				t.Errorf("application register check with full restore got unexpected error: %v", err)
+			}
+			if err := testutil.CheckTestRegs(regs, true); err != nil {
+				t.Errorf("application register check with full restore failed: %v", err)
+			}
+			break // Done.
 		}
 		return false
 	})
@@ -238,16 +256,21 @@ func TestRegistersFault(t *testing.T) {
 func TestSegments(t *testing.T) {
 	applicationTest(t, true, testutil.TwiddleSegments, func(c *vCPU, regs *syscall.PtraceRegs, pt *pagetables.PageTables) bool {
 		testutil.SetTestSegments(regs)
-		if _, _, err := c.SwitchToUser(ring0.SwitchOpts{
-			Registers:          regs,
-			FloatingPointState: dummyFPState,
-			PageTables:         pt,
-			FullRestore:        true,
-		}); err != nil {
-			t.Errorf("application segment check with full restore got unexpected error: %v", err)
-		}
-		if err := testutil.CheckTestSegments(regs); err != nil {
-			t.Errorf("application segment check with full restore failed: %v", err)
+		for {
+			if _, _, err := c.SwitchToUser(ring0.SwitchOpts{
+				Registers:          regs,
+				FloatingPointState: dummyFPState,
+				PageTables:         pt,
+				FullRestore:        true,
+			}); err == platform.ErrContextInterrupt {
+				continue // Retry.
+			} else if err != nil {
+				t.Errorf("application segment check with full restore got unexpected error: %v", err)
+			}
+			if err := testutil.CheckTestSegments(regs); err != nil {
+				t.Errorf("application segment check with full restore failed: %v", err)
+			}
+			break // Done.
 		}
 		return false
 	})
@@ -323,22 +346,32 @@ func TestInvalidate(t *testing.T) {
 	var data uintptr // Used below.
 	applicationTest(t, true, testutil.Touch, func(c *vCPU, regs *syscall.PtraceRegs, pt *pagetables.PageTables) bool {
 		testutil.SetTouchTarget(regs, &data) // Read legitimate value.
-		if _, _, err := c.SwitchToUser(ring0.SwitchOpts{
-			Registers:          regs,
-			FloatingPointState: dummyFPState,
-			PageTables:         pt,
-		}); err != nil {
-			t.Errorf("application partial restore: got %v, wanted nil", err)
+		for {
+			if _, _, err := c.SwitchToUser(ring0.SwitchOpts{
+				Registers:          regs,
+				FloatingPointState: dummyFPState,
+				PageTables:         pt,
+			}); err == platform.ErrContextInterrupt {
+				continue // Retry.
+			} else if err != nil {
+				t.Errorf("application partial restore: got %v, wanted nil", err)
+			}
+			break // Done.
 		}
 		// Unmap the page containing data & invalidate.
 		pt.Unmap(usermem.Addr(reflect.ValueOf(&data).Pointer() & ^uintptr(usermem.PageSize-1)), usermem.PageSize)
-		if _, _, err := c.SwitchToUser(ring0.SwitchOpts{
-			Registers:          regs,
-			FloatingPointState: dummyFPState,
-			PageTables:         pt,
-			Flush:              true,
-		}); err != platform.ErrContextSignal {
-			t.Errorf("application partial restore: got %v, wanted %v", err, platform.ErrContextSignal)
+		for {
+			if _, _, err := c.SwitchToUser(ring0.SwitchOpts{
+				Registers:          regs,
+				FloatingPointState: dummyFPState,
+				PageTables:         pt,
+				Flush:              true,
+			}); err == platform.ErrContextInterrupt {
+				continue // Retry.
+			} else if err != platform.ErrContextSignal {
+				t.Errorf("application partial restore: got %v, wanted %v", err, platform.ErrContextSignal)
+			}
+			break // Success.
 		}
 		return false
 	})
@@ -355,7 +388,9 @@ func TestEmptyAddressSpace(t *testing.T) {
 			Registers:          regs,
 			FloatingPointState: dummyFPState,
 			PageTables:         pt,
-		}); !IsFault(err, si) {
+		}); err == platform.ErrContextInterrupt {
+			return true // Retry.
+		} else if !IsFault(err, si) {
 			t.Errorf("first fault with partial restore failed got %v", err)
 			t.Logf("registers: %#v", &regs)
 		}
@@ -367,7 +402,9 @@ func TestEmptyAddressSpace(t *testing.T) {
 			FloatingPointState: dummyFPState,
 			PageTables:         pt,
 			FullRestore:        true,
-		}); !IsFault(err, si) {
+		}); err == platform.ErrContextInterrupt {
+			return true // Retry.
+		} else if !IsFault(err, si) {
 			t.Errorf("first fault with full restore failed got %v", err)
 			t.Logf("registers: %#v", &regs)
 		}
@@ -422,11 +459,10 @@ func BenchmarkApplicationSyscall(b *testing.B) {
 			Registers:          regs,
 			FloatingPointState: dummyFPState,
 			PageTables:         pt,
-		}); err != nil {
-			if err == platform.ErrContextInterrupt {
-				a++
-				return true // Ignore.
-			}
+		}); err == platform.ErrContextInterrupt {
+			a++
+			return true // Ignore.
+		} else if err != nil {
 			b.Fatalf("benchmark failed: %v", err)
 		}
 		i++
@@ -459,11 +495,10 @@ func BenchmarkWorldSwitchToUserRoundtrip(b *testing.B) {
 			Registers:          regs,
 			FloatingPointState: dummyFPState,
 			PageTables:         pt,
-		}); err != nil {
-			if err == platform.ErrContextInterrupt {
-				a++
-				return true // Ignore.
-			}
+		}); err == platform.ErrContextInterrupt {
+			a++
+			return true // Ignore.
+		} else if err != nil {
 			b.Fatalf("benchmark failed: %v", err)
 		}
 		// This will intentionally cause the world switch. By executing
@@ -474,6 +509,6 @@ func BenchmarkWorldSwitchToUserRoundtrip(b *testing.B) {
 		return i < b.N
 	})
 	if a != 0 {
-		b.Logf("EAGAIN occurred %d times (in %d iterations).", a, a+i)
+		b.Logf("ErrContextInterrupt occurred %d times (in %d iterations).", a, a+i)
 	}
 }
diff --git a/pkg/sentry/platform/kvm/virtual_map_test.go b/pkg/sentry/platform/kvm/virtual_map_test.go
index 31e5b0e61..7875bd3e9 100644
--- a/pkg/sentry/platform/kvm/virtual_map_test.go
+++ b/pkg/sentry/platform/kvm/virtual_map_test.go
@@ -22,14 +22,16 @@ import (
 )
 
 type checker struct {
-	ok bool
+	ok         bool
+	accessType usermem.AccessType
 }
 
-func (c *checker) Contains(addr uintptr) func(virtualRegion) {
+func (c *checker) Containing(addr uintptr) func(virtualRegion) {
 	c.ok = false // Reset for below calls.
 	return func(vr virtualRegion) {
 		if vr.virtual <= addr && addr < vr.virtual+vr.length {
 			c.ok = true
+			c.accessType = vr.accessType
 		}
 	}
 }
@@ -38,7 +40,7 @@ func TestParseMaps(t *testing.T) {
 	c := new(checker)
 
 	// Simple test.
-	if err := applyVirtualRegions(c.Contains(0)); err != nil {
+	if err := applyVirtualRegions(c.Containing(0)); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
@@ -52,7 +54,7 @@ func TestParseMaps(t *testing.T) {
 	}
 
 	// Re-parse maps.
-	if err := applyVirtualRegions(c.Contains(addr)); err != nil {
+	if err := applyVirtualRegions(c.Containing(addr)); err != nil {
 		syscall.RawSyscall(syscall.SYS_MUNMAP, addr, usermem.PageSize, 0)
 		t.Fatalf("unexpected error: %v", err)
 	}
@@ -63,16 +65,29 @@ func TestParseMaps(t *testing.T) {
 		t.Fatalf("updated map does not contain 0x%08x, expected true", addr)
 	}
 
-	// Unmap the region.
-	syscall.RawSyscall(syscall.SYS_MUNMAP, addr, usermem.PageSize, 0)
+	// Map the region as PROT_NONE.
+	newAddr, _, errno := syscall.RawSyscall6(
+		syscall.SYS_MMAP, addr, usermem.PageSize,
+		syscall.PROT_NONE,
+		syscall.MAP_ANONYMOUS|syscall.MAP_FIXED|syscall.MAP_PRIVATE, 0, 0)
+	if errno != 0 {
+		t.Fatalf("unexpected map error: %v", errno)
+	}
+	if newAddr != addr {
+		t.Fatalf("unable to remap address: got 0x%08x, wanted 0x%08x", newAddr, addr)
+	}
 
 	// Re-parse maps.
-	if err := applyVirtualRegions(c.Contains(addr)); err != nil {
+	if err := applyVirtualRegions(c.Containing(addr)); err != nil {
 		t.Fatalf("unexpected error: %v", err)
 	}
-
-	// Assert that it once again does _not_ contain the region.
-	if c.ok {
-		t.Fatalf("final map does contain 0x%08x, expected false", addr)
+	if !c.ok {
+		t.Fatalf("final map does not contain 0x%08x, expected true", addr)
+	}
+	if c.accessType.Read || c.accessType.Write || c.accessType.Execute {
+		t.Fatalf("final map has incorrect permissions for 0x%08x", addr)
 	}
+
+	// Unmap the region.
+	syscall.RawSyscall(syscall.SYS_MUNMAP, addr, usermem.PageSize, 0)
 }