Support multiple se_flags modules

Instead of centralized one se_flags module under system/sepolicy,
additional se_flags modules can be defined anywhere to support defining
downstream branches' own flagging.

Bug: 321875465
Test: TH
Test: soong test
Change-Id: I6e45c859b7f09e27ba1d60033b0db1424472cb63
diff --git a/build/soong/flags.go b/build/soong/flags.go
index b1aebac..c7aeb32 100644
--- a/build/soong/flags.go
+++ b/build/soong/flags.go
@@ -15,22 +15,141 @@
 package selinux
 
 import (
+	"maps"
+
 	"android/soong/android"
+
+	"github.com/google/blueprint"
 )
 
+var (
+	flagsDepTag      = dependencyTag{name: "flags"}
+	buildFlagsDepTag = dependencyTag{name: "build_flags"}
+)
+
+func init() {
+	ctx := android.InitRegistrationContext
+	ctx.RegisterModuleType("se_flags", flagsFactory)
+	ctx.RegisterModuleType("se_flags_collector", flagsCollectorFactory)
+}
+
 type flagsProperties struct {
-	// List of flags to be passed to M4 macro.
+	// List of build time flags for flag-guarding.
 	Flags []string
+
+	// List of se_flags_collector modules to export flags to.
+	Export_to []string
+}
+
+type flagsModule struct {
+	android.ModuleBase
+	properties flagsProperties
+}
+
+type flagsInfo struct {
+	Flags []string
+}
+
+var flagsProviderKey = blueprint.NewProvider[flagsInfo]()
+
+// se_flags contains a list of build time flags for sepolicy.  Build time flags are defined under
+// .scl files (e.g. build/release/build_flags.scl). By importing flags with se_flags modules,
+// sepolicy rules can be guarded by `is_flag_enabled` / `is_flag_disabled` macro.
+//
+// For example, an Android.bp file could have:
+//
+//	se_flags {
+//		name: "aosp_selinux_flags",
+//		flags: ["RELEASE_AVF_ENABLE_DEVICE_ASSIGNMENT"],
+//		export_to: ["all_selinux_flags"],
+//	}
+//
+// And then one could flag-guard .te file rules:
+//
+//	is_flag_enabled(RELEASE_AVF_ENABLE_DEVICE_ASSIGNMENT, `
+//		type vfio_handler, domain, coredomain;
+//		binder_use(vfio_handler)
+//	')
+//
+// or contexts entries:
+//
+//	is_flag_enabled(RELEASE_AVF_ENABLE_DEVICE_ASSIGNMENT, `
+//		android.system.virtualizationservice_internal.IVfioHandler u:object_r:vfio_handler_service:s0
+//	')
+func flagsFactory() android.Module {
+	module := &flagsModule{}
+	module.AddProperties(&module.properties)
+	android.InitAndroidModule(module)
+	return module
+}
+
+func (f *flagsModule) DepsMutator(ctx android.BottomUpMutatorContext) {
+	// dep se_flag_collector -> se_flags
+	for _, export := range f.properties.Export_to {
+		ctx.AddReverseDependency(ctx.Module(), flagsDepTag, export)
+	}
+}
+
+func (f *flagsModule) GenerateAndroidBuildActions(ctx android.ModuleContext) {
+	android.SetProvider(ctx, flagsProviderKey, flagsInfo{
+		Flags: f.properties.Flags,
+	})
+}
+
+type buildFlagsInfo struct {
+	BuildFlags map[string]string
+}
+
+var buildFlagsProviderKey = blueprint.NewProvider[buildFlagsInfo]()
+
+type flagsCollectorModule struct {
+	android.ModuleBase
+	buildFlags map[string]string
+}
+
+// se_flags_collector module collects flags from exported se_flags modules (see export_to property
+// of se_flags modules), and then converts them into build-time flags.  It will be used to generate
+// M4 macros to flag-guard sepolicy.
+func flagsCollectorFactory() android.Module {
+	module := &flagsCollectorModule{}
+	android.InitAndroidModule(module)
+	return module
+}
+
+func (f *flagsCollectorModule) GenerateAndroidBuildActions(ctx android.ModuleContext) {
+	var flags []string
+	ctx.VisitDirectDepsWithTag(flagsDepTag, func(m android.Module) {
+		if dep, ok := android.OtherModuleProvider(ctx, m, flagsProviderKey); ok {
+			flags = append(flags, dep.Flags...)
+		} else {
+			ctx.ModuleErrorf("unknown dependency %q", ctx.OtherModuleName(m))
+		}
+	})
+	buildFlags := make(map[string]string)
+	for _, flag := range android.SortedUniqueStrings(flags) {
+		if val, ok := ctx.Config().GetBuildFlag(flag); ok {
+			buildFlags[flag] = val
+		}
+	}
+	android.SetProvider(ctx, buildFlagsProviderKey, buildFlagsInfo{
+		BuildFlags: buildFlags,
+	})
+}
+
+type flaggableModuleProperties struct {
+	// List of se_flag_collector modules to be passed to M4 macro.
+	Build_flags []string
 }
 
 type flaggableModule interface {
 	android.Module
 	flagModuleBase() *flaggableModuleBase
+	flagDeps(ctx android.BottomUpMutatorContext)
 	getBuildFlags(ctx android.ModuleContext) map[string]string
 }
 
 type flaggableModuleBase struct {
-	properties flagsProperties
+	properties flaggableModuleProperties
 }
 
 func initFlaggableModule(m flaggableModule) {
@@ -42,13 +161,19 @@
 	return f
 }
 
+func (f *flaggableModuleBase) flagDeps(ctx android.BottomUpMutatorContext) {
+	ctx.AddDependency(ctx.Module(), buildFlagsDepTag, f.properties.Build_flags...)
+}
+
 // getBuildFlags returns a map from flag names to flag values.
 func (f *flaggableModuleBase) getBuildFlags(ctx android.ModuleContext) map[string]string {
 	ret := make(map[string]string)
-	for _, flag := range android.SortedUniqueStrings(f.properties.Flags) {
-		if val, ok := ctx.Config().GetBuildFlag(flag); ok {
-			ret[flag] = val
+	ctx.VisitDirectDepsWithTag(buildFlagsDepTag, func(m android.Module) {
+		if dep, ok := android.OtherModuleProvider(ctx, m, buildFlagsProviderKey); ok {
+			maps.Copy(ret, dep.BuildFlags)
+		} else {
+			ctx.PropertyErrorf("build_flags", "unknown dependency %q", ctx.OtherModuleName(m))
 		}
-	}
+	})
 	return ret
 }