diff --git a/firewall/privacy_mapper.go b/firewall/privacy_mapper.go index 7a2f8fe4..af4f3b0a 100644 --- a/firewall/privacy_mapper.go +++ b/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 } diff --git a/firewall/rule_enforcer.go b/firewall/rule_enforcer.go index 964baf32..008af72c 100644 --- a/firewall/rule_enforcer.go +++ b/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 } diff --git a/firewalldb/interface.go b/firewalldb/interface.go index 6e650957..86e638b5 100644 --- a/firewalldb/interface.go +++ b/firewalldb/interface.go @@ -1,6 +1,10 @@ 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. @@ -8,5 +12,5 @@ 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) } diff --git a/firewalldb/mock.go b/firewalldb/mock.go index 4030dde3..c13f0141 100644 --- a/firewalldb/mock.go +++ b/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 { diff --git a/session/interface.go b/session/interface.go index 8293cf71..bde8dc73 100644 --- a/session/interface.go +++ b/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. diff --git a/session/kvdb_store.go b/session/kvdb_store.go index 0c415ae6..8953a98f 100644 --- a/session/kvdb_store.go +++ b/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) diff --git a/session/store_test.go b/session/store_test.go index 6c70e4fa..72b37942 100644 --- a/session/store_test.go +++ b/session/store_test.go @@ -23,7 +23,7 @@ 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. @@ -31,7 +31,7 @@ func TestBasicSessionStore(t *testing.T) { 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 diff --git a/session_rpcserver.go b/session_rpcserver.go index 06832357..14472047 100644 --- a/session_rpcserver.go +++ b/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) }