Merge pull request #4880 from lubinszARM:pr_tlbi_02

PiperOrigin-RevId: 347890782

3 files changed, 25 insertions, 2 deletions

diff --git a/pkg/sentry/platform/ring0/kernel_arm64.go b/pkg/sentry/platform/ring0/kernel_arm64.go
index bffe27e5c..c05284641 100644
--- a/pkg/sentry/platform/ring0/kernel_arm64.go
+++ b/pkg/sentry/platform/ring0/kernel_arm64.go
@@ -63,7 +63,7 @@ func IsCanonical(addr uint64) bool {
 func (c *CPU) SwitchToUser(switchOpts SwitchOpts) (vector Vector) {
 	storeAppASID(uintptr(switchOpts.UserASID))
 	if switchOpts.Flush {
-		FlushTlbAll()
+		FlushTlbByASID(uintptr(switchOpts.UserASID))
 	}
 
 	regs := switchOpts.Registers
diff --git a/pkg/sentry/platform/ring0/lib_arm64.go b/pkg/sentry/platform/ring0/lib_arm64.go
index 29b3efd34..a490bf3af 100644
--- a/pkg/sentry/platform/ring0/lib_arm64.go
+++ b/pkg/sentry/platform/ring0/lib_arm64.go
@@ -22,7 +22,13 @@ func storeAppASID(asid uintptr)
 // LocalFlushTlbAll same as FlushTlbAll, but only applies to the calling CPU.
 func LocalFlushTlbAll()
 
-// FlushTlbAll flush all tlb.
+// FlushTlbByVA invalidates tlb by VA/Last-level/Inner-Shareable.
+func FlushTlbByVA(addr uintptr)
+
+// FlushTlbByASID invalidates tlb by ASID/Inner-Shareable.
+func FlushTlbByASID(asid uintptr)
+
+// FlushTlbAll invalidates all tlb.
 func FlushTlbAll()
 
 // CPACREL1 returns the value of the CPACR_EL1 register.
diff --git a/pkg/sentry/platform/ring0/lib_arm64.s b/pkg/sentry/platform/ring0/lib_arm64.s
index 6f4923539..e39b32841 100644
--- a/pkg/sentry/platform/ring0/lib_arm64.s
+++ b/pkg/sentry/platform/ring0/lib_arm64.s
@@ -15,6 +15,23 @@
 #include "funcdata.h"
 #include "textflag.h"
 
+#define TLBI_ASID_SHIFT		48
+
+TEXT ·FlushTlbByVA(SB),NOSPLIT,$0-8
+	MOVD addr+0(FP), R1
+	DSB $10                 // dsb(ishst)
+	WORD $0xd50883a1        // tlbi vale1is, x1
+	DSB $11                 // dsb(ish)
+	RET
+
+TEXT ·FlushTlbByASID(SB),NOSPLIT,$0-8
+	MOVD asid+0(FP), R1
+	LSL $TLBI_ASID_SHIFT, R1, R1
+	DSB $10                 // dsb(ishst)
+	WORD $0xd5088341        // tlbi aside1is, x1
+	DSB $11                 // dsb(ish)
+	RET
+
 TEXT ·LocalFlushTlbAll(SB),NOSPLIT,$0
 	DSB $6			// dsb(nshst)
 	WORD $0xd508871f	// __tlbi(vmalle1)