o/ifacestate: properly handle undo scenarios where auto-connections are dropped

overlord/ifacestate/handlers.go

@@ -337,6 +337,9 @@ func (d delayedEffectsForSnaps) EnqueueFor(snapName affectedSnap, backend interf
 	d[snapName][backend] = append(d[snapName][backend], item)
 }
 
+// refreshAppSetConnections refreshes repository connections for appSet and, on
+// the setup-profiles do path, records undo data for persisted connection state
+// that reloadConnections changed or dropped.
 func (m *InterfaceManager) refreshAppSetConnections(task *state.Task, appSet *interfaces.SnapAppSet) ([]string, []string, error) {
 	snapInfo := appSet.Info()
 	snapName := appSet.InstanceName()
@@ -369,11 +372,21 @@ func (m *InterfaceManager) refreshAppSetConnections(task *state.Task, appSet *in
 		task.Logf("%s", snap.BadInterfacesSummary(snapInfo))
 	}
 
-	affectedConnections, err := m.reloadConnections(snapName)
+	reloadedConns, changedOrDroppedConns, err := m.reloadConnections(snapName)
 	if err != nil {
 		return nil, nil, err
 	}
-	return disconnectedSnaps, affectedConnections, nil
+
+	// if this task modified any connection states, take a snapshot of the
+	// original connections so that setup-profiles' undo can restore them, if
+	// needed
+	if task.Status() != state.UndoingStatus {
+		if err := snapshotChangedConnectionsForUndo(task, snapName, changedOrDroppedConns); err != nil {
+			return nil, nil, err
+		}
+	}
+
+	return disconnectedSnaps, reloadedConns, nil
 }
 
 func (m *InterfaceManager) setupProfilesForAppSet(
@@ -385,7 +398,7 @@ func (m *InterfaceManager) setupProfilesForAppSet(
 	st := task.State()
 
 	snapName := appSet.InstanceName()
-	disconnectedSnaps, affectedConnections, err := m.refreshAppSetConnections(task, appSet)
+	disconnectedSnaps, reloadedConns, err := m.refreshAppSetConnections(task, appSet)
 	if err != nil {
 		return nil, err
 	}
@@ -399,7 +412,7 @@ func (m *InterfaceManager) setupProfilesForAppSet(
 	snapsWithConnectedSlots := make(map[string]bool)
 	newConnectedSnaps := make(map[string]bool)
 	// Identify affected snaps on either side of the connection.
-	for _, connID := range affectedConnections {
+	for _, connID := range reloadedConns {
 		connRef, err := interfaces.ParseConnRef(connID)
 		if err != nil {
 			return nil, fmt.Errorf("internal error: cannot parse existing connection: %w", err)
@@ -697,6 +710,13 @@ func (m *InterfaceManager) undoSetupProfiles(task *state.Task, tomb *tomb.Tomb)
 		return err
 	}
 
+	// restore any connection state snapshot saved by refreshAppSetConnections on
+	// the original setup-profiles do path before rebuilding profiles for the old
+	// revision
+	if err := restoreConnectionsForSetupProfiles(task); err != nil {
+		return err
+	}
+
 	// Get the name from SnapSetup and use it to find the current SideInfo
 	// about the snap, if there is one.
 	var snapst snapstate.SnapState
@@ -2003,7 +2023,7 @@ func (m *InterfaceManager) transitionConnectionsCoreMigration(st *state.State, o
 	// on disk are rewritten. This is ok because core/ubuntu-core have
 	// exactly the same profiles and nothing in the generated policies
 	// has the core snap-name encoded.
-	if _, err := m.reloadConnections(newName); err != nil {
+	if _, _, err := m.reloadConnections(newName); err != nil {
 		return err
 	}
 

overlord/ifacestate/helpers.go

@@ -426,15 +426,102 @@ func isBroken(st *state.State, snapName string) (bool, error) {
 	return false, nil
 }
 
+func cloneConnState(connState *schema.ConnState) *schema.ConnState {
+	clone := *connState
+
+	cloneAttrs := func(attrs map[string]any) map[string]any {
+		if attrs == nil {
+			return nil
+		}
+		return utils.CopyAttributes(attrs)
+	}
+
+	clone.StaticPlugAttrs = cloneAttrs(connState.StaticPlugAttrs)
+	clone.DynamicPlugAttrs = cloneAttrs(connState.DynamicPlugAttrs)
+	clone.StaticSlotAttrs = cloneAttrs(connState.StaticSlotAttrs)
+	clone.DynamicSlotAttrs = cloneAttrs(connState.DynamicSlotAttrs)
+
+	return &clone
+}
+
+// snapshotChangedConnectionsForUndo records original states for persisted
+// connections that setup-profiles changed or dropped so undo can restore them,
+// if needed.
+func snapshotChangedConnectionsForUndo(task *state.Task, instanceName string, changedConns map[string]*schema.ConnState) error {
+	if len(changedConns) == 0 {
+		return nil
+	}
+
+	// if this isn't the setup-profiles task that is going to handle the undo,
+	// then we don't need to keep track of these on the task
+	if !shouldUndoSetupProfiles(task, instanceName) {
+		return nil
+	}
+
+	var connectionSnapshot map[string]*schema.ConnState
+	err := task.Get("changed-or-dropped-connection-snapshot", &connectionSnapshot)
+	if err != nil && !errors.Is(err, state.ErrNoState) {
+		return err
+	}
+	if connectionSnapshot == nil {
+		connectionSnapshot = make(map[string]*schema.ConnState)
+	}
+
+	for connID, connState := range changedConns {
+		if connectionSnapshot[connID] != nil {
+			// a setup-profiles task can be retried after saving the connection
+			// states and unlocking for backend setup. keep the first snapshot.
+			continue
+		}
+		connectionSnapshot[connID] = connState
+	}
+
+	task.Set("changed-or-dropped-connection-snapshot", connectionSnapshot)
+
+	return nil
+}
+
+// restoreConnectionsForSetupProfiles restores connection states saved by
+// snapshotChangedConnectionsForUndo on a setup-profiles task.
+func restoreConnectionsForSetupProfiles(task *state.Task) error {
+	var connectionSnapshot map[string]*schema.ConnState
+	err := task.Get("changed-or-dropped-connection-snapshot", &connectionSnapshot)
+	if errors.Is(err, state.ErrNoState) {
+		return nil
+	}
+	if err != nil {
+		return err
+	}
+
+	st := task.State()
+
+	conns, err := getConns(st)
+	if err != nil {
+		return err
+	}
+
+	for connID, connState := range connectionSnapshot {
+		conns[connID] = connState
+	}
+	setConns(st, conns)
+
+	return nil
+}
+
 // reloadConnections reloads connections stored in the state in the repository.
 // Using non-empty snapName the operation can be scoped to connections
 // affecting a given snap.
 //
-// The return value is the list of affected snap names and their connection IDs.
-func (m *InterfaceManager) reloadConnections(snapName string) (reloadedConnectionIDs []string, err error) {
+// The return value is the list of reloaded connection IDs, plus the original
+// connection states whose persisted state was changed or dropped.
+func (m *InterfaceManager) reloadConnections(snapName string) (
+	reloadedConnectionIDs []string,
+	changedOrDroppedConns map[string]*schema.ConnState,
+	err error,
+) {
 	conns, err := getConns(m.state)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	var policyChecker interfaces.PolicyFunc
@@ -445,21 +532,22 @@ func (m *InterfaceManager) reloadConnections(snapName string) (reloadedConnectio
 	if errors.Is(err, state.ErrNoState) {
 		// everything else is a noop, as no model means no connections
 		// to reload
-		return nil, nil
+		return nil, nil, nil
 	} else if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	autoChecker, err = newAutoConnectChecker(m.state, m.repo, deviceCtx)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	connChecker, err = newConnectChecker(m.state, deviceCtx)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	connStateChanged := false
+	changedOrDroppedConns = make(map[string]*schema.ConnState)
 
 	var reloadedConnections []string
 ConnsLoop:
@@ -473,7 +561,7 @@ ConnsLoop:
 		}
 		connRef, err := interfaces.ParseConnRef(connId)
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 		// Apply filtering, this allows us to reload only a subset of
 		// connections (and similarly, refresh the static attributes of only a
@@ -497,13 +585,14 @@ ConnsLoop:
 				for _, snapName := range []string{connRef.PlugRef.Snap, connRef.SlotRef.Snap} {
 					broken, err := isBroken(m.state, snapName)
 					if err != nil {
-						return nil, err
+						return nil, nil, err
 					}
 					if broken {
 						logger.Noticef("Snap %q is broken, ignored by reloadConnections", snapName)
 						continue ConnsLoop
 					}
 				}
+				changedOrDroppedConns[connId] = cloneConnState(connState)
 				delete(conns, connId)
 				connStateChanged = true
 			}
@@ -540,11 +629,11 @@ ConnsLoop:
 
 		plugAppSet, err := interfaces.NewSnapAppSet(plugInfo.Snap, nil)
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 		slotAppSet, err := interfaces.NewSnapAppSet(slotInfo.Snap, nil)
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 
 		cplug := interfaces.NewConnectedPlug(plugInfo, plugAppSet, newStaticPlugAttrs, connState.DynamicPlugAttrs)
@@ -568,6 +657,7 @@ ConnsLoop:
 			reloadedConnections = append(reloadedConnections, connId)
 
 			if updateStaticAttrs {
+				changedOrDroppedConns[connId] = cloneConnState(connState)
 				connState.StaticPlugAttrs = staticPlugAttrs
 				connState.StaticSlotAttrs = staticSlotAttrs
 				connStateChanged = true
@@ -578,7 +668,7 @@ ConnsLoop:
 		setConns(m.state, conns)
 	}
 
-	return reloadedConnections, nil
+	return reloadedConnections, changedOrDroppedConns, nil
 }
 
 // removeConnections disconnects all connections of the snap in the repo. It should only be used if the snap

overlord/ifacestate/ifacemgr.go

@@ -206,7 +206,7 @@ func (m *InterfaceManager) StartUp() error {
 	if err := removeStaleConnections(m.state); err != nil {
 		return err
 	}
-	if _, err := m.reloadConnections(""); err != nil {
+	if _, _, err := m.reloadConnections(""); err != nil {
 		return err
 	}