Maintain a RevertStatus map in snapstate and allow Revert() with

NotBlocked flag - such revision isn't considered as blocked in Block() logic for refreshes. 

6 files changed, 263 insertions, 6 deletions

diff --git a/overlord/snapstate/flags.go b/overlord/snapstate/flags.go
index b117a4ff1d..e5a92410c8 100644
--- a/overlord/snapstate/flags.go
+++ b/overlord/snapstate/flags.go
@@ -35,6 +35,8 @@ type Flags struct {
 
	// Revert flags the SnapSetup as coming from a revert
	Revert bool `json:"revert,omitempty"`
+	// If reverting, set this status for the reverted revision.
+	RevertStatus RevertStatus `json:"revert-status,omitempty"`
 
	// RemoveSnapPath is used via InstallPath to flag that the file passed in is
	// temporary and should be removed
diff --git a/overlord/snapstate/handlers.go b/overlord/snapstate/handlers.go
index fec976c7e1..8e6ce72023 100644
--- a/overlord/snapstate/handlers.go
+++ b/overlord/snapstate/handlers.go
@@ -1342,6 +1342,21 @@ func (m *SnapManager) doLinkSnap(t *state.Task, _ *tomb.Tomb) (err error) {
	snapst.Sequence[len(snapst.Sequence)-1] = cand
	}
 
+	// update snapst.RevertStatus map if needed
+	if snapsup.Revert {
+	switch snapsup.RevertStatus {
+	case NotBlocked:
+	if snapst.RevertStatus == nil {
+	snapst.RevertStatus = make(map[int]RevertStatus)
+	}
+	snapst.RevertStatus[snapst.Current.N] = NotBlocked
+	default:
+	delete(snapst.RevertStatus, snapst.Current.N)
+	}
+	} else {
+	delete(snapst.RevertStatus, cand.Revision.N)
+	}
+
	oldCurrent := snapst.Current
	snapst.Current = cand.Revision
	snapst.Active = true
@@ -2392,6 +2407,9 @@ func (m *SnapManager) doDiscardSnap(t *state.Task, _ *tomb.Tomb) error {
	return fmt.Errorf("internal error: cannot discard snap %q: still active", snapsup.InstanceName())
	}
 
+	// drop any potential revert status for this revision
+	delete(snapst.RevertStatus, snapsup.Revision().N)
+
	if len(snapst.Sequence) == 1 {
	snapst.Sequence = nil
	snapst.Current = snap.Revision{}
diff --git a/overlord/snapstate/snapmgr.go b/overlord/snapstate/snapmgr.go
index 22e51a2032..ae48c7476f 100644
--- a/overlord/snapstate/snapmgr.go
+++ b/overlord/snapstate/snapmgr.go
@@ -131,11 +131,25 @@ func (snapsup *SnapSetup) MountFile() string {
	return snap.MountFile(snapsup.InstanceName(), snapsup.Revision())
 }
 
+// RevertStatus is a status of a snap revert; anything other than DefaultStatus
+// denotes a reverted snap revision that needs special handling in terms of
+// refresh blocking.
+type RevertStatus int
+
+const (
+	DefaultStatus RevertStatus = iota
+	NotBlocked
+)
+
 // SnapState holds the state for a snap installed in the system.
 type SnapState struct {
	SnapType string `json:"type"` // Use Type and SetType
	Sequence []*snap.SideInfo `json:"sequence"`
-	Active bool `json:"active,omitempty"`
+
+	// RevertStatus maps revisions to RevertStatus for revisions that
+	// need special handling in Block().
+	RevertStatus map[int]RevertStatus `json:"revert-status,omitempty"`
+	Active bool `json:"active,omitempty"`
 
	// LastActiveDisabledServices is a list of services that were disabled in
	// this snap when it was last active - i.e. when it was disabled, before
@@ -264,16 +278,23 @@ func (snapst *SnapState) LastIndex(revision snap.Revision) int {
 }
 
 // Block returns revisions that should be blocked on refreshes,
-// computed from Sequence[currentRevisionIndex+1:].
+// computed from Sequence[currentRevisionIndex+1:] and considering
+// special casing resulting from snapst.RevertStatus map.
 func (snapst *SnapState) Block() []snap.Revision {
-	// return revisions from Sequence[currentIndex:]
+	// return revisions from Sequence[currentIndex:], potentially excluding
+	// some of them based on RevertStatus.
	currentIndex := snapst.LastIndex(snapst.Current)
	if currentIndex < 0 || currentIndex+1 == len(snapst.Sequence) {
	return nil
	}
-	out := make([]snap.Revision, len(snapst.Sequence)-currentIndex-1)
-	for i, si := range snapst.Sequence[currentIndex+1:] {
-	out[i] = si.Revision
+	out := []snap.Revision{}
+	for _, si := range snapst.Sequence[currentIndex+1:] {
+	if status, ok := snapst.RevertStatus[si.Revision.N]; ok {
+	if status == NotBlocked {
+	continue
+	}
+	}
+	out = append(out, si.Revision)
	}
	return out
 }
diff --git a/overlord/snapstate/snapstate_remove_test.go b/overlord/snapstate/snapstate_remove_test.go
index 69c9f43c04..fcca0ffa86 100644
--- a/overlord/snapstate/snapstate_remove_test.go
+++ b/overlord/snapstate/snapstate_remove_test.go
@@ -871,6 +871,57 @@ func (s *snapmgrTestSuite) TestRemoveOneRevisionRunThrough(c *C) {
	c.Check(snapst.Sequence, HasLen, 2)
 }
 
+func (s *snapmgrTestSuite) TestRemoveOneRevisionDropsRevertStatus(c *C) {
+	si3 := snap.SideInfo{
+	RealName: "some-snap",
+	Revision: snap.R(3),
+	}
+
+	si5 := snap.SideInfo{
+	RealName: "some-snap",
+	Revision: snap.R(5),
+	}
+
+	si7 := snap.SideInfo{
+	RealName: "some-snap",
+	Revision: snap.R(7),
+	}
+
+	s.state.Lock()
+	defer s.state.Unlock()
+
+	snapstate.Set(s.state, "some-snap", &snapstate.SnapState{
+	Active: true,
+	Sequence: []*snap.SideInfo{&si5, &si3, &si7},
+	Current: si7.Revision,
+	RevertStatus: map[int]snapstate.RevertStatus{
+	3: snapstate.NotBlocked,
+	5: snapstate.NotBlocked,
+	},
+	SnapType: "app",
+	})
+
+	chg := s.state.NewChange("remove", "remove a snap")
+	ts, err := snapstate.Remove(s.state, "some-snap", snap.R(3), nil)
+	c.Assert(err, IsNil)
+	chg.AddAll(ts)
+
+	s.state.Unlock()
+	defer s.se.Stop()
+	s.settle(c)
+	s.state.Lock()
+
+	// verify snaps in the system state
+	var snapst snapstate.SnapState
+	err = snapstate.Get(s.state, "some-snap", &snapst)
+	c.Assert(err, IsNil)
+	c.Check(snapst.Sequence, HasLen, 2)
+	// revert status of revision 3 got dropped
+	c.Check(snapst.RevertStatus, DeepEquals, map[int]snapstate.RevertStatus{
+	5: snapstate.NotBlocked,
+	})
+}
+
 func (s *snapmgrTestSuite) TestRemoveLastRevisionRunThrough(c *C) {
	si := snap.SideInfo{
	RealName: "some-snap",
diff --git a/overlord/snapstate/snapstate_test.go b/overlord/snapstate/snapstate_test.go
index 4e9c335365..b888a79903 100644
--- a/overlord/snapstate/snapstate_test.go
+++ b/overlord/snapstate/snapstate_test.go
@@ -1525,9 +1525,68 @@ func (s *snapmgrTestSuite) TestRevertRunThrough(c *C) {
	Channel: "",
	Revision: snap.R(7),
	})
+	c.Check(snapst.RevertStatus, HasLen, 0)
	c.Assert(snapst.Block(), DeepEquals, []snap.Revision{snap.R(7)})
 }
 
+func (s *snapmgrTestSuite) TestRevertRevisionNotBlocked(c *C) {
+	si := snap.SideInfo{
+	RealName: "some-snap",
+	Revision: snap.R(7),
+	}
+	siOld := snap.SideInfo{
+	RealName: "some-snap",
+	Revision: snap.R(2),
+	}
+
+	s.state.Lock()
+	defer s.state.Unlock()
+
+	snapstate.Set(s.state, "some-snap", &snapstate.SnapState{
+	Active: true,
+	SnapType: "app",
+	Sequence: []*snap.SideInfo{&siOld, &si},
+	Current: si.Revision,
+	})
+
+	chg := s.state.NewChange("revert", "revert a snap backwards")
+	ts, err := snapstate.Revert(s.state, "some-snap", snapstate.Flags{RevertStatus: snapstate.NotBlocked})
+	c.Assert(err, IsNil)
+	chg.AddAll(ts)
+
+	s.state.Unlock()
+	defer s.se.Stop()
+	s.settle(c)
+	s.state.Lock()
+
+	// verify that the R(2) version is active now and R(7) is still there
+	var snapst snapstate.SnapState
+	err = snapstate.Get(s.state, "some-snap", &snapst)
+	c.Assert(err, IsNil)
+
+	// last refresh time shouldn't be modified on revert.
+	c.Check(snapst.LastRefreshTime, IsNil)
+	c.Assert(snapst.Active, Equals, true)
+	c.Assert(snapst.Current, Equals, snap.R(2))
+	c.Assert(snapst.Sequence, HasLen, 2)
+	c.Assert(snapst.Sequence[0], DeepEquals, &snap.SideInfo{
+	RealName: "some-snap",
+	Channel: "",
+	Revision: snap.R(2),
+	})
+	c.Assert(snapst.Sequence[1], DeepEquals, &snap.SideInfo{
+	RealName: "some-snap",
+	Channel: "",
+	Revision: snap.R(7),
+	})
+	// we have reverted from rev 7 to rev 2, but rev 7 is marked as not blocked
+	// due to revert.
+	c.Check(snapst.RevertStatus, DeepEquals, map[int]snapstate.RevertStatus{
+	7: snapstate.NotBlocked,
+	})
+	c.Assert(snapst.Block(), HasLen, 0)
+}
+
 func (s *snapmgrTestSuite) TestRevertWithBaseRunThrough(c *C) {
	si := snap.SideInfo{
	RealName: "some-snap-with-base",
diff --git a/overlord/snapstate/snapstate_update_test.go b/overlord/snapstate/snapstate_update_test.go
index 48ad6a14e7..195905fe8c 100644
--- a/overlord/snapstate/snapstate_update_test.go
+++ b/overlord/snapstate/snapstate_update_test.go
@@ -1006,6 +1006,68 @@ func (s *snapmgrTestSuite) TestUpdateRunThrough(c *C) {
	c.Check(snapstate.AuxStoreInfoFilename("services-snap-id"), testutil.FilePresent)
 }
 
+func (s *snapmgrTestSuite) TestUpdateDropsRevertStatus(c *C) {
+	si := snap.SideInfo{
+	RealName: "services-snap",
+	Revision: snap.R(7),
+	SnapID: "services-snap-id",
+	}
+	snaptest.MockSnap(c, `name: services-snap`, &si)
+
+	s.state.Lock()
+	defer s.state.Unlock()
+
+	si2 := snap.SideInfo{
+	RealName: "services-snap",
+	Revision: snap.R(11),
+	SnapID: "services-snap-id",
+	}
+	snapstate.Set(s.state, "services-snap", &snapstate.SnapState{
+	Active: true,
+	Sequence: []*snap.SideInfo{&si, &si2},
+	Current: si.Revision,
+	RevertStatus: map[int]snapstate.RevertStatus{
+	11: snapstate.NotBlocked,
+	},
+	SnapType: "app",
+	TrackingChannel: "latest/stable",
+	CohortKey: "embattled",
+	})
+
+	chg := s.state.NewChange("refresh", "refresh a snap")
+	ts, err := snapstate.Update(s.state, "services-snap", &snapstate.RevisionOptions{
+	Channel: "some-channel",
+	CohortKey: "some-cohort",
+	}, s.user.ID, snapstate.Flags{})
+	c.Assert(err, IsNil)
+	chg.AddAll(ts)
+
+	s.state.Unlock()
+	defer s.se.Stop()
+	s.settle(c)
+	s.state.Lock()
+
+	// verify snaps in the system state
+	var snapst snapstate.SnapState
+	c.Assert(snapstate.Get(s.state, "services-snap", &snapst), IsNil)
+	c.Assert(snapst.Active, Equals, true)
+	c.Assert(snapst.Current, Equals, snap.R(11))
+	c.Assert(snapst.Sequence, HasLen, 2)
+	c.Assert(snapst.Sequence[0], DeepEquals, &snap.SideInfo{
+	RealName: "services-snap",
+	SnapID: "services-snap-id",
+	Channel: "",
+	Revision: snap.R(7),
+	})
+	c.Assert(snapst.Sequence[1], DeepEquals, &snap.SideInfo{
+	RealName: "services-snap",
+	Channel: "some-channel",
+	SnapID: "services-snap-id",
+	Revision: snap.R(11),
+	})
+	c.Check(snapst.RevertStatus, HasLen, 0)
+}
+
 func (s *snapmgrTestSuite) TestUpdateResetsHoldState(c *C) {
	si := snap.SideInfo{
	RealName: "some-snap",
@@ -3438,6 +3500,50 @@ func (s *snapmgrTestSuite) TestAllUpdateBlockedRevision(c *C) {
	})
 }
 
+func (s *snapmgrTestSuite) TestAllUpdateRevisionNotBlocked(c *C) {
+	// update-all *should* set the block list
+	si7 := snap.SideInfo{
+	RealName: "some-snap",
+	SnapID: "some-snap-id",
+	Revision: snap.R(7),
+	}
+	si11 := snap.SideInfo{
+	RealName: "some-snap",
+	SnapID: "some-snap-id",
+	Revision: snap.R(11),
+	}
+
+	s.state.Lock()
+	defer s.state.Unlock()
+
+	snapstate.Set(s.state, "some-snap", &snapstate.SnapState{
+	Active: true,
+	Sequence: []*snap.SideInfo{&si7, &si11},
+	Current: si7.Revision,
+	RevertStatus: map[int]snapstate.RevertStatus{
+	si7.Revision.N: snapstate.NotBlocked,
+	},
+	})
+
+	updates, _, err := snapstate.UpdateMany(context.Background(), s.state, nil, s.user.ID, nil)
+	c.Check(err, IsNil)
+	c.Check(updates, HasLen, 0)
+
+	c.Assert(s.fakeBackend.ops, HasLen, 2)
+	c.Check(s.fakeBackend.ops[0], DeepEquals, fakeOp{
+	op: "storesvc-snap-action",
+	curSnaps: []store.CurrentSnap{{
+	InstanceName: "some-snap",
+	SnapID: "some-snap-id",
+	Revision: snap.R(7),
+	RefreshedDate: fakeRevDateEpoch.AddDate(0, 0, 7),
+	Block: []snap.Revision{snap.R(11)},
+	Epoch: snap.E("1*"),
+	}},
+	userID: 1,
+	})
+}
+
 func (s *snapmgrTestSuite) TestUpdateManyPartialFailureCheckRerefreshDone(c *C) {
	s.state.Lock()
	defer s.state.Unlock()