multi: thread contexts through privacy map interfaces

Update the PrivacyMapDB interface methods to take contexts (both the
methods themselves and the call-back params) and then ensure all
implementations are updated and all call-sites pass contexts through
correctly.
This commit is contained in:
Elle Mouton 2025-03-11 14:57:52 -05:00
parent e49a1c373f
commit 7ce36d7e7d
No known key found for this signature in database
GPG key ID: D7D916376026F177
16 changed files with 160 additions and 90 deletions

View file

@ -325,14 +325,16 @@ func handleGetInfoResponse(db firewalldb.PrivacyMapDB,
flags session.PrivacyFlags) func(ctx context.Context,
r *lnrpc.GetInfoResponse) (proto.Message, error) {
return func(_ context.Context, r *lnrpc.GetInfoResponse) (
return func(ctx context.Context, r *lnrpc.GetInfoResponse) (
proto.Message, error) {
// We hide the pubkey unless it is disabled.
pseudoPubKey := r.IdentityPubkey
if !flags.Contains(session.ClearPubkeys) {
err := db.Update(
func(tx firewalldb.PrivacyMapTx) error {
err := db.Update(ctx,
func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
var err error
pseudoPubKey, err = firewalldb.HideString(
tx, r.IdentityPubkey,
@ -377,14 +379,16 @@ func handleFwdHistoryResponse(db firewalldb.PrivacyMapDB,
randIntn func(int) (int, error)) func(ctx context.Context,
r *lnrpc.ForwardingHistoryResponse) (proto.Message, error) {
return func(_ context.Context, r *lnrpc.ForwardingHistoryResponse) (
return func(ctx context.Context, r *lnrpc.ForwardingHistoryResponse) (
proto.Message, error) {
fwdEvents := make(
[]*lnrpc.ForwardingEvent, len(r.ForwardingEvents),
)
err := db.Update(func(tx firewalldb.PrivacyMapTx) error {
err := db.Update(ctx, func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
for i, fe := range r.ForwardingEvents {
var err error
@ -487,7 +491,9 @@ func handleFeeReportResponse(db firewalldb.PrivacyMapDB,
chanFees := make([]*lnrpc.ChannelFeeReport, len(r.ChannelFees))
err := db.Update(func(tx firewalldb.PrivacyMapTx) error {
err := db.Update(ctx, func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
var err error
for i, c := range r.ChannelFees {
@ -550,7 +556,9 @@ func handleListChannelsRequest(db firewalldb.PrivacyMapDB,
return r, nil
}
err := db.View(func(tx firewalldb.PrivacyMapTx) error {
err := db.View(ctx, func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
peer, err := firewalldb.RevealBytes(tx, r.Peer)
if err != nil {
return err
@ -572,7 +580,7 @@ func handleListChannelsResponse(db firewalldb.PrivacyMapDB,
randIntn func(int) (int, error)) func(ctx context.Context,
r *lnrpc.ListChannelsResponse) (proto.Message, error) {
return func(_ context.Context, r *lnrpc.ListChannelsResponse) (
return func(ctx context.Context, r *lnrpc.ListChannelsResponse) (
proto.Message, error) {
hidePubkeys := !flags.Contains(session.ClearPubkeys)
@ -580,7 +588,9 @@ func handleListChannelsResponse(db firewalldb.PrivacyMapDB,
channels := make([]*lnrpc.Channel, len(r.Channels))
err := db.Update(func(tx firewalldb.PrivacyMapTx) error {
err := db.Update(ctx, func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
for i, c := range r.Channels {
var err error
@ -745,7 +755,7 @@ func handleUpdatePolicyRequest(db firewalldb.PrivacyMapDB,
flags session.PrivacyFlags) func(ctx context.Context,
r *lnrpc.PolicyUpdateRequest) (proto.Message, error) {
return func(_ context.Context, r *lnrpc.PolicyUpdateRequest) (
return func(ctx context.Context, r *lnrpc.PolicyUpdateRequest) (
proto.Message, error) {
chanPoint := r.GetChanPoint()
@ -764,7 +774,9 @@ func handleUpdatePolicyRequest(db firewalldb.PrivacyMapDB,
newTxid := txid.String()
newIndex := chanPoint.GetOutputIndex()
if !flags.Contains(session.ClearChanIDs) {
err = db.View(func(tx firewalldb.PrivacyMapTx) error {
err = db.View(ctx, func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
var err error
newTxid, newIndex, err = firewalldb.RevealChanPoint(
tx, newTxid, newIndex,
@ -793,7 +805,7 @@ func handleUpdatePolicyResponse(db firewalldb.PrivacyMapDB,
flags session.PrivacyFlags) func(ctx context.Context,
r *lnrpc.PolicyUpdateResponse) (proto.Message, error) {
return func(_ context.Context, r *lnrpc.PolicyUpdateResponse) (
return func(ctx context.Context, r *lnrpc.PolicyUpdateResponse) (
proto.Message, error) {
if flags.Contains(session.ClearChanIDs) {
@ -804,7 +816,9 @@ func handleUpdatePolicyResponse(db firewalldb.PrivacyMapDB,
[]*lnrpc.FailedUpdate, len(r.FailedUpdates),
)
err := db.Update(func(tx firewalldb.PrivacyMapTx) error {
err := db.Update(ctx, func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
for i, u := range r.FailedUpdates {
failedUpdates[i] = &lnrpc.FailedUpdate{
Reason: u.Reason,
@ -926,7 +940,7 @@ func handleClosedChannelsResponse(db firewalldb.PrivacyMapDB,
randIntn func(int) (int, error)) func(ctx context.Context,
r *lnrpc.ClosedChannelsResponse) (proto.Message, error) {
return func(_ context.Context, r *lnrpc.ClosedChannelsResponse) (
return func(ctx context.Context, r *lnrpc.ClosedChannelsResponse) (
proto.Message, error) {
closedChannels := make(
@ -934,7 +948,9 @@ func handleClosedChannelsResponse(db firewalldb.PrivacyMapDB,
len(r.Channels),
)
err := db.Update(func(tx firewalldb.PrivacyMapTx) error {
err := db.Update(ctx, func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
for i, c := range r.Channels {
var err error
@ -1117,7 +1133,7 @@ func handlePendingChannelsResponse(db firewalldb.PrivacyMapDB,
randIntn func(int) (int, error)) func(ctx context.Context,
r *lnrpc.PendingChannelsResponse) (proto.Message, error) {
return func(_ context.Context, r *lnrpc.PendingChannelsResponse) (
return func(ctx context.Context, r *lnrpc.PendingChannelsResponse) (
proto.Message, error) {
pendingOpens := make(
@ -1140,7 +1156,9 @@ func handlePendingChannelsResponse(db firewalldb.PrivacyMapDB,
len(r.WaitingCloseChannels),
)
err := db.Update(func(tx firewalldb.PrivacyMapTx) error {
err := db.Update(ctx, func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
for i, c := range r.PendingOpenChannels {
var err error
@ -1343,12 +1361,14 @@ func handleBatchOpenChannelRequest(db firewalldb.PrivacyMapDB,
flags session.PrivacyFlags) func(ctx context.Context,
r *lnrpc.BatchOpenChannelRequest) (proto.Message, error) {
return func(_ context.Context, r *lnrpc.BatchOpenChannelRequest) (
return func(ctx context.Context, r *lnrpc.BatchOpenChannelRequest) (
proto.Message, error) {
var reqs = make([]*lnrpc.BatchOpenChannel, len(r.Channels))
err := db.View(func(tx firewalldb.PrivacyMapTx) error {
err := db.View(ctx, func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
for i, c := range r.Channels {
var err error
@ -1414,12 +1434,14 @@ func handleBatchOpenChannelResponse(db firewalldb.PrivacyMapDB,
flags session.PrivacyFlags) func(ctx context.Context,
r *lnrpc.BatchOpenChannelResponse) (proto.Message, error) {
return func(_ context.Context, r *lnrpc.BatchOpenChannelResponse) (
return func(ctx context.Context, r *lnrpc.BatchOpenChannelResponse) (
proto.Message, error) {
resps := make([]*lnrpc.PendingUpdate, len(r.PendingChannels))
err := db.Update(func(tx firewalldb.PrivacyMapTx) error {
err := db.Update(ctx, func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
for i, p := range r.PendingChannels {
var (
txIdBytes = p.Txid
@ -1471,14 +1493,15 @@ func handleChannelOpenRequest(db firewalldb.PrivacyMapDB,
flags session.PrivacyFlags) func(ctx context.Context,
r *lnrpc.OpenChannelRequest) (proto.Message, error) {
return func(_ context.Context, r *lnrpc.OpenChannelRequest) (
return func(ctx context.Context, r *lnrpc.OpenChannelRequest) (
proto.Message, error) {
var nodePubkey []byte
err := db.View(func(tx firewalldb.PrivacyMapTx) error {
var err error
err := db.View(ctx, func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
var err error
// We use the byte slice representation of the
// pubkey and fall back to the hex string if present.
nodePubkey = r.NodePubkey
@ -1548,7 +1571,7 @@ func handleChannelOpenResponse(db firewalldb.PrivacyMapDB,
flags session.PrivacyFlags) func(ctx context.Context,
r *lnrpc.ChannelPoint) (proto.Message, error) {
return func(_ context.Context, r *lnrpc.ChannelPoint) (
return func(ctx context.Context, r *lnrpc.ChannelPoint) (
proto.Message, error) {
var (
@ -1556,7 +1579,9 @@ func handleChannelOpenResponse(db firewalldb.PrivacyMapDB,
index uint32
)
err := db.Update(func(tx firewalldb.PrivacyMapTx) error {
err := db.Update(ctx, func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
var err error
txid = r.GetFundingTxidStr()
@ -1622,12 +1647,14 @@ func handleConnectPeerRequest(db firewalldb.PrivacyMapDB,
flags session.PrivacyFlags) func(ctx context.Context,
r *lnrpc.ConnectPeerRequest) (proto.Message, error) {
return func(_ context.Context, r *lnrpc.ConnectPeerRequest) (
return func(ctx context.Context, r *lnrpc.ConnectPeerRequest) (
proto.Message, error) {
var addr *lnrpc.LightningAddress
err := db.View(func(tx firewalldb.PrivacyMapTx) error {
err := db.View(ctx, func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
var err error
// Note, this only works if the pubkey alias was

View file

@ -1073,7 +1073,9 @@ func newMockDB(t *testing.T, preloadRealToPseudo map[string]string,
db := mockDB{privDB: make(map[string]*mockPrivacyMapDB)}
sessDB := db.NewSessionDB(sessID)
_ = sessDB.Update(func(tx firewalldb.PrivacyMapTx) error {
_ = sessDB.Update(context.Background(), func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
for r, p := range preloadRealToPseudo {
require.NoError(t, tx.NewPair(r, p))
}
@ -1107,16 +1109,16 @@ type mockPrivacyMapDB struct {
p2r map[string]string
}
func (m *mockPrivacyMapDB) Update(
f func(tx firewalldb.PrivacyMapTx) error) error {
func (m *mockPrivacyMapDB) Update(ctx context.Context,
f func(ctx context.Context, tx firewalldb.PrivacyMapTx) error) error {
return f(m)
return f(ctx, m)
}
func (m *mockPrivacyMapDB) View(
f func(tx firewalldb.PrivacyMapTx) error) error {
func (m *mockPrivacyMapDB) View(ctx context.Context,
f func(ctx context.Context, tx firewalldb.PrivacyMapTx) error) error {
return f(m)
return f(ctx, m)
}
func (m *mockPrivacyMapDB) NewPair(real, pseudo string) error {

View file

@ -395,7 +395,7 @@ func (r *RuleEnforcer) initRule(ctx context.Context, reqID uint64, name string,
privMap := r.newPrivMap(session.GroupID)
ruleValues, err = ruleValues.PseudoToReal(
privMap, session.PrivacyFlags,
ctx, privMap, session.PrivacyFlags,
)
if err != nil {
return nil, fmt.Errorf("could not prepare rule "+

View file

@ -1,6 +1,7 @@
package firewalldb
import (
"context"
"crypto/rand"
"encoding/binary"
"encoding/hex"
@ -57,13 +58,13 @@ type PrivacyMapDB interface {
// error, the transaction is rolled back. If the rollback fails, the
// original error returned by f is still returned. If the commit fails,
// the commit error is returned.
Update(f func(tx PrivacyMapTx) error) error
Update(context.Context, func(context.Context, PrivacyMapTx) error) error
// View opens a database read transaction and executes the function f
// with the transaction passed as a parameter. After f exits, the
// transaction is rolled back. If f errors, its error is returned, not a
// rollback error (if any occur).
View(f func(tx PrivacyMapTx) error) error
View(context.Context, func(context.Context, PrivacyMapTx) error) error
}
// PrivacyMapTx represents a db that can be used to create, store and fetch
@ -112,7 +113,9 @@ func (p *privacyMapDB) beginTx(writable bool) (*privacyMapTx, error) {
// returned.
//
// NOTE: this is part of the PrivacyMapDB interface.
func (p *privacyMapDB) Update(f func(tx PrivacyMapTx) error) error {
func (p *privacyMapDB) Update(ctx context.Context, f func(ctx context.Context,
tx PrivacyMapTx) error) error {
tx, err := p.beginTx(true)
if err != nil {
return err
@ -125,7 +128,7 @@ func (p *privacyMapDB) Update(f func(tx PrivacyMapTx) error) error {
}
}()
err = f(tx)
err = f(ctx, tx)
if err != nil {
// Want to return the original error, not a rollback error if
// any occur.
@ -142,7 +145,9 @@ func (p *privacyMapDB) Update(f func(tx PrivacyMapTx) error) error {
// occur).
//
// NOTE: this is part of the PrivacyMapDB interface.
func (p *privacyMapDB) View(f func(tx PrivacyMapTx) error) error {
func (p *privacyMapDB) View(ctx context.Context, f func(ctx context.Context,
tx PrivacyMapTx) error) error {
tx, err := p.beginTx(false)
if err != nil {
return err
@ -155,7 +160,7 @@ func (p *privacyMapDB) View(f func(tx PrivacyMapTx) error) error {
}
}()
err = f(tx)
err = f(ctx, tx)
rollbackErr := tx.boltTx.Rollback()
if err != nil {
return err

View file

@ -1,6 +1,7 @@
package firewalldb
import (
"context"
"fmt"
"testing"
@ -9,6 +10,9 @@ import (
// TestPrivacyMapStorage tests the privacy mapper CRUD logic.
func TestPrivacyMapStorage(t *testing.T) {
t.Parallel()
ctx := context.Background()
tmpDir := t.TempDir()
db, err := NewDB(tmpDir, "test.db", nil)
require.NoError(t, err)
@ -18,7 +22,7 @@ func TestPrivacyMapStorage(t *testing.T) {
pdb1 := db.PrivacyDB([4]byte{1, 1, 1, 1})
_ = pdb1.Update(func(tx PrivacyMapTx) error {
_ = pdb1.Update(ctx, func(ctx context.Context, tx PrivacyMapTx) error {
_, err = tx.RealToPseudo("real")
require.ErrorIs(t, err, ErrNoSuchKeyFound)
@ -48,7 +52,7 @@ func TestPrivacyMapStorage(t *testing.T) {
pdb2 := db.PrivacyDB([4]byte{2, 2, 2, 2})
_ = pdb2.Update(func(tx PrivacyMapTx) error {
_ = pdb2.Update(ctx, func(ctx context.Context, tx PrivacyMapTx) error {
_, err = tx.RealToPseudo("real")
require.ErrorIs(t, err, ErrNoSuchKeyFound)
@ -78,7 +82,7 @@ func TestPrivacyMapStorage(t *testing.T) {
pdb3 := db.PrivacyDB([4]byte{3, 3, 3, 3})
_ = pdb3.Update(func(tx PrivacyMapTx) error {
_ = 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()
@ -180,6 +184,9 @@ func TestPrivacyMapStorage(t *testing.T) {
// provide atomic access to the db. If anything fails in the middle of an
// `Update` function, then all the changes prior should be rolled back.
func TestPrivacyMapTxs(t *testing.T) {
t.Parallel()
ctx := context.Background()
tmpDir := t.TempDir()
db, err := NewDB(tmpDir, "test.db", nil)
require.NoError(t, err)
@ -191,7 +198,9 @@ func TestPrivacyMapTxs(t *testing.T) {
// Test that if an action fails midway through the transaction, then
// it is rolled back.
err = pdb1.Update(func(tx PrivacyMapTx) error {
err = pdb1.Update(ctx, func(ctx context.Context,
tx PrivacyMapTx) error {
err := tx.NewPair("real", "pseudo")
if err != nil {
return err
@ -208,7 +217,7 @@ func TestPrivacyMapTxs(t *testing.T) {
})
require.Error(t, err)
err = pdb1.View(func(tx PrivacyMapTx) error {
err = pdb1.View(ctx, func(ctx context.Context, tx PrivacyMapTx) error {
_, err := tx.RealToPseudo("real")
return err
})

View file

@ -396,8 +396,8 @@ func (f *ChanPolicyBounds) RuleName() string {
// find the real values. This is a no-op for the ChanPolicyBounds rule.
//
// NOTE: this is part of the Values interface.
func (f *ChanPolicyBounds) PseudoToReal(_ firewalldb.PrivacyMapDB,
_ session.PrivacyFlags) (Values, error) {
func (f *ChanPolicyBounds) PseudoToReal(_ context.Context,
_ firewalldb.PrivacyMapDB, _ session.PrivacyFlags) (Values, error) {
return f, nil
}
@ -407,8 +407,9 @@ func (f *ChanPolicyBounds) PseudoToReal(_ firewalldb.PrivacyMapDB,
// 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(_ firewalldb.PrivacyMapReader,
_ session.PrivacyFlags) (Values, map[string]string, error) {
func (f *ChanPolicyBounds) RealToPseudo(_ context.Context,
_ firewalldb.PrivacyMapReader, _ session.PrivacyFlags) (Values,
map[string]string, error) {
return f, nil, nil
}

View file

@ -333,8 +333,8 @@ func (v *ChannelConstraint) RuleName() string {
// find the real values. This is a no-op for the ChannelConstraint rule.
//
// NOTE: this is part of the Values interface.
func (v *ChannelConstraint) PseudoToReal(_ firewalldb.PrivacyMapDB,
_ session.PrivacyFlags) (Values, error) {
func (v *ChannelConstraint) PseudoToReal(_ context.Context,
_ firewalldb.PrivacyMapDB, _ session.PrivacyFlags) (Values, error) {
return v, nil
}
@ -344,8 +344,9 @@ func (v *ChannelConstraint) PseudoToReal(_ firewalldb.PrivacyMapDB,
// that should be persisted. This is a no-op for the ChannelConstraint rule.
//
// NOTE: this is part of the Values interface.
func (v *ChannelConstraint) RealToPseudo(_ firewalldb.PrivacyMapReader,
_ session.PrivacyFlags) (Values, map[string]string, error) {
func (v *ChannelConstraint) RealToPseudo(_ context.Context,
_ firewalldb.PrivacyMapReader, _ session.PrivacyFlags) (Values,
map[string]string, error) {
return v, nil, nil
}

View file

@ -336,8 +336,9 @@ func (c *ChannelRestrict) ToProto() *litrpc.RuleValue {
// It constructs a new ChannelRestrict instance with these real channel IDs.
//
// NOTE: this is part of the Values interface.
func (c *ChannelRestrict) PseudoToReal(db firewalldb.PrivacyMapDB,
flags session.PrivacyFlags) (Values, error) {
func (c *ChannelRestrict) PseudoToReal(ctx context.Context,
db firewalldb.PrivacyMapDB, flags session.PrivacyFlags) (Values,
error) {
restrictList := make([]uint64, len(c.DenyList))
@ -348,7 +349,9 @@ func (c *ChannelRestrict) PseudoToReal(db firewalldb.PrivacyMapDB,
return &ChannelRestrict{DenyList: restrictList}, nil
}
err := db.View(func(tx firewalldb.PrivacyMapTx) error {
err := db.View(ctx, func(ctx context.Context,
tx firewalldb.PrivacyMapTx) error {
for i, chanID := range c.DenyList {
real, err := firewalldb.RevealUint64(tx, chanID)
if err != nil {
@ -372,7 +375,8 @@ func (c *ChannelRestrict) PseudoToReal(db firewalldb.PrivacyMapDB,
// not find in the given PrivacyMapReader.
//
// NOTE: this is part of the Values interface.
func (c *ChannelRestrict) RealToPseudo(db firewalldb.PrivacyMapReader,
func (c *ChannelRestrict) RealToPseudo(_ context.Context,
db firewalldb.PrivacyMapReader,
flags session.PrivacyFlags) (Values, map[string]string, error) {
pseudoIDs := make([]uint64, len(c.DenyList))

View file

@ -167,6 +167,9 @@ func (m *mockLndClient) ListChannels(_ context.Context, _, _ bool) (
// 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) {
t.Parallel()
ctx := context.Background()
chanID1 := firewalldb.Uint64ToStr(1)
chanID2 := firewalldb.Uint64ToStr(2)
chanID3 := firewalldb.Uint64ToStr(3)
@ -249,7 +252,7 @@ func TestChannelRestrictRealToPseudo(t *testing.T) {
// form along with any new privacy map pairs that should
// be added to the DB.
v, newPairs, err := cr.RealToPseudo(
privMapPairDB, test.privacyFlags,
ctx, privMapPairDB, test.privacyFlags,
)
require.NoError(t, err)
require.Len(t, newPairs, len(test.expectNewPairs))

View file

@ -256,8 +256,8 @@ func (h *HistoryLimit) GetStartDate() time.Time {
// find the real values. This is a no-op for the HistoryLimit rule.
//
// NOTE: this is part of the Values interface.
func (h *HistoryLimit) PseudoToReal(_ firewalldb.PrivacyMapDB,
_ session.PrivacyFlags) (Values, error) {
func (h *HistoryLimit) PseudoToReal(_ context.Context,
_ firewalldb.PrivacyMapDB, _ session.PrivacyFlags) (Values, error) {
return h, nil
}
@ -267,8 +267,9 @@ func (h *HistoryLimit) PseudoToReal(_ firewalldb.PrivacyMapDB,
// 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(_ firewalldb.PrivacyMapReader,
_ session.PrivacyFlags) (Values, map[string]string, error) {
func (h *HistoryLimit) RealToPseudo(_ context.Context,
_ firewalldb.PrivacyMapReader, _ session.PrivacyFlags) (Values,
map[string]string, error) {
return h, nil, nil
}

View file

@ -64,13 +64,13 @@ type Values interface {
// 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,
RealToPseudo(ctx context.Context, db firewalldb.PrivacyMapReader,
flags session.PrivacyFlags) (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
// passed PrivacyMapDB to find the real values.
PseudoToReal(db firewalldb.PrivacyMapDB,
PseudoToReal(ctx context.Context, db firewalldb.PrivacyMapDB,
flags session.PrivacyFlags) (Values, error)
}

View file

@ -363,8 +363,8 @@ func (o *OnChainBudget) ToProto() *litrpc.RuleValue {
// find the real values. This is a no-op for the OnChainBudget rule.
//
// NOTE: this is part of the Values interface.
func (o *OnChainBudget) PseudoToReal(_ firewalldb.PrivacyMapDB,
_ session.PrivacyFlags) (Values, error) {
func (o *OnChainBudget) PseudoToReal(_ context.Context,
_ firewalldb.PrivacyMapDB, _ session.PrivacyFlags) (Values, error) {
return o, nil
}
@ -374,8 +374,9 @@ func (o *OnChainBudget) PseudoToReal(_ firewalldb.PrivacyMapDB,
// that should be persisted. This is a no-op for the OnChainBudget rule.
//
// NOTE: this is part of the Values interface.
func (o *OnChainBudget) RealToPseudo(db firewalldb.PrivacyMapReader,
flags session.PrivacyFlags) (Values, map[string]string, error) {
func (o *OnChainBudget) RealToPseudo(_ context.Context,
_ firewalldb.PrivacyMapReader, _ session.PrivacyFlags) (Values,
map[string]string, error) {
return o, nil, nil
}

View file

@ -381,8 +381,9 @@ func (c *PeerRestrict) ToProto() *litrpc.RuleValue {
// It constructs a new PeerRestrict instance with these real peer IDs.
//
// NOTE: this is part of the Values interface.
func (c *PeerRestrict) PseudoToReal(db firewalldb.PrivacyMapDB,
flags session.PrivacyFlags) (Values, error) {
func (c *PeerRestrict) PseudoToReal(ctx context.Context,
db firewalldb.PrivacyMapDB, flags session.PrivacyFlags) (Values,
error) {
restrictList := make([]string, len(c.DenyList))
@ -393,7 +394,9 @@ func (c *PeerRestrict) PseudoToReal(db firewalldb.PrivacyMapDB,
return &PeerRestrict{DenyList: restrictList}, nil
}
err := db.View(func(tx firewalldb.PrivacyMapTx) error {
err := db.View(ctx, func(_ context.Context,
tx firewalldb.PrivacyMapTx) error {
for i, peerPubKey := range c.DenyList {
real, err := firewalldb.RevealString(tx, peerPubKey)
if err != nil {
@ -418,7 +421,8 @@ func (c *PeerRestrict) PseudoToReal(db firewalldb.PrivacyMapDB,
// find in the given PrivacyMapReader.
//
// NOTE: this is part of the Values interface.
func (c *PeerRestrict) RealToPseudo(db firewalldb.PrivacyMapReader,
func (c *PeerRestrict) RealToPseudo(_ context.Context,
db firewalldb.PrivacyMapReader,
flags session.PrivacyFlags) (Values, map[string]string, error) {
pseudoIDs := make([]string, len(c.DenyList))

View file

@ -204,6 +204,9 @@ func TestPeerRestrictCheckRequest(t *testing.T) {
// 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) {
t.Parallel()
ctx := context.Background()
tests := []struct {
name string
privacyFlags session.PrivacyFlags
@ -276,7 +279,7 @@ func TestPeerRestrictRealToPseudo(t *testing.T) {
// form along with any new privacy map pairs that should
// be added to the DB.
v, newPairs, err := pr.RealToPseudo(
privMapPairDB, test.privacyFlags,
ctx, privMapPairDB, test.privacyFlags,
)
require.NoError(t, err)
require.Len(t, newPairs, len(test.expectNewPairs))

View file

@ -267,8 +267,8 @@ func (r *RateLimit) ToProto() *litrpc.RuleValue {
// find the real values. This is a no-op for the RateLimit rule.
//
// NOTE: this is part of the Values interface.
func (r *RateLimit) PseudoToReal(_ firewalldb.PrivacyMapDB,
_ session.PrivacyFlags) (Values, error) {
func (r *RateLimit) PseudoToReal(_ context.Context,
_ firewalldb.PrivacyMapDB, _ session.PrivacyFlags) (Values, error) {
return r, nil
}
@ -278,8 +278,9 @@ func (r *RateLimit) PseudoToReal(_ firewalldb.PrivacyMapDB,
// 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(_ firewalldb.PrivacyMapReader,
flags session.PrivacyFlags) (Values, map[string]string, error) {
func (r *RateLimit) RealToPseudo(_ context.Context,
_ firewalldb.PrivacyMapReader, flags session.PrivacyFlags) (Values,
map[string]string, error) {
return r, nil, nil
}

View file

@ -355,7 +355,7 @@ func (s *sessionRpcServer) AddSession(ctx context.Context,
return nil, fmt.Errorf("error fetching session: %v", err)
}
rpcSession, err := s.marshalRPCSession(sess)
rpcSession, err := s.marshalRPCSession(ctx, sess)
if err != nil {
return nil, fmt.Errorf("error marshaling session: %v", err)
}
@ -557,7 +557,7 @@ func (s *sessionRpcServer) ListSessions(ctx context.Context,
Sessions: make([]*litrpc.Session, len(sessions)),
}
for idx, sess := range sessions {
response.Sessions[idx], err = s.marshalRPCSession(sess)
response.Sessions[idx], err = s.marshalRPCSession(ctx, sess)
if err != nil {
return nil, fmt.Errorf("error marshaling session: %v",
err)
@ -629,7 +629,9 @@ func (s *sessionRpcServer) PrivacyMapConversion(ctx context.Context,
var res string
privMap := s.cfg.privMap(groupID)
err = privMap.View(func(tx firewalldb.PrivacyMapTx) error {
err = privMap.View(ctx, func(_ context.Context,
tx firewalldb.PrivacyMapTx) error {
var err error
if req.RealToPseudo {
res, err = tx.RealToPseudo(req.Input)
@ -899,7 +901,9 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
linkedGroupSession = groupSess
privDB := s.cfg.privMap(groupID)
err = privDB.View(func(tx firewalldb.PrivacyMapTx) error {
err = privDB.View(ctx, func(_ context.Context,
tx firewalldb.PrivacyMapTx) error {
knownPrivMapPairs, err = tx.FetchAllPairs()
return err
@ -1002,7 +1006,8 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
if privacy {
var privMapPairs map[string]string
v, privMapPairs, err = v.RealToPseudo(
knownPrivMapPairs, privacyFlags,
ctx, knownPrivMapPairs,
privacyFlags,
)
if err != nil {
return nil, err
@ -1221,7 +1226,9 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
// Register all the privacy map pairs for this session ID.
privDB := s.cfg.privMap(sess.GroupID)
err = privDB.Update(func(tx firewalldb.PrivacyMapTx) error {
err = privDB.Update(ctx, func(_ context.Context,
tx firewalldb.PrivacyMapTx) error {
for r, p := range newPrivMapPairs {
err := tx.NewPair(r, p)
if err != nil {
@ -1272,7 +1279,7 @@ func (s *sessionRpcServer) AddAutopilotSession(ctx context.Context,
return nil, fmt.Errorf("error fetching session: %v", err)
}
rpcSession, err := s.marshalRPCSession(sess)
rpcSession, err := s.marshalRPCSession(ctx, sess)
if err != nil {
return nil, fmt.Errorf("error marshaling session: %v", err)
}
@ -1297,7 +1304,7 @@ func (s *sessionRpcServer) ListAutopilotSessions(ctx context.Context,
Sessions: make([]*litrpc.Session, len(sessions)),
}
for idx, sess := range sessions {
response.Sessions[idx], err = s.marshalRPCSession(sess)
response.Sessions[idx], err = s.marshalRPCSession(ctx, sess)
if err != nil {
return nil, fmt.Errorf("error marshaling session: %v",
err)
@ -1426,8 +1433,8 @@ func marshalPerms(perms map[string][]bakery.Op) []*litrpc.Permissions {
}
// marshalRPCSession converts a session into its RPC counterpart.
func (s *sessionRpcServer) marshalRPCSession(sess *session.Session) (
*litrpc.Session, error) {
func (s *sessionRpcServer) marshalRPCSession(ctx context.Context,
sess *session.Session) (*litrpc.Session, error) {
rpcState, err := marshalRPCState(sess.State)
if err != nil {
@ -1484,7 +1491,8 @@ func (s *sessionRpcServer) marshalRPCSession(sess *session.Session) (
sess.GroupID,
)
val, err = val.PseudoToReal(
db, sess.PrivacyFlags,
ctx, db,
sess.PrivacyFlags,
)
if err != nil {
return nil, err