session+firewall: pass context to GetSessionByID

Author: ellemouton

firewall/privacy_mapper.go

@@ -190,7 +190,7 @@ func (p *PrivacyMapper) checkAndReplaceIncomingRequest(ctx context.Context,
 	uri string, req proto.Message, sessionID session.ID) (proto.Message,
 	error) {
 
-	session, err := p.sessionDB.GetSessionByID(sessionID)
+	session, err := p.sessionDB.GetSessionByID(ctx, sessionID)
 	if err != nil {
 		return nil, err
 	}
@@ -220,7 +220,7 @@ func (p *PrivacyMapper) checkAndReplaceIncomingRequest(ctx context.Context,
 func (p *PrivacyMapper) replaceOutgoingResponse(ctx context.Context, uri string,
 	resp proto.Message, sessionID session.ID) (proto.Message, error) {
 
-	session, err := p.sessionDB.GetSessionByID(sessionID)
+	session, err := p.sessionDB.GetSessionByID(ctx, sessionID)
 	if err != nil {
 		return nil, err
 	}

firewall/rule_enforcer.go

@@ -386,7 +386,7 @@ func (r *RuleEnforcer) initRule(ctx context.Context, reqID uint64, name string,
 		return nil, err
 	}
 
-	session, err := r.sessionDB.GetSessionByID(sessionID)
+	session, err := r.sessionDB.GetSessionByID(ctx, sessionID)
 	if err != nil {
 		return nil, err
 	}

firewalldb/interface.go

@@ -1,12 +1,16 @@
 package firewalldb
 
-import "github.com/lightninglabs/lightning-terminal/session"
+import (
+	"context"
+
+	"github.com/lightninglabs/lightning-terminal/session"
+)
 
 // SessionDB is an interface that abstracts the database operations needed for
 // the privacy mapper to function.
 type SessionDB interface {
 	session.IDToGroupIndex
 
 	// GetSessionByID returns the session for a specific id.
-	GetSessionByID(session.ID) (*session.Session, error)
+	GetSessionByID(context.Context, session.ID) (*session.Session, error)
 }

firewalldb/mock.go

@@ -1,6 +1,7 @@
 package firewalldb
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/lightninglabs/lightning-terminal/session"
@@ -54,8 +55,8 @@ func (m *mockSessionDB) GetSessionIDs(groupID session.ID) ([]session.ID, error)
 }
 
 // GetSessionByID returns the session for a specific id.
-func (m *mockSessionDB) GetSessionByID(sessionID session.ID) (*session.Session,
-	error) {
+func (m *mockSessionDB) GetSessionByID(_ context.Context,
+	sessionID session.ID) (*session.Session, error) {
 
 	s, ok := m.sessionToGroupID[sessionID]
 	if !ok {

session/interface.go

@@ -298,7 +298,7 @@ type Store interface {
 		remotePubKey *btcec.PublicKey) error
 
 	// GetSessionByID fetches the session with the given ID.
-	GetSessionByID(id ID) (*Session, error)
+	GetSessionByID(ctx context.Context, id ID) (*Session, error)
 
 	// DeleteReservedSessions deletes all sessions that are in the
 	// StateReserved state.

session/kvdb_store.go

@@ -569,7 +569,9 @@ func (db *BoltStore) ShiftState(id ID, dest State) error {
 // GetSessionByID fetches the session with the given ID.
 //
 // NOTE: this is part of the Store interface.
-func (db *BoltStore) GetSessionByID(id ID) (*Session, error) {
+func (db *BoltStore) GetSessionByID(_ context.Context, id ID) (*Session,
+	error) {
+
 	var session *Session
 	err := db.View(func(tx *bbolt.Tx) error {
 		sessionBucket, err := getBucket(tx, sessionBucketKey)

session/store_test.go

@@ -23,15 +23,15 @@ func TestBasicSessionStore(t *testing.T) {
 	db := NewTestDB(t, clock)
 
 	// Try fetch a session that doesn't exist yet.
-	_, err := db.GetSessionByID(ID{1, 3, 4, 4})
+	_, err := db.GetSessionByID(ctx, ID{1, 3, 4, 4})
 	require.ErrorIs(t, err, ErrSessionNotFound)
 
 	// Reserve a session. This should succeed.
 	s1, err := reserveSession(db, "session 1")
 	require.NoError(t, err)
 
 	// Show that the session starts in the reserved state.
-	s1, err = db.GetSessionByID(s1.ID)
+	s1, err = db.GetSessionByID(ctx, s1.ID)
 	require.NoError(t, err)
 	require.Equal(t, StateReserved, s1.State)
 
@@ -40,7 +40,7 @@ func TestBasicSessionStore(t *testing.T) {
 	require.NoError(t, err)
 
 	// Show that the session is now in the created state.
-	s1, err = db.GetSessionByID(s1.ID)
+	s1, err = db.GetSessionByID(ctx, s1.ID)
 	require.NoError(t, err)
 	require.Equal(t, StateCreated, s1.State)
 
@@ -80,7 +80,7 @@ func TestBasicSessionStore(t *testing.T) {
 		require.NoError(t, err)
 		assertEqualSessions(t, s, session)
 
-		session, err = db.GetSessionByID(s.ID)
+		session, err = db.GetSessionByID(ctx, s.ID)
 		require.NoError(t, err)
 		assertEqualSessions(t, s, session)
 	}
@@ -386,7 +386,7 @@ func createSession(t *testing.T, db Store, label string,
 	err = db.ShiftState(s.ID, StateCreated)
 	require.NoError(t, err)
 
-	s, err = db.GetSessionByID(s.ID)
+	s, err = db.GetSessionByID(context.Background(), s.ID)
 	require.NoError(t, err)
 
 	return s

session_rpcserver.go

@@ -335,7 +335,7 @@ func (s *sessionRpcServer) AddSession(ctx context.Context,
 
 	// Re-fetch the session to get the latest state of it before marshaling
 	// it.
-	sess, err = s.cfg.db.GetSessionByID(sess.ID)
+	sess, err = s.cfg.db.GetSessionByID(ctx, sess.ID)
 	if err != nil {
 		return nil, fmt.Errorf("error fetching session: %v", err)
 	}
@@ -867,7 +867,7 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
 		copy(groupID[:], req.LinkedGroupId)
 
 		// Check that the group actually does exist.
-		groupSess, err := s.cfg.db.GetSessionByID(groupID)
+		groupSess, err := s.cfg.db.GetSessionByID(ctx, groupID)
 		if err != nil {
 			return nil, err
 		}
@@ -1252,7 +1252,7 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
 
 	// Re-fetch the session to get the latest state of it before marshaling
 	// it.
-	sess, err = s.cfg.db.GetSessionByID(sess.ID)
+	sess, err = s.cfg.db.GetSessionByID(ctx, sess.ID)
 	if err != nil {
 		return nil, fmt.Errorf("error fetching session: %v", err)
 	}