Make memtag_stack a variant

Otherwise we cannot disable memtag_stack on first stage init, because it
statically links in a lot of libraries.

This was benign, but if we want to add a stack history buffer, we do not
want to emit these instructions for first stage init.

Bug: 309446520
Change-Id: Ifbdcbe5ba06eb46970bc1e9b64a26c1ca59ba5f2
diff --git a/cc/sanitize.go b/cc/sanitize.go
index db046ec..188b793 100644
--- a/cc/sanitize.go
+++ b/cc/sanitize.go
@@ -176,11 +176,11 @@
 
 func (t SanitizerType) registerMutators(ctx android.RegisterMutatorsContext) {
 	switch t {
-	case cfi, Hwasan, Asan, tsan, Fuzzer, scs:
+	case cfi, Hwasan, Asan, tsan, Fuzzer, scs, Memtag_stack:
 		sanitizer := &sanitizerSplitMutator{t}
 		ctx.TopDown(t.variationName()+"_markapexes", sanitizer.markSanitizableApexesMutator)
 		ctx.Transition(t.variationName(), sanitizer)
-	case Memtag_heap, Memtag_stack, Memtag_globals, intOverflow:
+	case Memtag_heap, Memtag_globals, intOverflow:
 		// do nothing
 	default:
 		panic(fmt.Errorf("unknown SanitizerType %d", t))
@@ -407,6 +407,7 @@
 
 	android.RegisterMakeVarsProvider(pctx, cfiMakeVarsProvider)
 	android.RegisterMakeVarsProvider(pctx, hwasanMakeVarsProvider)
+	android.RegisterMakeVarsProvider(pctx, memtagStackMakeVarsProvider)
 }
 
 func (sanitize *sanitize) props() []interface{} {
@@ -683,10 +684,14 @@
 		s.Diag.Cfi = nil
 	}
 
-	// HWASan ramdisk (which is built from recovery) goes over some bootloader limit.
-	// Keep libc instrumented so that ramdisk / vendor_ramdisk / recovery can run hwasan-instrumented code if necessary.
-	if (ctx.inRamdisk() || ctx.inVendorRamdisk() || ctx.inRecovery()) && !strings.HasPrefix(ctx.ModuleDir(), "bionic/libc") {
-		s.Hwaddress = nil
+	if ctx.inRamdisk() || ctx.inVendorRamdisk() || ctx.inRecovery() {
+		// HWASan ramdisk (which is built from recovery) goes over some bootloader limit.
+		// Keep libc instrumented so that ramdisk / vendor_ramdisk / recovery can run hwasan-instrumented code if necessary.
+		if !strings.HasPrefix(ctx.ModuleDir(), "bionic/libc") {
+			s.Hwaddress = nil
+		}
+		// Memtag stack in ramdisk makes pKVM unhappy.
+		s.Memtag_stack = nil
 	}
 
 	if ctx.staticBinary() {
@@ -1303,6 +1308,8 @@
 					hwasanStaticLibs(mctx.Config()).add(c, c.Module().Name())
 				} else if s.sanitizer == cfi {
 					cfiStaticLibs(mctx.Config()).add(c, c.Module().Name())
+				} else if s.sanitizer == Memtag_stack {
+					memtagStackStaticLibs(mctx.Config()).add(c, c.Module().Name());
 				}
 			}
 		} else if c.IsSanitizerEnabled(s.sanitizer) {
@@ -1715,6 +1722,14 @@
 	}).(*sanitizerStaticLibsMap)
 }
 
+var memtagStackStaticLibsKey = android.NewOnceKey("memtagStackStaticLibs")
+
+func memtagStackStaticLibs(config android.Config) *sanitizerStaticLibsMap {
+	return config.Once(memtagStackStaticLibsKey, func() interface{} {
+		return newSanitizerStaticLibsMap(Memtag_stack)
+	}).(*sanitizerStaticLibsMap)
+}
+
 func enableMinimalRuntime(sanitize *sanitize) bool {
 	if sanitize.isSanitizerEnabled(Asan) {
 		return false
@@ -1761,3 +1776,7 @@
 func hwasanMakeVarsProvider(ctx android.MakeVarsContext) {
 	hwasanStaticLibs(ctx.Config()).exportToMake(ctx)
 }
+
+func memtagStackMakeVarsProvider(ctx android.MakeVarsContext) {
+	memtagStackStaticLibs(ctx.Config()).exportToMake(ctx)
+}