session: add context to GetSession

This commit is contained in:
Elle Mouton 2025-02-25 13:16:37 +02:00
parent 2dda9308b6
commit 310f11acda
No known key found for this signature in database
GPG key ID: D7D916376026F177
4 changed files with 18 additions and 10 deletions

View file

@ -279,7 +279,7 @@ type Store interface {
error)
// GetSession fetches the session with the given key.
GetSession(key *btcec.PublicKey) (*Session, error)
GetSession(ctx context.Context, key *btcec.PublicKey) (*Session, error)
// ListAllSessions returns all sessions currently known to the store.
ListAllSessions() ([]*Session, error)

View file

@ -319,7 +319,9 @@ func (db *BoltStore) UpdateSessionRemotePubKey(localPubKey,
// GetSession fetches the session with the given key.
//
// NOTE: this is part of the Store interface.
func (db *BoltStore) GetSession(key *btcec.PublicKey) (*Session, error) {
func (db *BoltStore) GetSession(_ context.Context, key *btcec.PublicKey) (
*Session, error) {
var session *Session
err := db.View(func(tx *bbolt.Tx) error {
sessionBucket, err := getBucket(tx, sessionBucketKey)

View file

@ -15,6 +15,9 @@ var testTime = time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC)
// TestBasicSessionStore tests the basic getters and setters of the session
// store.
func TestBasicSessionStore(t *testing.T) {
t.Parallel()
ctx := context.Background()
// Set up a new DB.
clock := clock.NewTestClock(testTime)
db := NewTestDB(t, clock)
@ -73,7 +76,7 @@ func TestBasicSessionStore(t *testing.T) {
// Ensure that we can retrieve each session by both its local pub key
// and by its ID.
for _, s := range []*Session{s1, s2, s3} {
session, err := db.GetSession(s.LocalPublicKey)
session, err := db.GetSession(ctx, s.LocalPublicKey)
require.NoError(t, err)
assertEqualSessions(t, s, session)
@ -83,7 +86,7 @@ func TestBasicSessionStore(t *testing.T) {
}
// Fetch session 1 and assert that it currently has no remote pub key.
session1, err := db.GetSession(s1.LocalPublicKey)
session1, err := db.GetSession(ctx, s1.LocalPublicKey)
require.NoError(t, err)
require.Nil(t, session1.RemotePublicKey)
@ -96,7 +99,7 @@ func TestBasicSessionStore(t *testing.T) {
require.NoError(t, err)
// Assert that the session now does have the remote pub key.
session1, err = db.GetSession(s1.LocalPublicKey)
session1, err = db.GetSession(ctx, s1.LocalPublicKey)
require.NoError(t, err)
require.True(t, remotePub.IsEqual(session1.RemotePublicKey))
@ -105,7 +108,7 @@ func TestBasicSessionStore(t *testing.T) {
// Now revoke the session and assert that the state is revoked.
require.NoError(t, db.ShiftState(s1.ID, StateRevoked))
s1, err = db.GetSession(s1.LocalPublicKey)
s1, err = db.GetSession(ctx, s1.LocalPublicKey)
require.NoError(t, err)
require.Equal(t, s1.State, StateRevoked)
@ -285,6 +288,9 @@ func TestLinkedSessions(t *testing.T) {
// TestStateShift tests that the ShiftState method works as expected.
func TestStateShift(t *testing.T) {
t.Parallel()
ctx := context.Background()
// Set up a new DB.
clock := clock.NewTestClock(testTime)
db := NewTestDB(t, clock)
@ -294,7 +300,7 @@ func TestStateShift(t *testing.T) {
// Check that the session is in the StateCreated state. Also check that
// the "RevokedAt" time has not yet been set.
s1, err := db.GetSession(s1.LocalPublicKey)
s1, err := db.GetSession(ctx, s1.LocalPublicKey)
require.NoError(t, err)
require.Equal(t, StateCreated, s1.State)
require.Equal(t, time.Time{}, s1.RevokedAt)
@ -305,7 +311,7 @@ func TestStateShift(t *testing.T) {
// This should have worked. Since it is now in a terminal state, the
// "RevokedAt" time should be set.
s1, err = db.GetSession(s1.LocalPublicKey)
s1, err = db.GetSession(ctx, s1.LocalPublicKey)
require.NoError(t, err)
require.Equal(t, StateRevoked, s1.State)
require.True(t, clock.Now().Equal(s1.RevokedAt))

View file

@ -563,7 +563,7 @@ func (s *sessionRpcServer) RevokeSession(ctx context.Context,
return nil, fmt.Errorf("error parsing public key: %v", err)
}
sess, err := s.cfg.db.GetSession(pubKey)
sess, err := s.cfg.db.GetSession(ctx, pubKey)
if err != nil {
return nil, fmt.Errorf("error fetching session: %v", err)
}
@ -1303,7 +1303,7 @@ func (s *sessionRpcServer) RevokeAutopilotSession(ctx context.Context,
return nil, fmt.Errorf("error parsing public key: %v", err)
}
sess, err := s.cfg.db.GetSession(pubKey)
sess, err := s.cfg.db.GetSession(ctx, pubKey)
if err != nil {
return nil, err
}