2 files changed, 24 insertions, 13 deletions

diff --git a/pkg/sentry/platform/kvm/address_space.go b/pkg/sentry/platform/kvm/address_space.go
index a777533c5..4c76883ad 100644
--- a/pkg/sentry/platform/kvm/address_space.go
+++ b/pkg/sentry/platform/kvm/address_space.go
@@ -16,6 +16,7 @@ package kvm
 
 import (
 	"reflect"
+	"sync"
 
 	"gvisor.googlesource.com/gvisor/pkg/sentry/platform"
 	"gvisor.googlesource.com/gvisor/pkg/sentry/platform/filemem"
@@ -27,6 +28,11 @@ import (
 type addressSpace struct {
 	platform.NoAddressSpaceIO
 
+	// mu is the lock for modifications to the address space.
+	//
+	// Note that the page tables themselves are not locked.
+	mu sync.Mutex
+
 	// filemem is the memory instance.
 	filemem *filemem.FileMem
 
@@ -49,8 +55,8 @@ type addressSpace struct {
 	files hostMap
 }
 
-// Invalidate interrupts all dirty contexts.
-func (as *addressSpace) Invalidate() {
+// invalidate is the implementation for Invalidate.
+func (as *addressSpace) invalidate() {
 	for i := 0; i < as.dirtySet.size(); i++ {
 		if c := as.dirtySet.swap(i, nil); c != nil && c.active.get() == as {
 			c.BounceToKernel() // Force a kernel transition.
@@ -58,6 +64,13 @@ func (as *addressSpace) Invalidate() {
 	}
 }
 
+// Invalidate interrupts all dirty contexts.
+func (as *addressSpace) Invalidate() {
+	as.mu.Lock()
+	defer as.mu.Unlock()
+	as.invalidate()
+}
+
 // Touch adds the given vCPU to the dirty list.
 //
 // The return value indicates whether a flush is required.
@@ -120,7 +133,7 @@ func (as *addressSpace) mapHostFile(addr usermem.Addr, fd int, fr platform.FileR
 		addr += usermem.Addr(m.length)
 	}
 	if inv {
-		as.Invalidate()
+		as.invalidate()
 	}
 
 	return nil
@@ -169,7 +182,7 @@ func (as *addressSpace) mapFilemem(addr usermem.Addr, fr platform.FileRange, at
 		addr += usermem.Addr(len(s))
 	}
 	if inv {
-		as.Invalidate()
+		as.invalidate()
 		as.files.DeleteMapping(orig)
 	}
 
@@ -178,6 +191,9 @@ func (as *addressSpace) mapFilemem(addr usermem.Addr, fr platform.FileRange, at
 
 // MapFile implements platform.AddressSpace.MapFile.
 func (as *addressSpace) MapFile(addr usermem.Addr, fd int, fr platform.FileRange, at usermem.AccessType, precommit bool) error {
+	as.mu.Lock()
+	defer as.mu.Unlock()
+
 	// Create an appropriate mapping. If this is filemem, we don't create
 	// custom mappings for each in-application mapping. For files however,
 	// we create distinct mappings for each address space. Unfortunately,
@@ -195,8 +211,11 @@ func (as *addressSpace) MapFile(addr usermem.Addr, fd int, fr platform.FileRange
 
 // Unmap unmaps the given range by calling pagetables.PageTables.Unmap.
 func (as *addressSpace) Unmap(addr usermem.Addr, length uint64) {
+	as.mu.Lock()
+	defer as.mu.Unlock()
+
 	if prev := as.pageTables.Unmap(addr, uintptr(length)); prev {
-		as.Invalidate()
+		as.invalidate()
 		as.files.DeleteMapping(usermem.AddrRange{
 			Start: addr,
 			End:   addr + usermem.Addr(length),
diff --git a/pkg/sentry/platform/ring0/pagetables/pagetables.go b/pkg/sentry/platform/ring0/pagetables/pagetables.go
index 2df6792f7..2a83bbff2 100644
--- a/pkg/sentry/platform/ring0/pagetables/pagetables.go
+++ b/pkg/sentry/platform/ring0/pagetables/pagetables.go
@@ -16,8 +16,6 @@
 package pagetables
 
 import (
-	"sync"
-
 	"gvisor.googlesource.com/gvisor/pkg/sentry/usermem"
 )
 
@@ -39,8 +37,6 @@ type Node struct {
 
 // PageTables is a set of page tables.
 type PageTables struct {
-	mu sync.Mutex
-
 	// root is the pagetable root.
 	root *Node
 
@@ -122,7 +118,6 @@ func (p *PageTables) Map(addr usermem.Addr, length uintptr, opts MapOpts, physic
 		return p.Unmap(addr, length)
 	}
 	prev := false
-	p.mu.Lock()
 	end, ok := addr.AddLength(uint64(length))
 	if !ok {
 		panic("pagetables.Map: overflow")
@@ -139,7 +134,6 @@ func (p *PageTables) Map(addr usermem.Addr, length uintptr, opts MapOpts, physic
 		}
 		pte.Set(p, opts)
 	})
-	p.mu.Unlock()
 	return prev
 }
 
@@ -147,13 +141,11 @@ func (p *PageTables) Map(addr usermem.Addr, length uintptr, opts MapOpts, physic
 //
 // True is returned iff there was a previous mapping in the range.
 func (p *PageTables) Unmap(addr usermem.Addr, length uintptr) bool {
-	p.mu.Lock()
 	count := 0
 	p.iterateRange(uintptr(addr), uintptr(addr)+length, false, func(s, e uintptr, pte *PTE, align uintptr) {
 		pte.Clear()
 		count++
 	})
-	p.mu.Unlock()
 	return count > 0
 }