session: add context to UpdateSessionRemotePubKey

Author: ellemouton

session/interface.go

@@ -294,7 +294,7 @@ type Store interface {
 
 	// UpdateSessionRemotePubKey can be used to add the given remote pub key
 	// to the session with the given local pub key.
-	UpdateSessionRemotePubKey(localPubKey,
+	UpdateSessionRemotePubKey(ctx context.Context, localPubKey,
 		remotePubKey *btcec.PublicKey) error
 
 	// GetSessionByID fetches the session with the given ID.

session/kvdb_store.go

@@ -286,7 +286,7 @@ func (db *BoltStore) NewSession(ctx context.Context, label string, typ Type,
 // to the session with the given local pub key.
 //
 // NOTE: this is part of the Store interface.
-func (db *BoltStore) UpdateSessionRemotePubKey(localPubKey,
+func (db *BoltStore) UpdateSessionRemotePubKey(_ context.Context, localPubKey,
 	remotePubKey *btcec.PublicKey) error {
 
 	key := localPubKey.SerializeCompressed()

session/server.go

@@ -1,13 +1,15 @@
 package session
 
 import (
+	"context"
 	"crypto/tls"
 	"fmt"
 	"sync"
 	"time"
 
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/lightninglabs/lightning-node-connect/mailbox"
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/keychain"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
@@ -21,8 +23,9 @@ type GRPCServerCreator func(opts ...grpc.ServerOption) *grpc.Server
 type mailboxSession struct {
 	server *grpc.Server
 
-	wg   sync.WaitGroup
-	quit chan struct{}
+	cancel fn.Option[context.CancelFunc]
+	wg     sync.WaitGroup
+	quit   chan struct{}
 }
 
 func newMailboxSession() *mailboxSession {
@@ -33,7 +36,8 @@ func newMailboxSession() *mailboxSession {
 
 func (m *mailboxSession) start(session *Session,
 	serverCreator GRPCServerCreator, authData []byte,
-	onUpdate func(local, remote *btcec.PublicKey) error,
+	onUpdate func(ctx context.Context, local,
+		remote *btcec.PublicKey) error,
 	onNewStatus func(s mailbox.ServerStatus)) error {
 
 	tlsConfig := &tls.Config{}
@@ -43,10 +47,13 @@ func (m *mailboxSession) start(session *Session,
 
 	ecdh := &keychain.PrivKeyECDH{PrivKey: session.LocalPrivateKey}
 
+	ctx, cancel := context.WithCancel(context.Background())
+	m.cancel = fn.Some(cancel)
+
 	keys := mailbox.NewConnData(
 		ecdh, session.RemotePublicKey, session.PairingSecret[:],
 		authData, func(key *btcec.PublicKey) error {
-			return onUpdate(session.LocalPublicKey, key)
+			return onUpdate(ctx, session.LocalPublicKey, key)
 		}, nil,
 	)
 
@@ -81,6 +88,7 @@ func (m *mailboxSession) run(mailboxServer *mailbox.Server) {
 }
 
 func (m *mailboxSession) stop() {
+	m.cancel.WhenSome(func(fn context.CancelFunc) { fn() })
 	m.server.Stop()
 	close(m.quit)
 	m.wg.Wait()
@@ -104,7 +112,8 @@ func NewServer(serverCreator GRPCServerCreator) *Server {
 }
 
 func (s *Server) StartSession(session *Session, authData []byte,
-	onUpdate func(local, remote *btcec.PublicKey) error,
+	onUpdate func(ctx context.Context, local,
+		remote *btcec.PublicKey) error,
 	onNewStatus func(s mailbox.ServerStatus)) (chan struct{}, error) {
 
 	s.activeSessionsMtx.Lock()

session/store_test.go

@@ -95,7 +95,9 @@ func TestBasicSessionStore(t *testing.T) {
 	require.NoError(t, err)
 	remotePub := remotePriv.PubKey()
 
-	err = db.UpdateSessionRemotePubKey(session1.LocalPublicKey, remotePub)
+	err = db.UpdateSessionRemotePubKey(
+		ctx, session1.LocalPublicKey, remotePub,
+	)
 	require.NoError(t, err)
 
 	// Assert that the session now does have the remote pub key.

session_rpcserver.go

@@ -1230,7 +1230,9 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
 			"autopilot server: %v", err)
 	}
 
-	err = s.cfg.db.UpdateSessionRemotePubKey(sess.LocalPublicKey, remoteKey)
+	err = s.cfg.db.UpdateSessionRemotePubKey(
+		ctx, sess.LocalPublicKey, remoteKey,
+	)
 	if err != nil {
 		return nil, fmt.Errorf("error setting remote pubkey: %v", err)
 	}