session: add ShiftState method to session Store

And only allow legal state shifts.

Author: ellemouton

session/interface.go

@@ -27,9 +27,9 @@ const (
 type State uint8
 
 /*
-		/---> StateExpired
+		/---> StateExpired (terminal)
 StateCreated ---
-       		\---> StateRevoked
+       		\---> StateRevoked (terminal)
 */
 
 const (
@@ -59,6 +59,24 @@ const (
 	StateReserved State = 4
 )
 
+// Terminal returns true if the state is a terminal state.
+func (s State) Terminal() bool {
+	return s == StateExpired || s == StateRevoked
+}
+
+// legalStateShifts is a map that defines the legal State transitions that a
+// Session can be put through.
+var legalStateShifts = map[State]map[State]bool{
+	StateCreated: {
+		StateExpired: true,
+		StateRevoked: true,
+	},
+	StateInUse: {
+		StateRevoked: true,
+		StateExpired: true,
+	},
+}
+
 // MacaroonRecipe defines the permissions and caveats that should be used
 // to bake a macaroon.
 type MacaroonRecipe struct {
@@ -225,5 +243,9 @@ type Store interface {
 	// StateReserved state.
 	DeleteReservedSessions() error
 
+	// ShiftState updates the state of the session with the given ID to the
+	// "dest" state.
+	ShiftState(id ID, dest State) error
+
 	IDToGroupIndex
 }

session/kvdb_store.go

@@ -514,6 +514,47 @@ func (db *BoltStore) DeleteReservedSessions() error {
 	})
 }
 
+// ShiftState updates the state of the session with the given ID to the "dest"
+// state.
+//
+// NOTE: this is part of the Store interface.
+func (db *BoltStore) ShiftState(id ID, dest State) error {
+	return db.Update(func(tx *bbolt.Tx) error {
+		sessionBucket, err := getBucket(tx, sessionBucketKey)
+		if err != nil {
+			return err
+		}
+
+		session, err := getSessionByID(sessionBucket, id)
+		if err != nil {
+			return err
+		}
+
+		// If the session is already in the desired state, we return
+		// with no error to maintain idempotency.
+		if session.State == dest {
+			return nil
+		}
+
+		// Ensure that the wanted state change is allowed.
+		allowedDestinations, ok := legalStateShifts[session.State]
+		if !ok || !allowedDestinations[dest] {
+			return fmt.Errorf("illegal session state transition "+
+				"from %d to %d", session.State, dest)
+		}
+
+		session.State = dest
+
+		// If the session is terminal, we set the revoked at time to the
+		// current time.
+		if dest.Terminal() {
+			session.RevokedAt = db.clock.Now().UTC()
+		}
+
+		return putSession(sessionBucket, session)
+	})
+}
+
 // RevokeSession updates the state of the session with the given local
 // public key to be revoked.
 //

session/store_test.go

@@ -392,6 +392,53 @@ func TestCheckSessionGroupPredicate(t *testing.T) {
 	require.True(t, ok)
 }
 
+// TestStateShift tests that the ShiftState method works as expected.
+func TestStateShift(t *testing.T) {
+	// Set up a new DB.
+	clock := clock.NewTestClock(testTime)
+	db, err := NewDB(t.TempDir(), "test.db", clock)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = db.Close()
+	})
+
+	// Add a new session to the DB.
+	s1 := newSession(t, db, clock, "label 1")
+	require.NoError(t, db.CreateSession(s1))
+
+	// Check that the session is in the StateCreated state. Also check that
+	// the "RevokedAt" time has not yet been set.
+	s1, err = db.GetSession(s1.LocalPublicKey)
+	require.NoError(t, err)
+	require.Equal(t, StateCreated, s1.State)
+	require.Equal(t, time.Time{}, s1.RevokedAt)
+
+	// Shift the state of the session to StateRevoked.
+	err = db.ShiftState(s1.ID, StateRevoked)
+	require.NoError(t, err)
+
+	// This should have worked. Since it is now in a terminal state, the
+	// "RevokedAt" time should be set.
+	s1, err = db.GetSession(s1.LocalPublicKey)
+	require.NoError(t, err)
+	require.Equal(t, StateRevoked, s1.State)
+	require.True(t, clock.Now().Equal(s1.RevokedAt))
+
+	// Trying to do the same state shift again should succeed since the
+	// session is already in the expected "dest" state. The revoked-at time
+	// should not have changed though.
+	prevTime := clock.Now()
+	clock.SetTime(prevTime.Add(time.Second))
+	err = db.ShiftState(s1.ID, StateRevoked)
+	require.NoError(t, err)
+	require.True(t, prevTime.Equal(s1.RevokedAt))
+
+	// Trying to shift the state from a terminal state back to StateCreated
+	// should also fail since this is not a legal state transition.
+	err = db.ShiftState(s1.ID, StateCreated)
+	require.ErrorContains(t, err, "illegal session state transition")
+}
+
 // testSessionModifier is a functional option that can be used to modify the
 // default test session created by newSession.
 type testSessionModifier func(*Session)