From 01be50785936a7f80ffe31e37554bf7ff1ba8de4 Mon Sep 17 00:00:00 2001
From: Bolek Kulbabinski <1416262+bolekk@users.noreply.github.com>
Date: Mon, 5 Aug 2024 15:57:12 -0700
Subject: [PATCH] [KS-412] Validate called DON membership in TriggerPublisher

Sender needs to actually belong to the DON on behalf of which it is subscribing trigger events.
---
 core/capabilities/remote/trigger_publisher.go      | 14 ++++++++++++++
 core/capabilities/remote/trigger_publisher_test.go |  9 +++++++--
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/core/capabilities/remote/trigger_publisher.go b/core/capabilities/remote/trigger_publisher.go
index c1f2fb32c5a..146b8789689 100644
--- a/core/capabilities/remote/trigger_publisher.go
+++ b/core/capabilities/remote/trigger_publisher.go
@@ -26,6 +26,7 @@ type triggerPublisher struct {
 	capInfo       commoncap.CapabilityInfo
 	capDonInfo    commoncap.DON
 	workflowDONs  map[uint32]commoncap.DON
+	membersCache  map[uint32]map[p2ptypes.PeerID]bool
 	dispatcher    types.Dispatcher
 	messageCache  *messageCache[registrationKey, p2ptypes.PeerID]
 	registrations map[registrationKey]*pubRegState
@@ -54,12 +55,21 @@ func NewTriggerPublisher(config *capabilities.RemoteTriggerConfig, underlying co
 		config = &capabilities.RemoteTriggerConfig{}
 	}
 	config.ApplyDefaults()
+	membersCache := make(map[uint32]map[p2ptypes.PeerID]bool)
+	for id, don := range workflowDONs {
+		cache := make(map[p2ptypes.PeerID]bool)
+		for _, member := range don.Members {
+			cache[member] = true
+		}
+		membersCache[id] = cache
+	}
 	return &triggerPublisher{
 		config:        config,
 		underlying:    underlying,
 		capInfo:       capInfo,
 		capDonInfo:    capDonInfo,
 		workflowDONs:  workflowDONs,
+		membersCache:  membersCache,
 		dispatcher:    dispatcher,
 		messageCache:  NewMessageCache[registrationKey, p2ptypes.PeerID](),
 		registrations: make(map[registrationKey]*pubRegState),
@@ -88,6 +98,10 @@ func (p *triggerPublisher) Receive(_ context.Context, msg *types.MessageBody) {
 			p.lggr.Errorw("received a message from unsupported workflow DON", "capabilityId", p.capInfo.ID, "callerDonId", msg.CallerDonId)
 			return
 		}
+		if !p.membersCache[msg.CallerDonId][sender] {
+			p.lggr.Errorw("sender not a member of its workflow DON", "capabilityId", p.capInfo.ID, "callerDonId", msg.CallerDonId, "sender", sender)
+			return
+		}
 		p.lggr.Debugw("received trigger registration", "capabilityId", p.capInfo.ID, "workflowId", req.Metadata.WorkflowID, "sender", sender)
 		key := registrationKey{msg.CallerDonId, req.Metadata.WorkflowID}
 		nowMs := time.Now().UnixMilli()
diff --git a/core/capabilities/remote/trigger_publisher_test.go b/core/capabilities/remote/trigger_publisher_test.go
index 2c4a8518965..32de37a95aa 100644
--- a/core/capabilities/remote/trigger_publisher_test.go
+++ b/core/capabilities/remote/trigger_publisher_test.go
@@ -7,7 +7,6 @@ import (
 
 	"github.com/stretchr/testify/require"
 
-	"github.com/smartcontractkit/chainlink-common/pkg/capabilities"
 	commoncap "github.com/smartcontractkit/chainlink-common/pkg/capabilities"
 	"github.com/smartcontractkit/chainlink-common/pkg/capabilities/pb"
 	"github.com/smartcontractkit/chainlink/v2/core/capabilities/remote"
@@ -42,7 +41,7 @@ func TestTriggerPublisher_Register(t *testing.T) {
 	}
 
 	dispatcher := remoteMocks.NewDispatcher(t)
-	config := &capabilities.RemoteTriggerConfig{
+	config := &commoncap.RemoteTriggerConfig{
 		RegistrationRefresh:     100 * time.Millisecond,
 		RegistrationExpiry:      100 * time.Second,
 		MinResponsesToAggregate: 1,
@@ -73,6 +72,12 @@ func TestTriggerPublisher_Register(t *testing.T) {
 		Payload:     marshaled,
 	}
 	publisher.Receive(ctx, regEvent)
+	// node p1 is not a member of the workflow DON so registration shoudn't happen
+	require.Empty(t, underlying.registrationsCh)
+
+	regEvent.Sender = p2[:]
+	publisher.Receive(ctx, regEvent)
+	require.NotEmpty(t, underlying.registrationsCh)
 	forwarded := <-underlying.registrationsCh
 	require.Equal(t, capRequest.Metadata.WorkflowID, forwarded.Metadata.WorkflowID)