From ef936114892f2cd9daeeedf03af070cb757a1c32 Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Thu, 13 Mar 2025 16:32:23 -0500
Subject: [PATCH] firewalldb: thread contexts to FetchAllPairs

Update the FetchAllPairs method of the PrivacyMapTx interface to take a
context.
---
 firewall/privacy_mapper_test.go   | 4 ++--
 firewalldb/privacy_mapper.go      | 6 ++++--
 firewalldb/privacy_mapper_test.go | 8 ++++----
 session_rpcserver.go              | 4 ++--
 4 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/firewall/privacy_mapper_test.go b/firewall/privacy_mapper_test.go
index 685c7cb3..1998d128 100644
--- a/firewall/privacy_mapper_test.go
+++ b/firewall/privacy_mapper_test.go
@@ -1151,8 +1151,8 @@ func (m *mockPrivacyMapDB) RealToPseudo(_ context.Context, real string) (string,
 	return p, nil
 }
 
-func (m *mockPrivacyMapDB) FetchAllPairs() (*firewalldb.PrivacyMapPairs,
-	error) {
+func (m *mockPrivacyMapDB) FetchAllPairs(_ context.Context) (
+	*firewalldb.PrivacyMapPairs, error) {
 
 	return firewalldb.NewPrivacyMapPairs(m.r2p), nil
 }
diff --git a/firewalldb/privacy_mapper.go b/firewalldb/privacy_mapper.go
index 6ae8b9fa..eadb4933 100644
--- a/firewalldb/privacy_mapper.go
+++ b/firewalldb/privacy_mapper.go
@@ -83,7 +83,7 @@ type PrivacyMapTx interface {
 
 	// FetchAllPairs loads and returns the real-to-pseudo pairs in the form
 	// of a PrivacyMapPairs struct.
-	FetchAllPairs() (*PrivacyMapPairs, error)
+	FetchAllPairs(ctx context.Context) (*PrivacyMapPairs, error)
 }
 
 // privacyMapDB is an implementation of PrivacyMapDB.
@@ -287,7 +287,9 @@ func (p *privacyMapTx) RealToPseudo(_ context.Context, real string) (string,
 // FetchAllPairs loads and returns the real-to-pseudo pairs.
 //
 // NOTE: this is part of the PrivacyMapTx interface.
-func (p *privacyMapTx) FetchAllPairs() (*PrivacyMapPairs, error) {
+func (p *privacyMapTx) FetchAllPairs(_ context.Context) (*PrivacyMapPairs,
+	error) {
+
 	privacyBucket, err := getBucket(p.boltTx, privacyBucketKey)
 	if err != nil {
 		return nil, err
diff --git a/firewalldb/privacy_mapper_test.go b/firewalldb/privacy_mapper_test.go
index 95c08aeb..7be4d3b6 100644
--- a/firewalldb/privacy_mapper_test.go
+++ b/firewalldb/privacy_mapper_test.go
@@ -40,7 +40,7 @@ func TestPrivacyMapStorage(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, "real", real)
 
-		pairs, err := tx.FetchAllPairs()
+		pairs, err := tx.FetchAllPairs(ctx)
 		require.NoError(t, err)
 
 		require.EqualValues(t, pairs.pairs, map[string]string{
@@ -70,7 +70,7 @@ func TestPrivacyMapStorage(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, "real 2", real)
 
-		pairs, err := tx.FetchAllPairs()
+		pairs, err := tx.FetchAllPairs(ctx)
 		require.NoError(t, err)
 
 		require.EqualValues(t, pairs.pairs, map[string]string{
@@ -85,7 +85,7 @@ func TestPrivacyMapStorage(t *testing.T) {
 	_ = pdb3.Update(ctx, func(ctx context.Context, tx PrivacyMapTx) error {
 		// Check that calling FetchAllPairs returns an empty map if
 		// nothing exists in the DB yet.
-		m, err := tx.FetchAllPairs()
+		m, err := tx.FetchAllPairs(ctx)
 		require.NoError(t, err)
 		require.Empty(t, m.pairs)
 
@@ -116,7 +116,7 @@ func TestPrivacyMapStorage(t *testing.T) {
 		require.NoError(t, err)
 
 		// Check that FetchAllPairs correctly returns all the pairs.
-		pairs, err := tx.FetchAllPairs()
+		pairs, err := tx.FetchAllPairs(ctx)
 		require.NoError(t, err)
 
 		require.EqualValues(t, pairs.pairs, map[string]string{
diff --git a/session_rpcserver.go b/session_rpcserver.go
index a3a3b339..652196f5 100644
--- a/session_rpcserver.go
+++ b/session_rpcserver.go
@@ -901,10 +901,10 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
 		linkedGroupSession = groupSess
 
 		privDB := s.cfg.privMap(groupID)
-		err = privDB.View(ctx, func(_ context.Context,
+		err = privDB.View(ctx, func(ctx context.Context,
 			tx firewalldb.PrivacyMapTx) error {
 
-			knownPrivMapPairs, err = tx.FetchAllPairs()
+			knownPrivMapPairs, err = tx.FetchAllPairs(ctx)
 
 			return err
 		})