4 files changed, 51 insertions, 9 deletions
diff --git a/pkg/sentry/platform/kvm/BUILD b/pkg/sentry/platform/kvm/BUILD
index b307832fd..77079cae9 100644
--- a/pkg/sentry/platform/kvm/BUILD
+++ b/pkg/sentry/platform/kvm/BUILD
@@ -6,6 +6,8 @@ go_library(
     name = "kvm",
     srcs = [
         "address_space.go",
+        "address_space_amd64.go",
+        "address_space_arm64.go",
         "bluepill.go",
         "bluepill_allocator.go",
         "bluepill_amd64.go",
diff --git a/pkg/sentry/platform/kvm/address_space.go b/pkg/sentry/platform/kvm/address_space.go
index 5524e8727..9929caebb 100644
--- a/pkg/sentry/platform/kvm/address_space.go
+++ b/pkg/sentry/platform/kvm/address_space.go
@@ -85,15 +85,6 @@ type addressSpace struct {
 	dirtySet *dirtySet
 }
 
-// invalidate is the implementation for Invalidate.
-func (as *addressSpace) invalidate() {
-	as.dirtySet.forEach(as.machine, func(c *vCPU) {
-		if c.active.get() == as { // If this happens to be active,
-			c.BounceToKernel() // ... force a kernel transition.
-		}
-	})
-}
-
 // Invalidate interrupts all dirty contexts.
 func (as *addressSpace) Invalidate() {
 	as.mu.Lock()
diff --git a/pkg/sentry/platform/kvm/address_space_amd64.go b/pkg/sentry/platform/kvm/address_space_amd64.go
new file mode 100644
index 000000000..d11d38679
--- /dev/null
+++ b/pkg/sentry/platform/kvm/address_space_amd64.go
@@ -0,0 +1,24 @@
+// Copyright 2021 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package kvm
+
+// invalidate is the implementation for Invalidate.
+func (as *addressSpace) invalidate() {
+	as.dirtySet.forEach(as.machine, func(c *vCPU) {
+		if c.active.get() == as { // If this happens to be active,
+			c.BounceToKernel() // ... force a kernel transition.
+		}
+	})
+}
diff --git a/pkg/sentry/platform/kvm/address_space_arm64.go b/pkg/sentry/platform/kvm/address_space_arm64.go
new file mode 100644
index 000000000..fb954418b
--- /dev/null
+++ b/pkg/sentry/platform/kvm/address_space_arm64.go
@@ -0,0 +1,25 @@
+// Copyright 2021 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package kvm
+
+import (
+	"gvisor.dev/gvisor/pkg/ring0"
+)
+
+// invalidate is the implementation for Invalidate.
+func (as *addressSpace) invalidate() {
+	bluepill(as.pageTables.Allocator.(*allocator).cpu)
+	ring0.FlushTlbAll()
+}