Skip to content

Commit 8b52899

Browse files
ellemoutonellemouton
committed
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.
1 parent d0bc3c3 commit 8b52899

File tree

9 files changed

+259
-15
lines changed

9 files changed

+259
-15
lines changed

‎rules/chan_policy_bounds.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -329,6 +329,8 @@ func (f *ChanPolicyBounds) PseudoToReal(_ firewalldb.PrivacyMapDB) (Values,
329329
// that should be persisted. This is a no-op for the ChanPolicyBounds rule.
330330
//
331331
// NOTE: this is part of the Values interface.
332-
func (f *ChanPolicyBounds) RealToPseudo() (Values, map[string]string, error) {
332+
func (f *ChanPolicyBounds) RealToPseudo(_ firewalldb.PrivacyMapReader) (Values,
333+
map[string]string, error) {
334+
333335
return f, nil, nil
334336
}

‎rules/channel_restrictions.go‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -360,17 +360,22 @@ func (c *ChannelRestrict) PseudoToReal(db firewalldb.PrivacyMapDB) (Values,
360360
}, nil
361361
}
362362

363-
// RealToPseudo converts all the channel IDs into pseudo IDs.
363+
// RealToPseudo converts all the real channel IDs into pseudo IDs. It returns a
364+
// map of any new real to pseudo strings that should be persisted that it did
365+
// not find in the given PrivacyMapReader.
364366
//
365367
// NOTE: this is part of the Values interface.
366-
func (c *ChannelRestrict) RealToPseudo() (Values, map[string]string, error) {
368+
func (c *ChannelRestrict) RealToPseudo(db firewalldb.PrivacyMapReader) (Values,
369+
map[string]string, error) {
370+
367371
pseudoIDs := make([]uint64, len(c.DenyList))
368372
privMapPairs := make(map[string]string)
369373
for i, c := range c.DenyList {
370374
// TODO(elle): check that this channel actually exists
371375

372376
chanID := firewalldb.Uint64ToStr(c)
373-
if pseudo, ok := privMapPairs[chanID]; ok {
377+
pseudo, ok := pseudoFromReal(db, privMapPairs, chanID)
378+
if ok {
374379
p, err := firewalldb.StrToUint64(pseudo)
375380
if err != nil {
376381
return nil, nil, err

‎rules/channel_restrictions_test.go‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -161,3 +161,111 @@ func (m *mockLndClient) ListChannels(_ context.Context, _, _ bool) (
161161

162162
return m.channels, nil
163163
}
164+
165+
// TestChannelRestrictRealToPseudo tests that the ChannelRestrict's RealToPseudo
166+
// method correctly determines which real strings to generate pseudo pairs for
167+
// based on the privacy map db passed to it.
168+
func TestChannelRestrictRealToPseudo(t *testing.T) {
169+
chanID1 := firewalldb.Uint64ToStr(1)
170+
chanID2 := firewalldb.Uint64ToStr(2)
171+
chanID3 := firewalldb.Uint64ToStr(3)
172+
chanID2Obfuscated := firewalldb.Uint64ToStr(200)
173+
174+
tests := []struct {
175+
name string
176+
dbPreLoad map[string]string
177+
expectNewPairs map[string]bool
178+
}{
179+
{
180+
// If there is no preloaded DB, then we expect all the
181+
// values in the deny list to be returned from the
182+
// RealToPseudo method.
183+
name: "no pre loaded db",
184+
expectNewPairs: map[string]bool{
185+
chanID1: true,
186+
chanID2: true,
187+
chanID3: true,
188+
},
189+
},
190+
{
191+
// If the DB is preloaded with an entry for "channel 2"
192+
// then we don't expect that entry to be returned in the
193+
// set of new pairs.
194+
name: "partially pre-loaded DB",
195+
dbPreLoad: map[string]string{
196+
chanID2: chanID2Obfuscated,
197+
},
198+
expectNewPairs: map[string]bool{
199+
chanID1: true,
200+
chanID3: true,
201+
},
202+
},
203+
}
204+
205+
// Construct the ChannelRestrict deny list. Note that we repeat one of
206+
// the entries here in order to ensure that the RealToPseudo method is
207+
// forced to look up any real-to-pseudo pairs that it already
208+
// generated.
209+
cr := &ChannelRestrict{
210+
DenyList: []uint64{
211+
1,
212+
2,
213+
3,
214+
3,
215+
},
216+
}
217+
218+
for _, test := range tests {
219+
test := test
220+
t.Run(test.name, func(t *testing.T) {
221+
t.Parallel()
222+
223+
privMapPairDB := firewalldb.NewPrivacyMapPairs(
224+
test.dbPreLoad,
225+
)
226+
227+
// Iterate over the preload key value pairs and load
228+
// them into the DB.
229+
expectedDenyList := make(map[uint64]bool)
230+
for _, p := range test.dbPreLoad {
231+
// Add the pseudo value to the expected deny
232+
// list.
233+
pInt, err := firewalldb.StrToUint64(p)
234+
require.NoError(t, err)
235+
236+
expectedDenyList[pInt] = true
237+
}
238+
239+
// Call the RealToPseudo method on the ChannelRestrict
240+
// rule. This will return the rule value in its pseudo
241+
// form along with any new privacy map pairs that should
242+
// be added to the DB.
243+
v, newPairs, err := cr.RealToPseudo(privMapPairDB)
244+
require.NoError(t, err)
245+
require.Len(t, newPairs, len(test.expectNewPairs))
246+
247+
// We add each new pair to the expected deny list too.
248+
for r, p := range newPairs {
249+
require.True(t, test.expectNewPairs[r])
250+
251+
pInt, err := firewalldb.StrToUint64(p)
252+
require.NoError(t, err)
253+
254+
expectedDenyList[pInt] = true
255+
}
256+
257+
denyList, ok := v.(*ChannelRestrict)
258+
require.True(t, ok)
259+
260+
// Assert that the resulting deny list is the same
261+
// length as the un-obfuscated one.
262+
require.Len(t, denyList.DenyList, len(cr.DenyList))
263+
264+
// Now iterate over the deny list and assert that each
265+
// value appears in our expected deny list.
266+
for _, channel := range denyList.DenyList {
267+
require.True(t, expectedDenyList[channel])
268+
}
269+
})
270+
}
271+
}

‎rules/history_limit.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -266,6 +266,8 @@ func (h *HistoryLimit) PseudoToReal(_ firewalldb.PrivacyMapDB) (Values,
266266
// that should be persisted. This is a no-op for the HistoryLimit rule.
267267
//
268268
// NOTE: this is part of the Values interface.
269-
func (h *HistoryLimit) RealToPseudo() (Values, map[string]string, error) {
269+
func (h *HistoryLimit) RealToPseudo(_ firewalldb.PrivacyMapReader) (Values,
270+
map[string]string, error) {
271+
270272
return h, nil, nil
271273
}

‎rules/interfaces.go‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -59,9 +59,11 @@ type Values interface {
5959
ToProto() *litrpc.RuleValue
6060

6161
// RealToPseudo converts the rule Values to a new one that uses pseudo
62-
// keys, channel IDs, channel points etc. It returns a map of real to
63-
// pseudo strings that should be persisted.
64-
RealToPseudo() (Values, map[string]string, error)
62+
// keys, channel IDs, channel points etc. It returns a map of any new
63+
// real to pseudo strings that should be persisted that it did not find
64+
// in the given PrivacyMapReader.
65+
RealToPseudo(db firewalldb.PrivacyMapReader) (Values, map[string]string,
66+
error)
6567

6668
// PseudoToReal attempts to convert any appropriate pseudo fields in
6769
// the rule Values to their corresponding real values. It uses the

‎rules/peer_restrictions.go‎

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -367,17 +367,22 @@ func (c *PeerRestrict) PseudoToReal(db firewalldb.PrivacyMapDB) (Values,
367367
}, nil
368368
}
369369

370-
// RealToPseudo converts all the real peer IDs into pseudo IDs.
370+
// RealToPseudo converts all the real peer IDs into pseudo IDs. It returns a map
371+
// of any new real to pseudo strings that should be persisted that it did not
372+
// find in the given PrivacyMapReader.
371373
//
372374
// NOTE: this is part of the Values interface.
373-
func (c *PeerRestrict) RealToPseudo() (Values, map[string]string, error) {
375+
func (c *PeerRestrict) RealToPseudo(db firewalldb.PrivacyMapReader) (Values,
376+
map[string]string, error) {
377+
374378
pseudoIDs := make([]string, len(c.DenyList))
375379
privMapPairs := make(map[string]string)
376380
for i, id := range c.DenyList {
377381
// TODO(elle): check that this peer is actually one of our
378382
// channel peers.
379383

380-
if pseudo, ok := privMapPairs[id]; ok {
384+
pseudo, ok := pseudoFromReal(db, privMapPairs, id)
385+
if ok {
381386
pseudoIDs[i] = pseudo
382387
continue
383388
}
@@ -393,3 +398,19 @@ func (c *PeerRestrict) RealToPseudo() (Values, map[string]string, error) {
393398

394399
return &PeerRestrict{DenyList: pseudoIDs}, privMapPairs, nil
395400
}
401+
402+
// pseudoFromReal is a helper that can be used to get the associated pseudo
403+
// value for a given real value from either the privacy map db if it is defined
404+
// or from a set of real-to-pseudo pairs.
405+
func pseudoFromReal(db firewalldb.PrivacyMapReader,
406+
privMapPairs map[string]string, real string) (string, bool) {
407+
408+
// First check the map.
409+
pseudo, ok := privMapPairs[real]
410+
if ok {
411+
return pseudo, true
412+
}
413+
414+
// Then check the DB reader.
415+
return db.GetPseudo(real)
416+
}

‎rules/peer_restrictions_test.go‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,3 +150,100 @@ func TestPeerRestrictCheckRequest(t *testing.T) {
150150
)
151151
require.NoError(t, err)
152152
}
153+
154+
// TestPeerRestrictionRealToPseudo tests that the PeerRestriction's RealToPseudo
155+
// method correctly determines which real strings to generate pseudo pairs for
156+
// based on the privacy map db passed to it.
157+
func TestPeerRestrictRealToPseudo(t *testing.T) {
158+
tests := []struct {
159+
name string
160+
dbPreLoad map[string]string
161+
expectNewPairs map[string]bool
162+
}{
163+
{
164+
// If there is no preloaded DB, then we expect all the
165+
// values in the deny list to be returned from the
166+
// RealToPseudo method.
167+
name: "no pre loaded db",
168+
expectNewPairs: map[string]bool{
169+
"peer 1": true,
170+
"peer 2": true,
171+
"peer 3": true,
172+
},
173+
},
174+
{
175+
// If the DB is preloaded with an entry for "peer 2"
176+
// then we don't expect that entry to be returned in the
177+
// set of new pairs.
178+
name: "partially pre-loaded DB",
179+
dbPreLoad: map[string]string{
180+
"peer 2": "obfuscated peer 2",
181+
},
182+
expectNewPairs: map[string]bool{
183+
"peer 1": true,
184+
"peer 3": true,
185+
},
186+
},
187+
}
188+
189+
// Construct the PeerRestrict deny list. Note that we repeat one of
190+
// the entries here in order to ensure that the RealToPseudo method is
191+
// forced to look up any real-to-pseudo pairs that it already
192+
// generated.
193+
pr := &PeerRestrict{
194+
DenyList: []string{
195+
"peer 1",
196+
"peer 2",
197+
"peer 2",
198+
"peer 3",
199+
},
200+
}
201+
202+
for _, test := range tests {
203+
test := test
204+
t.Run(test.name, func(t *testing.T) {
205+
t.Parallel()
206+
207+
privMapPairDB := firewalldb.NewPrivacyMapPairs(
208+
test.dbPreLoad,
209+
)
210+
211+
// Add the pseudo values from the preloaded DB to the
212+
// expected deny list.
213+
expectedDenyList := make(map[string]bool)
214+
for _, p := range test.dbPreLoad {
215+
expectedDenyList[p] = true
216+
}
217+
218+
// Call the RealToPseudo method on the PeerRestrict
219+
// rule. This will return the rule value in its pseudo
220+
// form along with any new privacy map pairs that should
221+
// be added to the DB.
222+
v, newPairs, err := pr.RealToPseudo(privMapPairDB)
223+
require.NoError(t, err)
224+
require.Len(t, newPairs, len(test.expectNewPairs))
225+
226+
// We add each new pair to the expected deny list too.
227+
for r, p := range newPairs {
228+
require.True(t, test.expectNewPairs[r])
229+
230+
expectedDenyList[p] = true
231+
}
232+
233+
// Assert that the element in the resulting deny list
234+
// matches all the elements in our expected deny list.
235+
denyList, ok := v.(*PeerRestrict)
236+
require.True(t, ok)
237+
238+
// Assert that the resulting deny list is the same
239+
// length as the un-obfuscated one.
240+
require.Len(t, denyList.DenyList, len(pr.DenyList))
241+
242+
// Now iterate over the deny list and assert that each
243+
// value appears in our expected deny list.
244+
for _, channel := range denyList.DenyList {
245+
require.True(t, expectedDenyList[channel])
246+
}
247+
})
248+
}
249+
}

‎rules/rate_limit.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -277,6 +277,8 @@ func (r *RateLimit) PseudoToReal(_ firewalldb.PrivacyMapDB) (Values,
277277
// that should be persisted. This is a no-op for the RateLimit rule.
278278
//
279279
// NOTE: this is part of the Values interface.
280-
func (r *RateLimit) RealToPseudo() (Values, map[string]string, error) {
280+
func (r *RateLimit) RealToPseudo(_ firewalldb.PrivacyMapReader) (Values,
281+
map[string]string, error) {
282+
281283
return r, nil, nil
282284
}

‎session_rpcserver.go‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -838,8 +838,11 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
838838
return nil, fmt.Errorf("expiry must be in the future")
839839
}
840840

841-
privacy := !req.NoPrivacyMapper
842-
privacyMapPairs := make(map[string]string)
841+
var (
842+
privacy = !req.NoPrivacyMapper
843+
privacyMapPairs = make(map[string]string)
844+
knownPrivMapPairs = firewalldb.NewPrivacyMapPairs(nil)
845+
)
843846

844847
// First need to fetch all the perms that need to be baked into this
845848
// mac based on the features.
@@ -882,7 +885,9 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
882885

883886
if privacy {
884887
var privMapPairs map[string]string
885-
v, privMapPairs, err = v.RealToPseudo()
888+
v, privMapPairs, err = v.RealToPseudo(
889+
knownPrivMapPairs,
890+
)
886891
if err != nil {
887892
return nil, err
888893
}

0 commit comments

Comments
 (0)