apex: use proper partition tag for validations

module.PartitionTag(DeviceConfig) doesn't work for validations because
it tells the physical partition where the apex is installed. It returns
"system" even for vendor APEXes when "VendorPath" is set to
"system/vendor".

For validations, we need logical partition tag. So, "vendor:true" apexes
should be "vendor" even when it's installed in /system/vendor.

Bug: 384988749
Test: m --no-skip-soong-tests
Change-Id: Ib478caa62c98808bc94fe350a2f182f1d2b42a46
diff --git a/apex/apex_test.go b/apex/apex_test.go
index e1a9582..c123d03 100644
--- a/apex/apex_test.go
+++ b/apex/apex_test.go
@@ -7090,6 +7090,34 @@
 	`)
 }
 
+func TestApexValidation_UsesProperPartitionTag(t *testing.T) {
+	t.Parallel()
+	ctx := testApex(t, `
+		apex {
+			name: "myapex",
+			key: "myapex.key",
+			updatable: false,
+			vendor: true,
+		}
+		apex_key {
+			name: "myapex.key",
+			public_key: "testkey.avbpubkey",
+			private_key: "testkey.pem",
+		}
+	`, android.FixtureModifyProductVariables(func(variables android.FixtureProductVariables) {
+		// vendor path should not affect "partition tag"
+		variables.VendorPath = proptools.StringPtr("system/vendor")
+	}))
+
+	module := ctx.ModuleForTests("myapex", "android_common_myapex")
+	android.AssertStringEquals(t, "partition tag for host_apex_verifier",
+		"vendor",
+		module.Output("host_apex_verifier.timestamp").Args["partition_tag"])
+	android.AssertStringEquals(t, "partition tag for apex_sepolicy_tests",
+		"vendor",
+		module.Output("apex_sepolicy_tests.timestamp").Args["partition_tag"])
+}
+
 func TestApexValidation_TestApexCanSkipInitRcCheck(t *testing.T) {
 	t.Parallel()
 	ctx := testApex(t, `
diff --git a/apex/builder.go b/apex/builder.go
index d9348c5..daba6f1 100644
--- a/apex/builder.go
+++ b/apex/builder.go
@@ -1202,6 +1202,22 @@
 	return timestamp
 }
 
+// Can't use PartitionTag() because PartitionTag() returns the partition this module is actually
+// installed (e.g. PartitionTag() may return "system" for vendor apex when vendor is linked to /system/vendor)
+func (a *apexBundle) partition() string {
+	if a.SocSpecific() {
+		return "vendor"
+	} else if a.DeviceSpecific() {
+		return "odm"
+	} else if a.ProductSpecific() {
+		return "product"
+	} else if a.SystemExtSpecific() {
+		return "system_ext"
+	} else {
+		return "system"
+	}
+}
+
 // Runs apex_sepolicy_tests
 //
 // $ apex-ls -Z {apex_file} > {file_contexts}
@@ -1213,7 +1229,7 @@
 		Input:  apexFile,
 		Output: timestamp,
 		Args: map[string]string{
-			"partition_tag": a.PartitionTag(ctx.DeviceConfig()),
+			"partition_tag": a.partition(),
 		},
 	})
 	return timestamp
@@ -1240,7 +1256,7 @@
 		Input:  apexFile,
 		Output: timestamp,
 		Args: map[string]string{
-			"partition_tag": a.PartitionTag(ctx.DeviceConfig()),
+			"partition_tag": a.partition(),
 		},
 	})
 	return timestamp