session: add context to various session List methods

This commit is contained in:
Elle Mouton 2025-02-25 13:19:34 +02:00
parent 310f11acda
commit 7c30c510a1
No known key found for this signature in database
GPG key ID: D7D916376026F177
4 changed files with 28 additions and 24 deletions

View file

@ -282,14 +282,15 @@ type Store interface {
GetSession(ctx context.Context, key *btcec.PublicKey) (*Session, error)
// ListAllSessions returns all sessions currently known to the store.
ListAllSessions() ([]*Session, error)
ListAllSessions(ctx context.Context) ([]*Session, error)
// ListSessionsByType returns all sessions of the given type.
ListSessionsByType(t Type) ([]*Session, error)
ListSessionsByType(ctx context.Context, t Type) ([]*Session, error)
// ListSessionsByState returns all sessions currently known to the store
// that are in the given states.
ListSessionsByState(...State) ([]*Session, error)
ListSessionsByState(ctx context.Context, state ...State) ([]*Session,
error)
// UpdateSessionRemotePubKey can be used to add the given remote pub key
// to the session with the given local pub key.

View file

@ -351,7 +351,7 @@ func (db *BoltStore) GetSession(_ context.Context, key *btcec.PublicKey) (
// ListAllSessions returns all sessions currently known to the store.
//
// NOTE: this is part of the Store interface.
func (db *BoltStore) ListAllSessions() ([]*Session, error) {
func (db *BoltStore) ListAllSessions(_ context.Context) ([]*Session, error) {
return db.listSessions(func(s *Session) bool {
return true
})
@ -361,7 +361,9 @@ func (db *BoltStore) ListAllSessions() ([]*Session, error) {
// have the given type.
//
// NOTE: this is part of the Store interface.
func (db *BoltStore) ListSessionsByType(t Type) ([]*Session, error) {
func (db *BoltStore) ListSessionsByType(_ context.Context, t Type) ([]*Session,
error) {
return db.listSessions(func(s *Session) bool {
return s.Type == t
})
@ -371,7 +373,9 @@ func (db *BoltStore) ListSessionsByType(t Type) ([]*Session, error) {
// are in the given states.
//
// NOTE: this is part of the Store interface.
func (db *BoltStore) ListSessionsByState(states ...State) ([]*Session, error) {
func (db *BoltStore) ListSessionsByState(_ context.Context, states ...State) (
[]*Session, error) {
return db.listSessions(func(s *Session) bool {
for _, state := range states {
if s.State == state {

View file

@ -58,18 +58,18 @@ func TestBasicSessionStore(t *testing.T) {
s3 := createSession(t, db, "session 3", withType(TypeAutopilot))
// Test the ListSessionsByType method.
sessions, err := db.ListSessionsByType(TypeMacaroonAdmin)
sessions, err := db.ListSessionsByType(ctx, TypeMacaroonAdmin)
require.NoError(t, err)
require.Equal(t, 2, len(sessions))
assertEqualSessions(t, s1, sessions[0])
assertEqualSessions(t, s2, sessions[1])
sessions, err = db.ListSessionsByType(TypeAutopilot)
sessions, err = db.ListSessionsByType(ctx, TypeAutopilot)
require.NoError(t, err)
require.Equal(t, 1, len(sessions))
assertEqualSessions(t, s3, sessions[0])
sessions, err = db.ListSessionsByType(TypeMacaroonReadonly)
sessions, err = db.ListSessionsByType(ctx, TypeMacaroonReadonly)
require.NoError(t, err)
require.Empty(t, sessions)
@ -113,7 +113,7 @@ func TestBasicSessionStore(t *testing.T) {
require.Equal(t, s1.State, StateRevoked)
// Test that ListAllSessions works.
sessions, err = db.ListAllSessions()
sessions, err = db.ListAllSessions(ctx)
require.NoError(t, err)
require.Equal(t, 3, len(sessions))
assertEqualSessions(t, s1, sessions[0])
@ -121,29 +121,29 @@ func TestBasicSessionStore(t *testing.T) {
assertEqualSessions(t, s3, sessions[2])
// Test that ListSessionsByState works.
sessions, err = db.ListSessionsByState(StateRevoked)
sessions, err = db.ListSessionsByState(ctx, StateRevoked)
require.NoError(t, err)
require.Equal(t, 1, len(sessions))
assertEqualSessions(t, s1, sessions[0])
sessions, err = db.ListSessionsByState(StateCreated)
sessions, err = db.ListSessionsByState(ctx, StateCreated)
require.NoError(t, err)
require.Equal(t, 2, len(sessions))
assertEqualSessions(t, s2, sessions[0])
assertEqualSessions(t, s3, sessions[1])
sessions, err = db.ListSessionsByState(StateCreated, StateRevoked)
sessions, err = db.ListSessionsByState(ctx, StateCreated, StateRevoked)
require.NoError(t, err)
require.Equal(t, 3, len(sessions))
assertEqualSessions(t, s1, sessions[0])
assertEqualSessions(t, s2, sessions[1])
assertEqualSessions(t, s3, sessions[2])
sessions, err = db.ListSessionsByState()
sessions, err = db.ListSessionsByState(ctx)
require.NoError(t, err)
require.Empty(t, sessions)
sessions, err = db.ListSessionsByState(StateReserved)
sessions, err = db.ListSessionsByState(ctx, StateReserved)
require.NoError(t, err)
require.Empty(t, sessions)
@ -153,7 +153,7 @@ func TestBasicSessionStore(t *testing.T) {
// of the sessions are reserved.
require.NoError(t, db.DeleteReservedSessions())
sessions, err = db.ListSessionsByState(StateReserved)
sessions, err = db.ListSessionsByState(ctx, StateReserved)
require.NoError(t, err)
require.Empty(t, sessions)
@ -163,7 +163,7 @@ func TestBasicSessionStore(t *testing.T) {
)
require.NoError(t, err)
sessions, err = db.ListSessionsByState(StateReserved)
sessions, err = db.ListSessionsByState(ctx, StateReserved)
require.NoError(t, err)
require.Equal(t, 1, len(sessions))
assertEqualSessions(t, s4, sessions[0])
@ -182,7 +182,7 @@ func TestBasicSessionStore(t *testing.T) {
// database and no longer in the group ID/session ID index.
require.NoError(t, db.DeleteReservedSessions())
sessions, err = db.ListSessionsByState(StateReserved)
sessions, err = db.ListSessionsByState(ctx, StateReserved)
require.NoError(t, err)
require.Empty(t, sessions)

View file

@ -103,8 +103,7 @@ func (s *sessionRpcServer) start(ctx context.Context) error {
// Start up all previously created sessions.
sessions, err := s.cfg.db.ListSessionsByState(
session.StateCreated,
session.StateInUse,
ctx, session.StateCreated, session.StateInUse,
)
if err != nil {
return fmt.Errorf("error listing sessions: %v", err)
@ -531,10 +530,10 @@ func (s *sessionRpcServer) resumeSession(ctx context.Context,
}
// ListSessions returns all sessions known to the session store.
func (s *sessionRpcServer) ListSessions(_ context.Context,
func (s *sessionRpcServer) ListSessions(ctx context.Context,
_ *litrpc.ListSessionsRequest) (*litrpc.ListSessionsResponse, error) {
sessions, err := s.cfg.db.ListAllSessions()
sessions, err := s.cfg.db.ListAllSessions(ctx)
if err != nil {
return nil, fmt.Errorf("error fetching sessions: %v", err)
}
@ -1270,11 +1269,11 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
// ListAutopilotSessions fetches and returns all the sessions from the DB that
// are of type TypeAutopilot.
func (s *sessionRpcServer) ListAutopilotSessions(_ context.Context,
func (s *sessionRpcServer) ListAutopilotSessions(ctx context.Context,
_ *litrpc.ListAutopilotSessionsRequest) (
*litrpc.ListAutopilotSessionsResponse, error) {
sessions, err := s.cfg.db.ListSessionsByType(session.TypeAutopilot)
sessions, err := s.cfg.db.ListSessionsByType(ctx, session.TypeAutopilot)
if err != nil {
return nil, fmt.Errorf("error fetching sessions: %v", err)
}