rules: let RealToPseudo take in a priv map DB

This commit expands the RealToPseudo methods to take in a privacy map
db. This allows the methods to check if the privacy map db already
contains an entry for a "real" string before generating a new one.

Author: ellemouton

rules/chan_policy_bounds.go

@@ -329,6 +329,8 @@ func (f *ChanPolicyBounds) PseudoToReal(_ firewalldb.PrivacyMapDB) (Values,
 // that should be persisted. This is a no-op for the ChanPolicyBounds rule.
 //
 // NOTE: this is part of the Values interface.
-func (f *ChanPolicyBounds) RealToPseudo() (Values, map[string]string, error) {
+func (f *ChanPolicyBounds) RealToPseudo(_ firewalldb.PrivacyMapDB) (Values,
+	map[string]string, error) {
+
 	return f, nil, nil
 }

rules/channel_restrictions.go

@@ -360,17 +360,24 @@ func (c *ChannelRestrict) PseudoToReal(db firewalldb.PrivacyMapDB) (Values,
 	}, nil
 }
 
-// RealToPseudo converts all the channel IDs into pseudo IDs.
+// RealToPseudo converts all the channel IDs into pseudo IDs. It returns a map
+// of any _new_ real to pseudo strings that should be persisted.
 //
 // NOTE: this is part of the Values interface.
-func (c *ChannelRestrict) RealToPseudo() (Values, map[string]string, error) {
+func (c *ChannelRestrict) RealToPseudo(db firewalldb.PrivacyMapDB) (Values,
+	map[string]string, error) {
+
 	pseudoIDs := make([]uint64, len(c.DenyList))
 	privMapPairs := make(map[string]string)
 	for i, c := range c.DenyList {
 		// TODO(elle): check that this channel actually exists
 
 		chanID := firewalldb.Uint64ToStr(c)
-		if pseudo, ok := privMapPairs[chanID]; ok {
+
+		pseudo, ok, err := pseudoFromReal(db, privMapPairs, chanID)
+		if err != nil {
+			return nil, nil, err
+		} else if ok {
 			p, err := firewalldb.StrToUint64(pseudo)
 			if err != nil {
 				return nil, nil, err

rules/channel_restrictions_test.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
+	"github.com/lightninglabs/lightning-terminal/firewall/mock"
 	"github.com/lightninglabs/lightning-terminal/firewalldb"
 	"github.com/lightninglabs/lndclient"
 	"github.com/lightningnetwork/lnd/lnrpc"
@@ -161,3 +162,91 @@ func (m *mockLndClient) ListChannels(_ context.Context, _, _ bool) (
 
 	return m.channels, nil
 }
+
+// TestChannelRestrictRealToPseudo tests that the ChannelRestrict's RealToPseudo
+// method correctly determines which real strings to generate pseudo pairs for
+// based on the privacy map db passed to it.
+func TestChannelRestrictRealToPseudo(t *testing.T) {
+	chanID1 := firewalldb.Uint64ToStr(1)
+	chanID2 := firewalldb.Uint64ToStr(2)
+	chanID3 := firewalldb.Uint64ToStr(3)
+	chanID2Obfuscated := firewalldb.Uint64ToStr(200)
+
+	tests := []struct {
+		name           string
+		dbPreLoad      map[string]string
+		expectNewPairs map[string]bool
+	}{
+		{
+			// If there is no preloaded DB, then we expect all the
+			// values in the deny list to be returned from the
+			// RealToPseudo method.
+			name: "no pre loaded db",
+			expectNewPairs: map[string]bool{
+				chanID1: true,
+				chanID2: true,
+				chanID3: true,
+			},
+		},
+		{
+			// If the DB is preloaded with an entry for "channel 2"
+			// then we don't expect that entry to be returned in the
+			// set of new pairs.
+			name: "partially pre-loaded DB",
+			dbPreLoad: map[string]string{
+				chanID2: chanID2Obfuscated,
+			},
+			expectNewPairs: map[string]bool{
+				chanID1: true,
+				chanID3: true,
+			},
+		},
+	}
+
+	cr := &ChannelRestrict{
+		DenyList: []uint64{
+			1,
+			2,
+			3,
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			var privDB firewalldb.PrivacyMapDB
+			if len(test.dbPreLoad) != 0 {
+				privDB = mock.NewPrivacyMapDB()
+			}
+
+			var expectedDenyList []string
+			for r, p := range test.dbPreLoad {
+				err := privDB.View(
+					func(tx firewalldb.PrivacyMapTx) error {
+						return tx.NewPair(r, p)
+					},
+				)
+				require.NoError(t, err)
+
+				expectedDenyList = append(expectedDenyList, p)
+			}
+
+			v, newPairs, err := cr.RealToPseudo(privDB)
+			require.NoError(t, err)
+			require.Len(t, newPairs, len(test.expectNewPairs))
+
+			for r, p := range newPairs {
+				require.True(t, test.expectNewPairs[r])
+
+				expectedDenyList = append(expectedDenyList, p)
+			}
+
+			denyList, ok := v.(*ChannelRestrict)
+			require.True(t, ok)
+
+			require.EqualValues(t, v, denyList)
+		})
+	}
+}

rules/history_limit.go

@@ -266,6 +266,8 @@ func (h *HistoryLimit) PseudoToReal(_ firewalldb.PrivacyMapDB) (Values,
 // that should be persisted. This is a no-op for the HistoryLimit rule.
 //
 // NOTE: this is part of the Values interface.
-func (h *HistoryLimit) RealToPseudo() (Values, map[string]string, error) {
+func (h *HistoryLimit) RealToPseudo(_ firewalldb.PrivacyMapDB) (Values,
+	map[string]string, error) {
+
 	return h, nil, nil
 }

rules/interfaces.go

@@ -59,9 +59,10 @@ type Values interface {
 	ToProto() *litrpc.RuleValue
 
 	// RealToPseudo converts the rule Values to a new one that uses pseudo
-	// keys, channel IDs, channel points etc. It returns a map of real to
-	// pseudo strings that should be persisted.
-	RealToPseudo() (Values, map[string]string, error)
+	// keys, channel IDs, channel points etc. It returns a map of any _new_
+	// real to pseudo strings that should be persisted.
+	RealToPseudo(db firewalldb.PrivacyMapDB) (Values, map[string]string,
+		error)
 
 	// PseudoToReal attempts to convert any appropriate pseudo fields in
 	// the rule Values to their corresponding real values. It uses the

rules/peer_restrictions.go

@@ -2,6 +2,7 @@ package rules
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"sync"
 
@@ -367,22 +368,28 @@ func (c *PeerRestrict) PseudoToReal(db firewalldb.PrivacyMapDB) (Values,
 	}, nil
 }
 
-// RealToPseudo converts all the real peer IDs into pseudo IDs.
+// RealToPseudo converts all the real peer IDs into pseudo IDs. It returns a map
+// of any _new_ real to pseudo strings that should be persisted.
 //
 // NOTE: this is part of the Values interface.
-func (c *PeerRestrict) RealToPseudo() (Values, map[string]string, error) {
+func (c *PeerRestrict) RealToPseudo(db firewalldb.PrivacyMapDB) (Values,
+	map[string]string, error) {
+
 	pseudoIDs := make([]string, len(c.DenyList))
 	privMapPairs := make(map[string]string)
 	for i, id := range c.DenyList {
 		// TODO(elle): check that this peer is actually one of our
 		//  channel peers.
 
-		if pseudo, ok := privMapPairs[id]; ok {
+		pseudo, ok, err := pseudoFromReal(db, privMapPairs, id)
+		if err != nil {
+			return nil, nil, err
+		} else if ok {
 			pseudoIDs[i] = pseudo
 			continue
 		}
 
-		pseudo, err := firewalldb.NewPseudoStr(len(id))
+		pseudo, err = firewalldb.NewPseudoStr(len(id))
 		if err != nil {
 			return nil, nil, err
 		}
@@ -393,3 +400,52 @@ func (c *PeerRestrict) RealToPseudo() (Values, map[string]string, error) {
 
 	return &PeerRestrict{DenyList: pseudoIDs}, privMapPairs, nil
 }
+
+// pseudoFromReal is a helper that can be used to get the associated pseudo
+// value for a given real value from either the privacy map db if it is defined
+// or from a set of real-to-pseudo pairs.
+func pseudoFromReal(db firewalldb.PrivacyMapDB,
+	privMapPairs map[string]string, real string) (string, bool, error) {
+
+	pseudo, ok := privMapPairs[real]
+
+	switch {
+	// If the db is nil, then only the map needs to be checked.
+	case db == nil:
+		return pseudo, ok, nil
+
+	// If the value was found in the map, then return that.
+	case ok:
+		return pseudo, true, nil
+
+	// Otherwise, the DB is checked.
+	default:
+	}
+
+	err := db.View(func(tx firewalldb.PrivacyMapTx) error {
+		var err error
+		pseudo, err = tx.RealToPseudo(real)
+		switch {
+		// If the error is nil, an entry exists.
+		case err == nil:
+			ok = true
+
+			return nil
+
+		// Return any unexpected error.
+		case !errors.Is(err, firewalldb.ErrNoSuchKeyFound):
+			return err
+
+		// No key found.
+		default:
+			ok = false
+		}
+
+		return nil
+	})
+	if err != nil {
+		return "", false, err
+	}
+
+	return pseudo, ok, nil
+}

rules/peer_restrictions_test.go

@@ -7,6 +7,7 @@ import (
 	"testing"
 
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
+	"github.com/lightninglabs/lightning-terminal/firewall/mock"
 	"github.com/lightninglabs/lightning-terminal/firewalldb"
 	"github.com/lightninglabs/lndclient"
 	"github.com/lightningnetwork/lnd/lnrpc"
@@ -150,3 +151,86 @@ func TestPeerRestrictCheckRequest(t *testing.T) {
 	)
 	require.NoError(t, err)
 }
+
+// TestPeerRestrictionRealToPseudo tests that the PeerRestriction's RealToPseudo
+// method correctly determines which real strings to generate pseudo pairs for
+// based on the privacy map db passed to it.
+func TestPeerRestrictRealToPseudo(t *testing.T) {
+	tests := []struct {
+		name           string
+		dbPreLoad      map[string]string
+		expectNewPairs map[string]bool
+	}{
+		{
+			// If there is no preloaded DB, then we expect all the
+			// values in the deny list to be returned from the
+			// RealToPseudo method.
+			name: "no pre loaded db",
+			expectNewPairs: map[string]bool{
+				"peer 1": true,
+				"peer 2": true,
+				"peer 3": true,
+			},
+		},
+		{
+			// If the DB is preloaded with an entry for "peer 2"
+			// then we don't expect that entry to be returned in the
+			// set of new pairs.
+			name: "partially pre-loaded DB",
+			dbPreLoad: map[string]string{
+				"peer 2": "obfuscated peer 2",
+			},
+			expectNewPairs: map[string]bool{
+				"peer 1": true,
+				"peer 3": true,
+			},
+		},
+	}
+
+	pr := &PeerRestrict{
+		DenyList: []string{
+			"peer 1",
+			"peer 2",
+			"peer 3",
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			var privDB firewalldb.PrivacyMapDB
+			if len(test.dbPreLoad) != 0 {
+				privDB = mock.NewPrivacyMapDB()
+			}
+
+			var expectedDenyList []string
+			for r, p := range test.dbPreLoad {
+				err := privDB.View(
+					func(tx firewalldb.PrivacyMapTx) error {
+						return tx.NewPair(r, p)
+					},
+				)
+				require.NoError(t, err)
+
+				expectedDenyList = append(expectedDenyList, p)
+			}
+
+			v, newPairs, err := pr.RealToPseudo(privDB)
+			require.NoError(t, err)
+			require.Len(t, newPairs, len(test.expectNewPairs))
+
+			for r, p := range newPairs {
+				require.True(t, test.expectNewPairs[r])
+
+				expectedDenyList = append(expectedDenyList, p)
+			}
+
+			denyList, ok := v.(*PeerRestrict)
+			require.True(t, ok)
+
+			require.EqualValues(t, v, denyList)
+		})
+	}
+}

rules/rate_limit.go

@@ -273,10 +273,12 @@ func (r *RateLimit) PseudoToReal(_ firewalldb.PrivacyMapDB) (Values,
 }
 
 // RealToPseudo converts the rule Values to a new one that uses pseudo keys,
-// channel IDs, channel points etc. It returns a map of real to pseudo strings
-// that should be persisted. This is a no-op for the RateLimit rule.
+// channel IDs, channel points etc. It returns a map of any new real to pseudo
+// strings that should be persisted. This is a no-op for the RateLimit rule.
 //
 // NOTE: this is part of the Values interface.
-func (r *RateLimit) RealToPseudo() (Values, map[string]string, error) {
+func (r *RateLimit) RealToPseudo(_ firewalldb.PrivacyMapDB) (Values,
+	map[string]string, error) {
+
 	return r, nil, nil
 }

session_rpcserver.go

@@ -882,7 +882,9 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
 
 				if privacy {
 					var privMapPairs map[string]string
-					v, privMapPairs, err = v.RealToPseudo()
+					v, privMapPairs, err = v.RealToPseudo(
+						nil,
+					)
 					if err != nil {
 						return nil, err
 					}