session: add context to ShiftState

Author: ellemouton

session/interface.go

@@ -306,7 +306,7 @@ type Store interface {
 
 	// ShiftState updates the state of the session with the given ID to the
 	// "dest" state.
-	ShiftState(id ID, dest State) error
+	ShiftState(ctx context.Context, id ID, dest State) error
 
 	IDToGroupIndex
 }

session/kvdb_store.go

@@ -524,7 +524,7 @@ func (db *BoltStore) DeleteReservedSessions(_ context.Context) error {
 // state.
 //
 // NOTE: this is part of the Store interface.
-func (db *BoltStore) ShiftState(id ID, dest State) error {
+func (db *BoltStore) ShiftState(_ context.Context, id ID, dest State) error {
 	return db.Update(func(tx *bbolt.Tx) error {
 		sessionBucket, err := getBucket(tx, sessionBucketKey)
 		if err != nil {

session/store_test.go

@@ -36,7 +36,7 @@ func TestBasicSessionStore(t *testing.T) {
 	require.Equal(t, StateReserved, s1.State)
 
 	// Move session 1 to the created state. This should succeed.
-	err = db.ShiftState(s1.ID, StateCreated)
+	err = db.ShiftState(ctx, s1.ID, StateCreated)
 	require.NoError(t, err)
 
 	// Show that the session is now in the created state.
@@ -46,7 +46,7 @@ func TestBasicSessionStore(t *testing.T) {
 
 	// Trying to move session 1 again should have no effect since it is
 	// already in the created state.
-	require.NoError(t, db.ShiftState(s1.ID, StateCreated))
+	require.NoError(t, db.ShiftState(ctx, s1.ID, StateCreated))
 
 	// Reserve and create a few more sessions. We increment the time by one
 	// second between each session to ensure that the created at time is
@@ -107,7 +107,7 @@ func TestBasicSessionStore(t *testing.T) {
 	require.Equal(t, session1.State, StateCreated)
 
 	// Now revoke the session and assert that the state is revoked.
-	require.NoError(t, db.ShiftState(s1.ID, StateRevoked))
+	require.NoError(t, db.ShiftState(ctx, s1.ID, StateRevoked))
 	s1, err = db.GetSession(ctx, s1.LocalPublicKey)
 	require.NoError(t, err)
 	require.Equal(t, s1.State, StateRevoked)
@@ -198,6 +198,7 @@ func TestBasicSessionStore(t *testing.T) {
 // TestLinkingSessions tests that session linking works as expected.
 func TestLinkingSessions(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	// Set up a new DB.
 	clock := clock.NewTestClock(testTime)
@@ -227,7 +228,7 @@ func TestLinkingSessions(t *testing.T) {
 	require.ErrorContains(t, err, "is still active")
 
 	// Revoke the first session.
-	require.NoError(t, db.ShiftState(s1.ID, StateRevoked))
+	require.NoError(t, db.ShiftState(ctx, s1.ID, StateRevoked))
 
 	// Persisting the second linked session should now work.
 	_, err = reserveSession(db, "session 2", withLinkedGroupID(&s1.GroupID))
@@ -255,10 +256,10 @@ func TestLinkedSessions(t *testing.T) {
 	// first session.
 	s1 := createSession(t, db, "session 1")
 
-	require.NoError(t, db.ShiftState(s1.ID, StateRevoked))
+	require.NoError(t, db.ShiftState(ctx, s1.ID, StateRevoked))
 	s2 := createSession(t, db, "session 2", withLinkedGroupID(&s1.GroupID))
 
-	require.NoError(t, db.ShiftState(s2.ID, StateRevoked))
+	require.NoError(t, db.ShiftState(ctx, s2.ID, StateRevoked))
 	s3 := createSession(t, db, "session 3", withLinkedGroupID(&s2.GroupID))
 
 	// Assert that the session ID to group ID index works as expected.
@@ -277,7 +278,7 @@ func TestLinkedSessions(t *testing.T) {
 	// To ensure that different groups don't interfere with each other,
 	// let's add another set of linked sessions not linked to the first.
 	s4 := createSession(t, db, "session 4")
-	require.NoError(t, db.ShiftState(s4.ID, StateRevoked))
+	require.NoError(t, db.ShiftState(ctx, s4.ID, StateRevoked))
 	s5 := createSession(t, db, "session 5", withLinkedGroupID(&s4.GroupID))
 	require.NotEqual(t, s4.GroupID, s1.GroupID)
 
@@ -319,7 +320,7 @@ func TestStateShift(t *testing.T) {
 	require.Equal(t, time.Time{}, s1.RevokedAt)
 
 	// Shift the state of the session to StateRevoked.
-	err = db.ShiftState(s1.ID, StateRevoked)
+	err = db.ShiftState(ctx, s1.ID, StateRevoked)
 	require.NoError(t, err)
 
 	// This should have worked. Since it is now in a terminal state, the
@@ -334,13 +335,13 @@ func TestStateShift(t *testing.T) {
 	// should not have changed though.
 	prevTime := clock.Now()
 	clock.SetTime(prevTime.Add(time.Second))
-	err = db.ShiftState(s1.ID, StateRevoked)
+	err = db.ShiftState(ctx, 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)
+	err = db.ShiftState(ctx, s1.ID, StateCreated)
 	require.ErrorContains(t, err, "illegal session state transition")
 }
 
@@ -396,7 +397,7 @@ func createSession(t *testing.T, db Store, label string,
 	s, err := reserveSession(db, label, mods...)
 	require.NoError(t, err)
 
-	err = db.ShiftState(s.ID, StateCreated)
+	err = db.ShiftState(context.Background(), s.ID, StateCreated)
 	require.NoError(t, err)
 
 	s, err = db.GetSessionByID(context.Background(), s.ID)

session_rpcserver.go

@@ -149,7 +149,7 @@ func (s *sessionRpcServer) start(ctx context.Context) error {
 
 				if perm {
 					err := s.cfg.db.ShiftState(
-						sess.ID, session.StateRevoked,
+						ctx, sess.ID, session.StateRevoked,
 					)
 					if err != nil {
 						log.Errorf("error revoking "+
@@ -323,7 +323,7 @@ func (s *sessionRpcServer) AddSession(ctx context.Context,
 		return nil, fmt.Errorf("error creating new session: %v", err)
 	}
 
-	err = s.cfg.db.ShiftState(sess.ID, session.StateCreated)
+	err = s.cfg.db.ShiftState(ctx, sess.ID, session.StateCreated)
 	if err != nil {
 		return nil, fmt.Errorf("error shifting session state to "+
 			"Created: %v", err)
@@ -362,7 +362,7 @@ func (s *sessionRpcServer) resumeSession(ctx context.Context,
 		log.Debugf("Not resuming session %x with expiry %s",
 			pubKeyBytes, sess.Expiry)
 
-		err := s.cfg.db.ShiftState(sess.ID, session.StateExpired)
+		err := s.cfg.db.ShiftState(ctx, sess.ID, session.StateExpired)
 		if err != nil {
 			return fmt.Errorf("error revoking session: %v", err)
 		}
@@ -440,7 +440,7 @@ func (s *sessionRpcServer) resumeSession(ctx context.Context,
 				"passed. Revoking session", pubKeyBytes)
 
 			return s.cfg.db.ShiftState(
-				sess.ID, session.StateRevoked,
+				ctx, sess.ID, session.StateRevoked,
 			)
 		}
 
@@ -520,7 +520,7 @@ func (s *sessionRpcServer) resumeSession(ctx context.Context,
 			log.Debugf("Error stopping session: %v", err)
 		}
 
-		err = s.cfg.db.ShiftState(sess.ID, session.StateRevoked)
+		err = s.cfg.db.ShiftState(ctx, sess.ID, session.StateRevoked)
 		if err != nil {
 			log.Debugf("error revoking session: %v", err)
 		}
@@ -567,7 +567,7 @@ func (s *sessionRpcServer) RevokeSession(ctx context.Context,
 		return nil, fmt.Errorf("error fetching session: %v", err)
 	}
 
-	err = s.cfg.db.ShiftState(sess.ID, session.StateRevoked)
+	err = s.cfg.db.ShiftState(ctx, sess.ID, session.StateRevoked)
 	if err != nil {
 		return nil, fmt.Errorf("error revoking session: %v", err)
 	}
@@ -1240,7 +1240,7 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
 
 	// We only activate the session if the Autopilot server registration
 	// was successful.
-	err = s.cfg.db.ShiftState(sess.ID, session.StateCreated)
+	err = s.cfg.db.ShiftState(ctx, sess.ID, session.StateCreated)
 	if err != nil {
 		return nil, fmt.Errorf("error shifting session state to "+
 			"Created: %v", err)