session: add context to GetSession

Author: ellemouton

session/interface.go

@@ -198,7 +198,7 @@ type Store interface {
 		flags PrivacyFlags) (*Session, 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)

session/kvdb_store.go

@@ -320,7 +320,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)

session/store_test.go

@@ -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, err := NewDB(t.TempDir(), "test.db", 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)
 
@@ -293,6 +296,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, err := NewDB(t.TempDir(), "test.db", clock)
@@ -306,7 +312,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)
@@ -317,7 +323,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))

session_rpcserver.go

@@ -556,7 +556,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)
 	}
@@ -1276,7 +1276,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
 	}