session: pass contexts through to all IDToGroupIndex methods

Author: ellemouton

firewalldb/actions.go

@@ -391,7 +391,7 @@ func (db *DB) ListSessionActions(sessionID session.ID,
 // pass the filterFn requirements.
 //
 // TODO: update to allow for pagination.
-func (db *DB) ListGroupActions(_ context.Context, groupID session.ID,
+func (db *DB) ListGroupActions(ctx context.Context, groupID session.ID,
 	filterFn ListActionsFilterFn) ([]*Action, error) {
 
 	if filterFn == nil {
@@ -400,7 +400,7 @@ func (db *DB) ListGroupActions(_ context.Context, groupID session.ID,
 		}
 	}
 
-	sessionIDs, err := db.sessionIDIndex.GetSessionIDs(groupID)
+	sessionIDs, err := db.sessionIDIndex.GetSessionIDs(ctx, groupID)
 	if err != nil {
 		return nil, err
 	}

firewalldb/mock.go

@@ -34,7 +34,9 @@ func (m *mockSessionDB) AddPair(sessionID, groupID session.ID) {
 }
 
 // GetGroupID returns the group ID for the given session ID.
-func (m *mockSessionDB) GetGroupID(sessionID session.ID) (session.ID, error) {
+func (m *mockSessionDB) GetGroupID(_ context.Context, sessionID session.ID) (
+	session.ID, error) {
+
 	id, ok := m.sessionToGroupID[sessionID]
 	if !ok {
 		return session.ID{}, fmt.Errorf("no group ID found for " +
@@ -45,7 +47,9 @@ func (m *mockSessionDB) GetGroupID(sessionID session.ID) (session.ID, error) {
 }
 
 // GetSessionIDs returns the set of session IDs that are in the group
-func (m *mockSessionDB) GetSessionIDs(groupID session.ID) ([]session.ID, error) {
+func (m *mockSessionDB) GetSessionIDs(_ context.Context, groupID session.ID) (
+	[]session.ID, error) {
+
 	ids, ok := m.groupToSessionIDs[groupID]
 	if !ok {
 		return nil, fmt.Errorf("no session IDs found for group ID")

session/interface.go

@@ -261,11 +261,11 @@ func WithMacaroonRecipe(caveats []macaroon.Caveat, perms []bakery.Op) Option {
 // IDToGroupIndex defines an interface for the session ID to group ID index.
 type IDToGroupIndex interface {
 	// GetGroupID will return the group ID for the given session ID.
-	GetGroupID(sessionID ID) (ID, error)
+	GetGroupID(ctx context.Context, sessionID ID) (ID, error)
 
 	// GetSessionIDs will return the set of session IDs that are in the
 	// group with the given ID.
-	GetSessionIDs(groupID ID) ([]ID, error)
+	GetSessionIDs(ctx context.Context, groupID ID) ([]ID, error)
 }
 
 // Store is the interface a persistent storage must implement for storing and

session/kvdb_store.go

@@ -624,7 +624,7 @@ func getUnusedIDAndKeyPair(bucket *bbolt.Bucket) (ID, *btcec.PrivateKey,
 // GetGroupID will return the group ID for the given session ID.
 //
 // NOTE: this is part of the IDToGroupIndex interface.
-func (db *BoltStore) GetGroupID(sessionID ID) (ID, error) {
+func (db *BoltStore) GetGroupID(_ context.Context, sessionID ID) (ID, error) {
 	var groupID ID
 	err := db.View(func(tx *bbolt.Tx) error {
 		sessionBkt, err := getBucket(tx, sessionBucketKey)
@@ -664,7 +664,9 @@ func (db *BoltStore) GetGroupID(sessionID ID) (ID, error) {
 // group with the given ID.
 //
 // NOTE: this is part of the IDToGroupIndex interface.
-func (db *BoltStore) GetSessionIDs(groupID ID) ([]ID, error) {
+func (db *BoltStore) GetSessionIDs(_ context.Context, groupID ID) ([]ID,
+	error) {
+
 	var (
 		sessionIDs []ID
 		err        error

session/store_test.go

@@ -170,11 +170,11 @@ func TestBasicSessionStore(t *testing.T) {
 
 	// Show that the group ID/session ID index has also been populated with
 	// this session.
-	groupID, err := db.GetGroupID(s4.ID)
+	groupID, err := db.GetGroupID(ctx, s4.ID)
 	require.NoError(t, err)
 	require.Equal(t, s1.ID, groupID)
 
-	sessIDs, err := db.GetSessionIDs(s4.GroupID)
+	sessIDs, err := db.GetSessionIDs(ctx, s4.GroupID)
 	require.NoError(t, err)
 	require.ElementsMatch(t, []ID{s4.ID, s1.ID}, sessIDs)
 
@@ -186,11 +186,11 @@ func TestBasicSessionStore(t *testing.T) {
 	require.NoError(t, err)
 	require.Empty(t, sessions)
 
-	_, err = db.GetGroupID(s4.ID)
+	_, err = db.GetGroupID(ctx, s4.ID)
 	require.ErrorIs(t, err, ErrUnknownGroup)
 
 	// Only session 1 should remain in this group.
-	sessIDs, err = db.GetSessionIDs(s4.GroupID)
+	sessIDs, err = db.GetSessionIDs(ctx, s4.GroupID)
 	require.NoError(t, err)
 	require.ElementsMatch(t, []ID{s1.ID}, sessIDs)
 }
@@ -235,6 +235,7 @@ func TestLinkingSessions(t *testing.T) {
 // of the GetGroupID and GetSessionIDs methods.
 func TestLinkedSessions(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	// Set up a new DB.
 	clock := clock.NewTestClock(testTime)
@@ -254,14 +255,14 @@ func TestLinkedSessions(t *testing.T) {
 
 	// Assert that the session ID to group ID index works as expected.
 	for _, s := range []*Session{s1, s2, s3} {
-		groupID, err := db.GetGroupID(s.ID)
+		groupID, err := db.GetGroupID(ctx, s.ID)
 		require.NoError(t, err)
 		require.Equal(t, s1.ID, groupID)
 		require.Equal(t, s.GroupID, groupID)
 	}
 
 	// Assert that the group ID to session ID index works as expected.
-	sIDs, err := db.GetSessionIDs(s1.GroupID)
+	sIDs, err := db.GetSessionIDs(ctx, s1.GroupID)
 	require.NoError(t, err)
 	require.EqualValues(t, []ID{s1.ID, s2.ID, s3.ID}, sIDs)
 
@@ -274,14 +275,14 @@ func TestLinkedSessions(t *testing.T) {
 
 	// Assert that the session ID to group ID index works as expected.
 	for _, s := range []*Session{s4, s5} {
-		groupID, err := db.GetGroupID(s.ID)
+		groupID, err := db.GetGroupID(ctx, s.ID)
 		require.NoError(t, err)
 		require.Equal(t, s4.ID, groupID)
 		require.Equal(t, s.GroupID, groupID)
 	}
 
 	// Assert that the group ID to session ID index works as expected.
-	sIDs, err = db.GetSessionIDs(s5.GroupID)
+	sIDs, err = db.GetSessionIDs(ctx, s5.GroupID)
 	require.NoError(t, err)
 	require.EqualValues(t, []ID{s4.ID, s5.ID}, sIDs)
 }

session_rpcserver.go

@@ -587,7 +587,7 @@ func (s *sessionRpcServer) RevokeSession(ctx context.Context,
 
 // PrivacyMapConversion can be used map real values to their pseudo counterpart
 // and vice versa.
-func (s *sessionRpcServer) PrivacyMapConversion(_ context.Context,
+func (s *sessionRpcServer) PrivacyMapConversion(ctx context.Context,
 	req *litrpc.PrivacyMapConversionRequest) (
 	*litrpc.PrivacyMapConversionResponse, error) {
 
@@ -606,7 +606,7 @@ func (s *sessionRpcServer) PrivacyMapConversion(_ context.Context,
 			return nil, err
 		}
 
-		groupID, err = s.cfg.db.GetGroupID(sessionID)
+		groupID, err = s.cfg.db.GetGroupID(ctx, sessionID)
 		if err != nil {
 			return nil, err
 		}