From 21983bab7564833fb78a98ad0d54f2baedbbce09 Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Mon, 13 Jan 2025 07:02:42 +0200
Subject: [PATCH] rules: thread context through

---
 firewall/rule_enforcer.go          | 18 +++++++++---------
 rules/chan_policy_bounds.go        |  4 ++--
 rules/channel_constraints.go       |  4 ++--
 rules/channel_restrictions.go      | 13 +++++++------
 rules/channel_restrictions_test.go |  2 +-
 rules/history_limit.go             |  4 ++--
 rules/interfaces.go                |  3 ++-
 rules/manager_set.go               |  5 +++--
 rules/onchain_budget.go            |  4 ++--
 rules/peer_restrictions.go         | 26 ++++++++++++++------------
 rules/peer_restrictions_test.go    |  2 +-
 rules/rate_limit.go                |  4 ++--
 12 files changed, 47 insertions(+), 42 deletions(-)

diff --git a/firewall/rule_enforcer.go b/firewall/rule_enforcer.go
index e68f3781..964baf32 100644
--- a/firewall/rule_enforcer.go
+++ b/firewall/rule_enforcer.go
@@ -238,7 +238,7 @@ func (r *RuleEnforcer) handleRequest(ctx context.Context,
 		return nil, fmt.Errorf("could not extract ID from macaroon")
 	}
 
-	rules, err := r.collectEnforcers(ri, sessionID)
+	rules, err := r.collectEnforcers(ctx, ri, sessionID)
 	if err != nil {
 		return nil, fmt.Errorf("error parsing rules: %v", err)
 	}
@@ -294,7 +294,7 @@ func (r *RuleEnforcer) handleResponse(ctx context.Context,
 		return nil, fmt.Errorf("could not extract ID from macaroon")
 	}
 
-	enforcers, err := r.collectEnforcers(ri, sessionID)
+	enforcers, err := r.collectEnforcers(ctx, ri, sessionID)
 	if err != nil {
 		return nil, fmt.Errorf("error parsing rules: %v", err)
 	}
@@ -328,7 +328,7 @@ func (r *RuleEnforcer) handleErrorResponse(ctx context.Context,
 		return nil, fmt.Errorf("could not extract ID from macaroon")
 	}
 
-	enforcers, err := r.collectEnforcers(ri, sessionID)
+	enforcers, err := r.collectEnforcers(ctx, ri, sessionID)
 	if err != nil {
 		return nil, fmt.Errorf("error parsing rules: %v", err)
 	}
@@ -353,7 +353,7 @@ func (r *RuleEnforcer) handleErrorResponse(ctx context.Context,
 
 // collectRule initialises and returns all the Rules that need to be enforced
 // for the given request.
-func (r *RuleEnforcer) collectEnforcers(ri *RequestInfo,
+func (r *RuleEnforcer) collectEnforcers(ctx context.Context, ri *RequestInfo,
 	sessionID session.ID) ([]rules.Enforcer, error) {
 
 	ruleEnforcers := make(
@@ -363,8 +363,8 @@ func (r *RuleEnforcer) collectEnforcers(ri *RequestInfo,
 
 	for rule, value := range ri.Rules.FeatureRules[ri.MetaInfo.Feature] {
 		r, err := r.initRule(
-			ri.RequestID, rule, []byte(value), ri.MetaInfo.Feature,
-			sessionID, false, ri.WithPrivacy,
+			ctx, ri.RequestID, rule, []byte(value),
+			ri.MetaInfo.Feature, sessionID, false, ri.WithPrivacy,
 		)
 		if err != nil {
 			return nil, err
@@ -377,8 +377,8 @@ func (r *RuleEnforcer) collectEnforcers(ri *RequestInfo,
 }
 
 // initRule initialises a rule.Rule with any required config values.
-func (r *RuleEnforcer) initRule(reqID uint64, name string, value []byte,
-	featureName string, sessionID session.ID,
+func (r *RuleEnforcer) initRule(ctx context.Context, reqID uint64, name string,
+	value []byte, featureName string, sessionID session.ID,
 	sessionRule, privacy bool) (rules.Enforcer, error) {
 
 	ruleValues, err := r.ruleMgrs.InitRuleValues(name, value)
@@ -425,5 +425,5 @@ func (r *RuleEnforcer) initRule(reqID uint64, name string, value []byte,
 		LndConnID:    r.lndConnID,
 	}
 
-	return r.ruleMgrs.InitEnforcer(cfg, name, ruleValues)
+	return r.ruleMgrs.InitEnforcer(ctx, cfg, name, ruleValues)
 }
diff --git a/rules/chan_policy_bounds.go b/rules/chan_policy_bounds.go
index 7ad3159c..9ba90ded 100644
--- a/rules/chan_policy_bounds.go
+++ b/rules/chan_policy_bounds.go
@@ -39,8 +39,8 @@ func (b *ChanPolicyBoundsMgr) Stop() error {
 // values and config.
 //
 // NOTE: This is part of the Manager interface.
-func (b *ChanPolicyBoundsMgr) NewEnforcer(_ Config, values Values) (Enforcer,
-	error) {
+func (b *ChanPolicyBoundsMgr) NewEnforcer(_ context.Context, _ Config,
+	values Values) (Enforcer, error) {
 
 	bounds, ok := values.(*ChanPolicyBounds)
 	if !ok {
diff --git a/rules/channel_constraints.go b/rules/channel_constraints.go
index 8e287b89..e50e30df 100644
--- a/rules/channel_constraints.go
+++ b/rules/channel_constraints.go
@@ -38,8 +38,8 @@ func (m *ChanConstraintMgr) Stop() error {
 // values and config.
 //
 // NOTE: This is part of the Manager interface.
-func (m *ChanConstraintMgr) NewEnforcer(_ Config, values Values) (Enforcer,
-	error) {
+func (m *ChanConstraintMgr) NewEnforcer(_ context.Context, _ Config,
+	values Values) (Enforcer, error) {
 
 	bounds, ok := values.(*ChannelConstraint)
 	if !ok {
diff --git a/rules/channel_restrictions.go b/rules/channel_restrictions.go
index 8594dde5..745ed85b 100644
--- a/rules/channel_restrictions.go
+++ b/rules/channel_restrictions.go
@@ -60,8 +60,8 @@ func (c *ChannelRestrictMgr) Stop() error {
 // values and config.
 //
 // NOTE: This is part of the Manager interface.
-func (c *ChannelRestrictMgr) NewEnforcer(cfg Config, values Values) (Enforcer,
-	error) {
+func (c *ChannelRestrictMgr) NewEnforcer(ctx context.Context, cfg Config,
+	values Values) (Enforcer, error) {
 
 	channels, ok := values.(*ChannelRestrict)
 	if !ok {
@@ -72,7 +72,8 @@ func (c *ChannelRestrictMgr) NewEnforcer(cfg Config, values Values) (Enforcer,
 	chanMap := make(map[uint64]bool, len(channels.DenyList))
 	for _, chanID := range channels.DenyList {
 		chanMap[chanID] = true
-		if err := c.maybeUpdateChannelMaps(cfg, chanID); err != nil {
+		err := c.maybeUpdateChannelMaps(ctx, cfg, chanID)
+		if err != nil {
 			return nil, err
 		}
 	}
@@ -118,8 +119,8 @@ func (c *ChannelRestrictMgr) EmptyValue() Values {
 
 // maybeUpdateChannelMaps updates the ChannelRestrictMgrs set of known channels
 // iff the channel given by the caller is not found in the current map set.
-func (c *ChannelRestrictMgr) maybeUpdateChannelMaps(cfg Config,
-	chanID uint64) error {
+func (c *ChannelRestrictMgr) maybeUpdateChannelMaps(ctx context.Context,
+	cfg Config, chanID uint64) error {
 
 	c.mu.Lock()
 	defer c.mu.Unlock()
@@ -133,7 +134,7 @@ func (c *ChannelRestrictMgr) maybeUpdateChannelMaps(cfg Config,
 
 	// Fetch a list of our open channels from LND.
 	lnd := cfg.GetLndClient()
-	chans, err := lnd.ListChannels(context.Background(), false, false)
+	chans, err := lnd.ListChannels(ctx, false, false)
 	if err != nil {
 		return err
 	}
diff --git a/rules/channel_restrictions_test.go b/rules/channel_restrictions_test.go
index d6ef6e8c..a12c8091 100644
--- a/rules/channel_restrictions_test.go
+++ b/rules/channel_restrictions_test.go
@@ -53,7 +53,7 @@ func TestChannelRestrictCheckRequest(t *testing.T) {
 			},
 		},
 	}
-	enf, err := mgr.NewEnforcer(cfg, &ChannelRestrict{
+	enf, err := mgr.NewEnforcer(ctx, cfg, &ChannelRestrict{
 		DenyList: []uint64{
 			chanID1, chanID2,
 		},
diff --git a/rules/history_limit.go b/rules/history_limit.go
index 8ca0270f..dccebef4 100644
--- a/rules/history_limit.go
+++ b/rules/history_limit.go
@@ -38,8 +38,8 @@ func (h *HistoryLimitMgr) Stop() error {
 // values and config.
 //
 // NOTE: This is part of the Manager interface.
-func (h *HistoryLimitMgr) NewEnforcer(_ Config, values Values) (Enforcer,
-	error) {
+func (h *HistoryLimitMgr) NewEnforcer(_ context.Context, _ Config,
+	values Values) (Enforcer, error) {
 
 	limit, ok := values.(*HistoryLimit)
 	if !ok {
diff --git a/rules/interfaces.go b/rules/interfaces.go
index 66fc27bf..a1683c4c 100644
--- a/rules/interfaces.go
+++ b/rules/interfaces.go
@@ -16,7 +16,8 @@ import (
 type Manager interface {
 	// NewEnforcer constructs a new rule enforcer using the passed values
 	// and config.
-	NewEnforcer(cfg Config, values Values) (Enforcer, error)
+	NewEnforcer(ctx context.Context, cfg Config, values Values) (Enforcer,
+		error)
 
 	// NewValueFromProto converts the given proto value into a Value object.
 	NewValueFromProto(p *litrpc.RuleValue) (Values, error)
diff --git a/rules/manager_set.go b/rules/manager_set.go
index edc9497d..fce0c95e 100644
--- a/rules/manager_set.go
+++ b/rules/manager_set.go
@@ -1,6 +1,7 @@
 package rules
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 
@@ -32,7 +33,7 @@ func NewRuleManagerSet() ManagerSet {
 
 // InitEnforcer gets the appropriate rule Manager for the given name and uses it
 // to create an appropriate rule Enforcer.
-func (m ManagerSet) InitEnforcer(cfg Config, name string,
+func (m ManagerSet) InitEnforcer(ctx context.Context, cfg Config, name string,
 	values Values) (Enforcer, error) {
 
 	mgr, ok := m[name]
@@ -41,7 +42,7 @@ func (m ManagerSet) InitEnforcer(cfg Config, name string,
 			name)
 	}
 
-	return mgr.NewEnforcer(cfg, values)
+	return mgr.NewEnforcer(ctx, cfg, values)
 }
 
 // GetAllRules returns a map of names of all the rules supported by rule
diff --git a/rules/onchain_budget.go b/rules/onchain_budget.go
index 7cfab158..1024dd08 100644
--- a/rules/onchain_budget.go
+++ b/rules/onchain_budget.go
@@ -63,8 +63,8 @@ func (o *OnChainBudgetMgr) Stop() error {
 // passed values and config.
 //
 // NOTE: This is part of the Manager interface.
-func (o *OnChainBudgetMgr) NewEnforcer(cfg Config, values Values) (Enforcer,
-	error) {
+func (o *OnChainBudgetMgr) NewEnforcer(_ context.Context, cfg Config,
+	values Values) (Enforcer, error) {
 
 	budget, ok := values.(*OnChainBudget)
 	if !ok {
diff --git a/rules/peer_restrictions.go b/rules/peer_restrictions.go
index 4f449f0c..fbaefe94 100644
--- a/rules/peer_restrictions.go
+++ b/rules/peer_restrictions.go
@@ -53,8 +53,8 @@ func (c *PeerRestrictMgr) Stop() error {
 // values and config.
 //
 // NOTE: This is part of the Manager interface.
-func (c *PeerRestrictMgr) NewEnforcer(cfg Config, values Values) (Enforcer,
-	error) {
+func (c *PeerRestrictMgr) NewEnforcer(ctx context.Context, cfg Config,
+	values Values) (Enforcer, error) {
 
 	peers, ok := values.(*PeerRestrict)
 	if !ok {
@@ -65,7 +65,7 @@ func (c *PeerRestrictMgr) NewEnforcer(cfg Config, values Values) (Enforcer,
 	peerMap := make(map[string]bool, len(peers.DenyList))
 	for _, peerID := range peers.DenyList {
 		peerMap[peerID] = true
-		if err := c.maybeUpdateMaps(cfg, peerID); err != nil {
+		if err := c.maybeUpdateMaps(ctx, cfg, peerID); err != nil {
 			return nil, err
 		}
 	}
@@ -112,8 +112,8 @@ func (c *PeerRestrictMgr) EmptyValue() Values {
 
 // maybeUpdateMaps updates the managers peer-to-channel and channel-to-peer maps
 // if the given peer ID is unknown to the manager.
-func (c *PeerRestrictMgr) maybeUpdateMaps(cfg peerRestrictCfg,
-	id string) error {
+func (c *PeerRestrictMgr) maybeUpdateMaps(ctx context.Context,
+	cfg peerRestrictCfg, id string) error {
 
 	c.mu.Lock()
 	defer c.mu.Unlock()
@@ -122,15 +122,17 @@ func (c *PeerRestrictMgr) maybeUpdateMaps(cfg peerRestrictCfg,
 		return nil
 	}
 
-	return c.updateMapsUnsafe(cfg)
+	return c.updateMapsUnsafe(ctx, cfg)
 }
 
 // updateMapsUnsafe updates the manager's peer-to-channel and channel-to-peer
 // maps. It is not thread safe and so must only be called if the manager's
 // mutex is being held.
-func (c *PeerRestrictMgr) updateMapsUnsafe(cfg peerRestrictCfg) error {
+func (c *PeerRestrictMgr) updateMapsUnsafe(ctx context.Context,
+	cfg peerRestrictCfg) error {
+
 	lnd := cfg.GetLndClient()
-	chans, err := lnd.ListChannels(context.Background(), false, false)
+	chans, err := lnd.ListChannels(ctx, false, false)
 	if err != nil {
 		return err
 	}
@@ -152,8 +154,8 @@ func (c *PeerRestrictMgr) updateMapsUnsafe(cfg peerRestrictCfg) error {
 	return nil
 }
 
-func (c *PeerRestrictMgr) getPeerFromChanPoint(cfg peerRestrictCfg,
-	cp string) (string, bool, error) {
+func (c *PeerRestrictMgr) getPeerFromChanPoint(ctx context.Context,
+	cfg peerRestrictCfg, cp string) (string, bool, error) {
 
 	c.mu.Lock()
 	defer c.mu.Unlock()
@@ -163,7 +165,7 @@ func (c *PeerRestrictMgr) getPeerFromChanPoint(cfg peerRestrictCfg,
 		return peer, ok, nil
 	}
 
-	err := c.updateMapsUnsafe(cfg)
+	err := c.updateMapsUnsafe(ctx, cfg)
 	if err != nil {
 		return "", false, err
 	}
@@ -295,7 +297,7 @@ func (c *PeerRestrictEnforcer) checkers() map[string]mid.RoundTripChecker {
 				point := fmt.Sprintf("%s:%d", txid, index)
 
 				peerID, ok, err := c.mgr.getPeerFromChanPoint(
-					c.cfg, point,
+					ctx, c.cfg, point,
 				)
 				if err != nil {
 					return err
diff --git a/rules/peer_restrictions_test.go b/rules/peer_restrictions_test.go
index cb9502e3..faa3c18d 100644
--- a/rules/peer_restrictions_test.go
+++ b/rules/peer_restrictions_test.go
@@ -68,7 +68,7 @@ func TestPeerRestrictCheckRequest(t *testing.T) {
 		},
 	}
 
-	enf, err := mgr.NewEnforcer(cfg, &PeerRestrict{
+	enf, err := mgr.NewEnforcer(ctx, cfg, &PeerRestrict{
 		DenyList: []string{
 			peerID1, peerID2,
 		},
diff --git a/rules/rate_limit.go b/rules/rate_limit.go
index 8f776dcc..4bff4bbe 100644
--- a/rules/rate_limit.go
+++ b/rules/rate_limit.go
@@ -38,8 +38,8 @@ func (r *RateLimitMgr) Stop() error {
 // and config.
 //
 // NOTE: This is part of the Manager interface.
-func (r *RateLimitMgr) NewEnforcer(cfg Config, values Values) (Enforcer,
-	error) {
+func (r *RateLimitMgr) NewEnforcer(_ context.Context, cfg Config,
+	values Values) (Enforcer, error) {
 
 	limits, ok := values.(*RateLimit)
 	if !ok {