o/devicestate: add seed allowlist to recovery-system creation

overlord/devicestate/devicestate.go

@@ -1673,6 +1673,9 @@ type recoverySystemSetup struct {
 	// LocalComponents is a list of components that should be used to create the
 	// recovery system.
 	LocalComponents []snapstate.PathComponent `json:"local-components,omitempty"`
+	// Allowlist identifies the snaps and components that may be used to create
+	// the recovery system.
+	Allowlist *SeedAllowlist `json:"allowlist,omitempty"`
 	// TestSystem is set to true if the new recovery system should
 	// not be verified by rebooting into the new system. Once the system is
 	// created, it will immediately be considered a valid recovery system.
@@ -1745,7 +1748,7 @@ func SeedRefreshTasks(
 		return nil, nil, nil
 	}
 
-	triggers := seedRefreshTriggers(st, dctx)
+	triggers, allowlist := seedRefreshPolicy(st, dctx)
 
 	var snapsups, compsups []string
 	added := make(map[string]bool, len(candidates))
@@ -1766,13 +1769,19 @@ func SeedRefreshTasks(
 		return nil, nil, nil
 	}
 
+	seedAllowlist, err := allowlist()
+	if err != nil {
+		return nil, nil, err
+	}
+
 	labelBase := timeNow().Format("20060102")
 	label, err := pickRecoverySystemLabel(labelBase)
 	if err != nil {
 		return nil, nil, fmt.Errorf("cannot select non-conflicting label for recovery system %q: %v", labelBase, err)
 	}
 
 	ts, err := createRecoverySystemTasks(st, label, snapsups, compsups, CreateRecoverySystemOptions{
+		Allowlist:   &seedAllowlist,
 		TestSystem:  true,
 		MarkDefault: true,
 		SeedRefresh: true,
@@ -1830,12 +1839,17 @@ func UpdateSeedRefreshChange(seedTS *snapstate.SeedRefreshTasks, dctx snapstate.
 		return false, nil
 	}
 
-	triggers := seedRefreshTriggers(seedTS.Create.State(), dctx)
-	ok, err := triggers(candidate.InstanceName)
+	setup, err := taskRecoverySystemSetup(seedTS.Create)
 	if err != nil {
 		return false, err
 	}
-	if !ok {
+	if setup.Allowlist == nil {
+		return false, errors.New("internal error: seed-refresh recovery system setup is missing seed allowlist")
+	}
+
+	// we've already calculated which candidates are allowed to go into the
+	// seed. avoid opening the seed again by using that list
+	if !strutil.ListContains(setup.Allowlist.Snaps, candidate.InstanceName) {
 		return false, nil
 	}
 
@@ -1864,7 +1878,7 @@ func appendSeedRefreshCandidate(create *state.Task, snapSetupTasks, compSetupTas
 // TODO:SEEDREFRESH: remove this once we support seed-refresh seeds
 // gaining/losing snaps
 func CheckSeedRefreshRemove(st *state.State, si *snap.Info, dctx snapstate.DeviceContext) error {
-	triggers := seedRefreshTriggers(st, dctx)
+	triggers, _ := seedRefreshPolicy(st, dctx)
 	ok, err := triggers(si.SnapName())
 	if err != nil {
 		return err
@@ -1876,69 +1890,139 @@ func CheckSeedRefreshRemove(st *state.State, si *snap.Info, dctx snapstate.Devic
 	return nil
 }
 
-// seedRefreshTriggers returns a closure that reports whether the given snap
-// should trigger a seed refresh. The seed is lazily loaded, and only opened
-// when required.
-func seedRefreshTriggers(st *state.State, dctx snapstate.DeviceContext) func(string) (bool, error) {
-	required := make(map[string]bool)
-	optional := make(map[string]bool)
+// seedRefreshPolicy returns closures that define the seed refresh policy. The
+// returned triggers function defines which snaps will trigger the seed refresh.
+// The allowlist function returns which snaps and components are allowed in the
+// final seed. This allowlist is composed of required model snaps/components and
+// optional model snaps/components that are present in the current seed.
+//
+// This function returns closures to help avoid opening the seed in as many
+// cases as possible. The closures share the state needed from the seed, so
+// we'll open the seed at most once.
+func seedRefreshPolicy(st *state.State, dctx snapstate.DeviceContext) (triggers func(string) (bool, error), allowlist func() (SeedAllowlist, error)) {
+	required := make(map[string]*asserts.ModelSnap)
+	optional := make(map[string]*asserts.ModelSnap)
+
 	for _, sn := range dctx.Model().AllSnaps() {
 		if sn.Presence == "required" {
-			required[sn.SnapName()] = true
+			required[sn.SnapName()] = sn
 		} else {
-			optional[sn.SnapName()] = true
+			optional[sn.SnapName()] = sn
 		}
 	}
 
-	// snapd should always be considered a part of the model. this is really a
-	// compatibility thing, and maybe should not be here since seed-refresh
-	// isn't gonna work on old models anyways.
-	required["snapd"] = true
+	// snapd is implicitly required, including models that don't list it
+	required["snapd"] = &asserts.ModelSnap{}
 
-	var optionalInSeed map[string]bool
+	var cachedOptionalContainers *seed.OptionalContainers
+	readOptionalContainers := func() (seed.OptionalContainers, error) {
+		if cachedOptionalContainers != nil {
+			return *cachedOptionalContainers, nil
+		}
 
-	return func(instanceName string) (bool, error) {
-		if required[instanceName] {
+		seedOCs, err := currentSeedOptionalContainers(st)
+		if err != nil {
+			return seed.OptionalContainers{}, err
+		}
+
+		cachedOptionalContainers = &seedOCs
+		return *cachedOptionalContainers, nil
+	}
+
+	triggers = func(instanceName string) (bool, error) {
+		if _, ok := required[instanceName]; ok {
 			return true, nil
 		}
 
-		if !optional[instanceName] {
+		if _, ok := optional[instanceName]; !ok {
 			return false, nil
 		}
 
-		if optionalInSeed == nil {
-			currentSystem, err := currentSeededSystem(st)
-			if err != nil {
-				return false, err
+		optionalInSeed, err := readOptionalContainers()
+		if err != nil {
+			return false, err
+		}
+
+		return strutil.ListContains(optionalInSeed.Snaps, instanceName), nil
+	}
+
+	allowlist = func() (SeedAllowlist, error) {
+		var snaps []string
+		components := make(map[string][]string)
+
+		for snapName, modelSnap := range required {
+			snaps = append(snaps, snapName)
+			for compName, comp := range modelSnap.Components {
+				if comp.Presence == "optional" {
+					optionalInSeed, err := readOptionalContainers()
+					if err != nil {
+						return SeedAllowlist{}, err
+					}
+
+					if !strutil.ListContains(optionalInSeed.Components[snapName], compName) {
+						continue
+					}
+				}
+
+				components[snapName] = append(components[snapName], compName)
 			}
+		}
 
-			current, err := seedOpen(dirs.SnapSeedDir, currentSystem.System)
+		for snapName, modelSnap := range optional {
+			optionalInSeed, err := readOptionalContainers()
 			if err != nil {
-				return false, err
-			}
-			if err := current.LoadAssertions(nil, nil); err != nil {
-				return false, err
+				return SeedAllowlist{}, err
 			}
 
-			copier, ok := current.(seed.Copier)
-			if !ok {
-				// this would only happen if the seed is pre-core20
-				return false, fmt.Errorf("internal error: seed %q does not support listing optional containers", currentSystem.System)
+			if !strutil.ListContains(optionalInSeed.Snaps, snapName) {
+				continue
 			}
 
-			oc, err := copier.OptionalContainers()
-			if err != nil {
-				return false, err
-			}
+			snaps = append(snaps, snapName)
+
+			for compName, comp := range modelSnap.Components {
+				if comp.Presence == "optional" && !strutil.ListContains(optionalInSeed.Components[snapName], compName) {
+					continue
+				}
 
-			optionalInSeed = make(map[string]bool, len(oc.Snaps))
-			for _, sn := range oc.Snaps {
-				optionalInSeed[sn] = true
+				components[snapName] = append(components[snapName], compName)
 			}
 		}
 
-		return optionalInSeed[instanceName], nil
+		sort.Strings(snaps)
+		for _, comps := range components {
+			sort.Strings(comps)
+		}
+
+		return SeedAllowlist{
+			Snaps:      snaps,
+			Components: components,
+		}, nil
 	}
+
+	return triggers, allowlist
+}
+
+func currentSeedOptionalContainers(st *state.State) (seed.OptionalContainers, error) {
+	currentSystem, err := currentSeededSystem(st)
+	if err != nil {
+		return seed.OptionalContainers{}, err
+	}
+
+	current, err := seedOpen(dirs.SnapSeedDir, currentSystem.System)
+	if err != nil {
+		return seed.OptionalContainers{}, err
+	}
+	if err := current.LoadAssertions(nil, nil); err != nil {
+		return seed.OptionalContainers{}, err
+	}
+
+	copier, ok := current.(seed.Copier)
+	if !ok {
+		return seed.OptionalContainers{}, fmt.Errorf("internal error: seed %q does not support listing optional containers", currentSystem.System)
+	}
+
+	return copier.OptionalContainers()
 }
 
 func findSeedRefreshTasks(ts *state.TaskSet) (*snapstate.SeedRefreshTasks, error) {
@@ -2074,6 +2158,7 @@ func createRecoverySystemTasks(st *state.State, label string, snapSetupTasks, co
 		ComponentSetupTasks: compSetupTasks,
 		LocalSnaps:          opts.LocalSnaps,
 		LocalComponents:     opts.LocalComponents,
+		Allowlist:           opts.Allowlist,
 		TestSystem:          opts.TestSystem,
 		MarkDefault:         opts.MarkDefault,
 		SeedRefresh:         opts.SeedRefresh,
@@ -2130,6 +2215,10 @@ type CreateRecoverySystemOptions struct {
 	// recovery system.
 	LocalComponents []snapstate.PathComponent
 
+	// Allowlist identifies the snaps and components that may be used to create
+	// the recovery system. A nil allowlist permits all containers.
+	Allowlist *SeedAllowlist
+
 	// TestSystem is set to true if the new recovery system should be verified
 	// by rebooting into the new system, prior to marking it as a valid recovery
 	// system. If false, the system will immediately be considered a valid
@@ -2149,6 +2238,13 @@ type CreateRecoverySystemOptions struct {
 	Offline bool
 }
 
+// SeedAllowlist identifies the snaps and components that may be used to
+// create a recovery system.
+type SeedAllowlist struct {
+	Snaps      []string            `json:"snaps,omitempty"`
+	Components map[string][]string `json:"components,omitempty"`
+}
+
 var ErrNoRecoverySystem = errors.New("recovery system does not exist")
 
 // RemoveRecoverySystem removes the recovery system with the given label. The