Enforce partition property on apex system server jars

This is a followup AI from https://r.android.com/3288083, which installs
dexpreopt files of apex system server jars in the same partition as the
top-level apex. This CL enforces that the partition properties of the
apex and the java library match.

This has been implemented in a bottomup manner.
systemserverclasspath_fragment module sets
`LibraryNameToPartitionInfoProvider`. The top-level apex uses this info
and raies an exception if the partition properties do not match.

This enforcement is done only for source apexes for now. It is not
needed for prebuilts since
- The dexpreopt rules of jars in prebuilt apexes are generated in the
  top-level `prebuiltApex`/ `apexSet`
- We do not have prebuilts of system_ext apexes today
  (com.android.compos is not part of mainline)

Test: m nothing --no-skip-soong-tests

Change-Id: Ib4df634457f315d5421681bdb0afebc2b1bf92d9

3 files changed, 93 insertions, 0 deletions

diff --git a/apex/apex.go b/apex/apex.go
index d7dc6d718..d3e7eee9d 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -29,6 +29,7 @@ import (
 	"android/soong/android"
 	"android/soong/bpf"
 	"android/soong/cc"
+	"android/soong/dexpreopt"
 	prebuilt_etc "android/soong/etc"
 	"android/soong/filesystem"
 	"android/soong/java"
@@ -1919,6 +1920,32 @@ func (vctx *visitorContext) normalizeFileInfo(mctx android.ModuleContext) {
 	})
 }
 
+// enforcePartitionTagOnApexSystemServerJar checks that the partition tags of an apex system server jar  matches
+// the partition tags of the top-level apex.
+// e.g. if the top-level apex sets system_ext_specific to true, the javalib must set this property to true as well.
+// This check ensures that the dexpreopt artifacts of the apex system server jar is installed in the same partition
+// as the apex.
+func (a *apexBundle) enforcePartitionTagOnApexSystemServerJar(ctx android.ModuleContext) {
+	global := dexpreopt.GetGlobalConfig(ctx)
+	ctx.VisitDirectDepsWithTag(sscpfTag, func(child android.Module) {
+		info, ok := android.OtherModuleProvider(ctx, child, java.LibraryNameToPartitionInfoProvider)
+		if !ok {
+			ctx.ModuleErrorf("Could not find partition info of apex system server jars.")
+		}
+		apexPartition := ctx.Module().PartitionTag(ctx.DeviceConfig())
+		for javalib, javalibPartition := range info.LibraryNameToPartition {
+			if !global.AllApexSystemServerJars(ctx).ContainsJar(javalib) {
+				continue // not an apex system server jar
+			}
+			if apexPartition != javalibPartition {
+				ctx.ModuleErrorf(`
+%s is an apex systemserver jar, but its partition does not match the partition of its containing apex. Expected %s, Got %s`,
+					javalib, apexPartition, javalibPartition)
+			}
+		}
+	})
+}
+
 func (a *apexBundle) depVisitor(vctx *visitorContext, ctx android.ModuleContext, child, parent android.Module) bool {
 	depTag := ctx.OtherModuleDependencyTag(child)
 	if _, ok := depTag.(android.ExcludeFromApexContentsTag); ok {
@@ -2341,6 +2368,7 @@ func (a *apexBundle) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 	a.required = append(a.required, a.VintfFragmentModuleNames(ctx)...)
 
 	a.setOutputFiles(ctx)
+	a.enforcePartitionTagOnApexSystemServerJar(ctx)
 }
 
 // Set prebuiltInfoProvider. This will be used by `apex_prebuiltinfo_singleton` to print out a metadata file
diff --git a/apex/apex_test.go b/apex/apex_test.go
index da6214db3..153ee3108 100644
--- a/apex/apex_test.go
+++ b/apex/apex_test.go
@@ -11670,3 +11670,48 @@ func TestSdkLibraryTransitiveClassLoaderContext(t *testing.T) {
 		}
        `)
 }
+
+// If an apex sets system_ext_specific: true, its systemserverclasspath libraries must set this property as well.
+func TestApexSSCPJarMustBeInSamePartitionAsApex(t *testing.T) {
+	testApexError(t, `foo is an apex systemserver jar, but its partition does not match the partition of its containing apex`, `
+		apex {
+			name: "myapex",
+			key: "myapex.key",
+			systemserverclasspath_fragments: [
+				"mysystemserverclasspathfragment",
+			],
+			min_sdk_version: "29",
+			updatable: true,
+			system_ext_specific: true,
+		}
+
+		apex_key {
+			name: "myapex.key",
+			public_key: "testkey.avbpubkey",
+			private_key: "testkey.pem",
+		}
+
+		java_library {
+			name: "foo",
+			srcs: ["b.java"],
+			min_sdk_version: "29",
+			installable: true,
+			apex_available: [
+				"myapex",
+			],
+			sdk_version: "current",
+		}
+
+		systemserverclasspath_fragment {
+			name: "mysystemserverclasspathfragment",
+			contents: [
+				"foo",
+			],
+			apex_available: [
+				"myapex",
+			],
+		}
+	`,
+		dexpreopt.FixtureSetApexSystemServerJars("myapex:foo"),
+	)
+}
diff --git a/java/systemserver_classpath_fragment.go b/java/systemserver_classpath_fragment.go
index 924abd460..aad106007 100644
--- a/java/systemserver_classpath_fragment.go
+++ b/java/systemserver_classpath_fragment.go
@@ -127,6 +127,26 @@ func (s *SystemServerClasspathModule) GenerateAndroidBuildActions(ctx android.Mo
 	configuredJars = configuredJars.AppendList(&standaloneConfiguredJars)
 	classpathJars = append(classpathJars, standaloneClasspathJars...)
 	s.classpathFragmentBase().generateClasspathProtoBuildActions(ctx, configuredJars, classpathJars)
+	s.setPartitionInfoOfLibraries(ctx)
+}
+
+// Map of java library name to their install partition.
+type LibraryNameToPartitionInfo struct {
+	LibraryNameToPartition map[string]string
+}
+
+// LibraryNameToPartitionInfoProvider will be used by the top-level apex to enforce that dexpreopt files
+// of apex system server jars are installed in the same partition as the top-level apex.
+var LibraryNameToPartitionInfoProvider = blueprint.NewProvider[LibraryNameToPartitionInfo]()
+
+func (s *SystemServerClasspathModule) setPartitionInfoOfLibraries(ctx android.ModuleContext) {
+	libraryNameToPartition := map[string]string{}
+	ctx.VisitDirectDepsWithTag(systemServerClasspathFragmentContentDepTag, func(m android.Module) {
+		libraryNameToPartition[m.Name()] = m.PartitionTag(ctx.DeviceConfig())
+	})
+	android.SetProvider(ctx, LibraryNameToPartitionInfoProvider, LibraryNameToPartitionInfo{
+		LibraryNameToPartition: libraryNameToPartition,
+	})
 }
 
 func (s *SystemServerClasspathModule) configuredJars(ctx android.ModuleContext) android.ConfiguredJarList {