rules: let RealToPseudo take in a PrivacyMapReader

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

For now, only an empty PrivacyMapReader is ever provided to the
RealToPseudo call. This will be changed in the following commit.

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.PrivacyMapReader) (Values,
+	map[string]string, error) {
+
 	return f, nil, nil
 }

rules/channel_restrictions.go

@@ -360,17 +360,22 @@ func (c *ChannelRestrict) PseudoToReal(db firewalldb.PrivacyMapDB) (Values,
 	}, nil
 }
 
-// RealToPseudo converts all the channel IDs into pseudo IDs.
+// RealToPseudo converts all the real channel IDs into pseudo IDs. It returns a
+// map of any new real to pseudo strings that should be persisted that it did
+// not find in the given PrivacyMapReader.
 //
 // NOTE: this is part of the Values interface.
-func (c *ChannelRestrict) RealToPseudo() (Values, map[string]string, error) {
+func (c *ChannelRestrict) RealToPseudo(db firewalldb.PrivacyMapReader) (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 := pseudoFromReal(db, privMapPairs, chanID)
+		if ok {
 			p, err := firewalldb.StrToUint64(pseudo)
 			if err != nil {
 				return nil, nil, err

rules/channel_restrictions_test.go

@@ -161,3 +161,111 @@ 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,
+			},
+		},
+	}
+
+	// Construct the ChannelRestrict deny list. Note that we repeat one of
+	// the entries here in order to ensure that the RealToPseudo method is
+	// forced to look up any real-to-pseudo pairs that it already
+	// generated.
+	cr := &ChannelRestrict{
+		DenyList: []uint64{
+			1,
+			2,
+			3,
+			3,
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			privMapPairDB := firewalldb.NewPrivacyMapPairs(
+				test.dbPreLoad,
+			)
+
+			// Iterate over the preload key value pairs and load
+			// them into the DB.
+			expectedDenyList := make(map[uint64]bool)
+			for _, p := range test.dbPreLoad {
+				// Add the pseudo value to the expected deny
+				// list.
+				pInt, err := firewalldb.StrToUint64(p)
+				require.NoError(t, err)
+
+				expectedDenyList[pInt] = true
+			}
+
+			// Call the RealToPseudo method on the ChannelRestrict
+			// rule. This will return the rule value in its pseudo
+			// form along with any new privacy map pairs that should
+			// be added to the DB.
+			v, newPairs, err := cr.RealToPseudo(privMapPairDB)
+			require.NoError(t, err)
+			require.Len(t, newPairs, len(test.expectNewPairs))
+
+			// We add each new pair to the expected deny list too.
+			for r, p := range newPairs {
+				require.True(t, test.expectNewPairs[r])
+
+				pInt, err := firewalldb.StrToUint64(p)
+				require.NoError(t, err)
+
+				expectedDenyList[pInt] = true
+			}
+
+			denyList, ok := v.(*ChannelRestrict)
+			require.True(t, ok)
+
+			// Assert that the resulting deny list is the same
+			// length as the un-obfuscated one.
+			require.Len(t, denyList.DenyList, len(cr.DenyList))
+
+			// Now iterate over the deny list and assert that each
+			// value appears in our expected deny list.
+			for _, channel := range denyList.DenyList {
+				require.True(t, expectedDenyList[channel])
+			}
+		})
+	}
+}

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.PrivacyMapReader) (Values,
+	map[string]string, error) {
+
 	return h, nil, nil
 }

rules/interfaces.go

@@ -59,9 +59,11 @@ 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 that it did not find
+	// in the given PrivacyMapReader.
+	RealToPseudo(db firewalldb.PrivacyMapReader) (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

@@ -367,17 +367,22 @@ 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 that it did not
+// find in the given PrivacyMapReader.
 //
 // NOTE: this is part of the Values interface.
-func (c *PeerRestrict) RealToPseudo() (Values, map[string]string, error) {
+func (c *PeerRestrict) RealToPseudo(db firewalldb.PrivacyMapReader) (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 := pseudoFromReal(db, privMapPairs, id)
+		if ok {
 			pseudoIDs[i] = pseudo
 			continue
 		}
@@ -393,3 +398,19 @@ 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.PrivacyMapReader,
+	privMapPairs map[string]string, real string) (string, bool) {
+
+	// First check the map.
+	pseudo, ok := privMapPairs[real]
+	if ok {
+		return pseudo, true
+	}
+
+	// Then check the DB reader.
+	return db.GetPseudo(real)
+}

rules/peer_restrictions_test.go

@@ -150,3 +150,100 @@ 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,
+			},
+		},
+	}
+
+	// Construct the PeerRestrict deny list. Note that we repeat one of
+	// the entries here in order to ensure that the RealToPseudo method is
+	// forced to look up any real-to-pseudo pairs that it already
+	// generated.
+	pr := &PeerRestrict{
+		DenyList: []string{
+			"peer 1",
+			"peer 2",
+			"peer 2",
+			"peer 3",
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			privMapPairDB := firewalldb.NewPrivacyMapPairs(
+				test.dbPreLoad,
+			)
+
+			// Add the pseudo values from the preloaded DB to the
+			// expected deny list.
+			expectedDenyList := make(map[string]bool)
+			for _, p := range test.dbPreLoad {
+				expectedDenyList[p] = true
+			}
+
+			// Call the RealToPseudo method on the PeerRestrict
+			// rule. This will return the rule value in its pseudo
+			// form along with any new privacy map pairs that should
+			// be added to the DB.
+			v, newPairs, err := pr.RealToPseudo(privMapPairDB)
+			require.NoError(t, err)
+			require.Len(t, newPairs, len(test.expectNewPairs))
+
+			// We add each new pair to the expected deny list too.
+			for r, p := range newPairs {
+				require.True(t, test.expectNewPairs[r])
+
+				expectedDenyList[p] = true
+			}
+
+			// Assert that the element in the resulting deny list
+			// matches all the elements in our expected deny list.
+			denyList, ok := v.(*PeerRestrict)
+			require.True(t, ok)
+
+			// Assert that the resulting deny list is the same
+			// length as the un-obfuscated one.
+			require.Len(t, denyList.DenyList, len(pr.DenyList))
+
+			// Now iterate over the deny list and assert that each
+			// value appears in our expected deny list.
+			for _, channel := range denyList.DenyList {
+				require.True(t, expectedDenyList[channel])
+			}
+		})
+	}
+}

rules/rate_limit.go

@@ -277,6 +277,8 @@ func (r *RateLimit) PseudoToReal(_ firewalldb.PrivacyMapDB) (Values,
 // 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.PrivacyMapReader) (Values,
+	map[string]string, error) {
+
 	return r, nil, nil
 }

session_rpcserver.go

@@ -838,8 +838,11 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
 		return nil, fmt.Errorf("expiry must be in the future")
 	}
 
-	privacy := !req.NoPrivacyMapper
-	privacyMapPairs := make(map[string]string)
+	var (
+		privacy           = !req.NoPrivacyMapper
+		privacyMapPairs   = make(map[string]string)
+		knownPrivMapPairs = firewalldb.NewPrivacyMapPairs(nil)
+	)
 
 	// First need to fetch all the perms that need to be baked into this
 	// mac based on the features.
@@ -882,7 +885,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(
+						knownPrivMapPairs,
+					)
 					if err != nil {
 						return nil, err
 					}