support mcms in forwarder contract config changeset (#15533)

* support mcms in ocr3 contract config changeset

* test working for ocr3 config with mcms

* migrate deployment test to setup test env

* fix test

* refactor forwarder configuration changeset for mcms

deployment/keystone/changeset/deploy_forwarder.go

@@ -11,7 +11,8 @@ var _ deployment.ChangeSet[uint64] = DeployForwarder
 
 // DeployForwarder deploys the KeystoneForwarder contract to all chains in the environment
 // callers must merge the output addressbook with the existing one
-func DeployForwarder(env deployment.Environment, registryChainSel uint64) (deployment.ChangesetOutput, error) {
+// TODO: add selectors to deploy only to specific chains
+func DeployForwarder(env deployment.Environment, _ uint64) (deployment.ChangesetOutput, error) {
 	lggr := env.Logger
 	ab := deployment.NewMemoryAddressBook()
 	for _, chain := range env.Chains {
@@ -25,3 +26,42 @@ func DeployForwarder(env deployment.Environment, registryChainSel uint64) (deplo
 
 	return deployment.ChangesetOutput{AddressBook: ab}, nil
 }
+
+var _ deployment.ChangeSet[ConfigureForwardContractsRequest] = ConfigureForwardContracts
+
+type ConfigureForwardContractsRequest struct {
+	WFDonName string
+	// workflow don node ids in the offchain client. Used to fetch and derive the signer keys
+	WFNodeIDs        []string
+	RegistryChainSel uint64
+
+	UseMCMS bool
+}
+
+func (r ConfigureForwardContractsRequest) Validate() error {
+	if len(r.WFNodeIDs) == 0 {
+		return fmt.Errorf("WFNodeIDs must not be empty")
+	}
+	return nil
+}
+
+func ConfigureForwardContracts(env deployment.Environment, req ConfigureForwardContractsRequest) (deployment.ChangesetOutput, error) {
+	wfDon, err := kslib.NewRegisteredDon(env, kslib.RegisteredDonConfig{
+		NodeIDs:          req.WFNodeIDs,
+		Name:             req.WFDonName,
+		RegistryChainSel: req.RegistryChainSel,
+	})
+	if err != nil {
+		return deployment.ChangesetOutput{}, fmt.Errorf("failed to create registered don: %w", err)
+	}
+	r, err := kslib.ConfigureForwardContracts(&env, kslib.ConfigureForwarderContractsRequest{
+		Dons:    []kslib.RegisteredDon{*wfDon},
+		UseMCMS: req.UseMCMS,
+	})
+	if err != nil {
+		return deployment.ChangesetOutput{}, fmt.Errorf("failed to configure forward contracts: %w", err)
+	}
+	return deployment.ChangesetOutput{
+		Proposals: r.Proposals,
+	}, nil
+}

deployment/keystone/changeset/deploy_forwarder_test.go

@@ -1,14 +1,17 @@
 package changeset_test
 
 import (
+	"fmt"
 	"testing"
 
 	"go.uber.org/zap/zapcore"
 
 	"github.com/stretchr/testify/require"
 
+	"github.com/smartcontractkit/ccip-owner-contracts/pkg/gethwrappers"
 	"github.com/smartcontractkit/chainlink-common/pkg/logger"
 	"github.com/smartcontractkit/chainlink/deployment"
+	commonchangeset "github.com/smartcontractkit/chainlink/deployment/common/changeset"
 	"github.com/smartcontractkit/chainlink/deployment/environment/memory"
 	"github.com/smartcontractkit/chainlink/deployment/keystone/changeset"
 )
@@ -45,3 +48,90 @@ func TestDeployForwarder(t *testing.T) {
 		require.Len(t, oaddrs, 1)
 	})
 }
+
+func TestConfigureForwarders(t *testing.T) {
+	t.Parallel()
+
+	t.Run("no mcms ", func(t *testing.T) {
+		for _, nChains := range []int{1, 3} {
+			name := fmt.Sprintf("nChains=%d", nChains)
+			t.Run(name, func(t *testing.T) {
+				te := SetupTestEnv(t, TestConfig{
+					WFDonConfig:     DonConfig{N: 4},
+					AssetDonConfig:  DonConfig{N: 4},
+					WriterDonConfig: DonConfig{N: 4},
+					NumChains:       nChains,
+				})
+
+				var wfNodes []string
+				for id, _ := range te.WFNodes {
+					wfNodes = append(wfNodes, id)
+				}
+
+				cfg := changeset.ConfigureForwardContractsRequest{
+					WFDonName:        "test-wf-don",
+					WFNodeIDs:        wfNodes,
+					RegistryChainSel: te.RegistrySelector,
+				}
+				csOut, err := changeset.ConfigureForwardContracts(te.Env, cfg)
+				require.NoError(t, err)
+				require.Nil(t, csOut.AddressBook)
+				require.Len(t, csOut.Proposals, 0)
+				// check that forwarder
+				// TODO set up a listener to check that the forwarder is configured
+				contractSet := te.ContractSets()
+				for selector := range te.Env.Chains {
+					cs, ok := contractSet[selector]
+					require.True(t, ok)
+					require.NotNil(t, cs.Forwarder)
+				}
+			})
+		}
+	})
+
+	t.Run("with mcms", func(t *testing.T) {
+		for _, nChains := range []int{1, 3} {
+			name := fmt.Sprintf("nChains=%d", nChains)
+			t.Run(name, func(t *testing.T) {
+				te := SetupTestEnv(t, TestConfig{
+					WFDonConfig:     DonConfig{N: 4},
+					AssetDonConfig:  DonConfig{N: 4},
+					WriterDonConfig: DonConfig{N: 4},
+					NumChains:       nChains,
+					UseMCMS:         true,
+				})
+
+				var wfNodes []string
+				for id, _ := range te.WFNodes {
+					wfNodes = append(wfNodes, id)
+				}
+
+				cfg := changeset.ConfigureForwardContractsRequest{
+					WFDonName:        "test-wf-don",
+					WFNodeIDs:        wfNodes,
+					RegistryChainSel: te.RegistrySelector,
+					UseMCMS:          true,
+				}
+				csOut, err := changeset.ConfigureForwardContracts(te.Env, cfg)
+				require.NoError(t, err)
+				require.Len(t, csOut.Proposals, nChains)
+				require.Nil(t, csOut.AddressBook)
+
+				timelocks := make(map[uint64]*gethwrappers.RBACTimelock)
+				for selector, contractSet := range te.ContractSets() {
+					require.NotNil(t, contractSet.Timelock)
+					timelocks[selector] = contractSet.Timelock
+				}
+				_, err = commonchangeset.ApplyChangesets(t, te.Env, timelocks, []commonchangeset.ChangesetApplication{
+					{
+						Changeset: commonchangeset.WrapChangeSet(changeset.ConfigureForwardContracts),
+						Config:    cfg,
+					},
+				})
+				require.NoError(t, err)
+
+			})
+		}
+	})
+
+}

deployment/keystone/changeset/helpers_test.go

@@ -96,6 +96,7 @@ func (c TestConfig) Validate() error {
 }
 
 type TestEnv struct {
+	t                *testing.T
 	Env              deployment.Environment
 	RegistrySelector uint64
 
@@ -104,6 +105,15 @@ type TestEnv struct {
 	AssetNodes map[string]memory.Node
 }
 
+func (te TestEnv) ContractSets() map[uint64]kslib.ContractSet {
+	r, err := kslib.GetContractSets(te.Env.Logger, &kslib.GetContractSetsRequest{
+		Chains:      te.Env.Chains,
+		AddressBook: te.Env.ExistingAddresses,
+	})
+	require.NoError(te.t, err)
+	return r.ContractSets
+}
+
 // SetupTestEnv sets up a keystone test environment with the given configuration
 func SetupTestEnv(t *testing.T, c TestConfig) TestEnv {
 	require.NoError(t, c.Validate())
@@ -250,57 +260,51 @@ func SetupTestEnv(t *testing.T, c TestConfig) TestEnv {
 
 	if c.UseMCMS {
 		// TODO: mcms on all the chains, currently only on the registry chain. need to fix this for forwarders
-		t.Logf("Enabling MCMS registry chain %d", registryChainSel)
 		// deploy, configure and xfer ownership of MCMS
+		timelockCfgs := make(map[uint64]commontypes.MCMSWithTimelockConfig)
+		for sel := range env.Chains {
+			t.Logf("Enabling MCMS on chain %d", sel)
+			timelockCfgs[sel] = commontypes.MCMSWithTimelockConfig{
+				Canceller:         commonchangeset.SingleGroupMCMS(t),
+				Bypasser:          commonchangeset.SingleGroupMCMS(t),
+				Proposer:          commonchangeset.SingleGroupMCMS(t),
+				TimelockExecutors: env.AllDeployerKeys(),
+				TimelockMinDelay:  big.NewInt(0),
+			}
+		}
 		env, err = commonchangeset.ApplyChangesets(t, env, nil, []commonchangeset.ChangesetApplication{
 			{
 				Changeset: commonchangeset.WrapChangeSet(commonchangeset.DeployMCMSWithTimelock),
-				Config: map[uint64]commontypes.MCMSWithTimelockConfig{
-					registryChainSel: {
-						Canceller:         commonchangeset.SingleGroupMCMS(t),
-						Bypasser:          commonchangeset.SingleGroupMCMS(t),
-						Proposer:          commonchangeset.SingleGroupMCMS(t),
-						TimelockExecutors: env.AllDeployerKeys(),
-						TimelockMinDelay:  big.NewInt(0),
-					},
-				},
+				Config:    timelockCfgs,
 			},
 		})
 		require.NoError(t, err)
 		// extract the MCMS address
 		r, err := kslib.GetContractSets(lggr, &kslib.GetContractSetsRequest{
-			Chains: map[uint64]deployment.Chain{
-				registryChainSel: env.Chains[registryChainSel],
-			},
+			Chains:      env.Chains,
 			AddressBook: env.ExistingAddresses,
 		})
 		require.NoError(t, err)
-		mcms := r.ContractSets[registryChainSel].MCMSWithTimelockState
-		require.NotNil(t, mcms)
-		// transfer ownership of all contracts to the MCMS
-		env, err = commonchangeset.ApplyChangesets(t, env, map[uint64]*gethwrappers.RBACTimelock{registryChainSel: mcms.Timelock}, []commonchangeset.ChangesetApplication{
-			{
-				Changeset: commonchangeset.WrapChangeSet(kschangeset.AcceptAllOwnershipsProposal),
-				Config: &kschangeset.AcceptAllOwnershipRequest{
-					ChainSelector: registryChainSel,
-					MinDelay:      0,
+		for sel := range env.Chains {
+			mcms := r.ContractSets[sel].MCMSWithTimelockState
+			require.NotNil(t, mcms, "MCMS not found on chain %d", sel)
+			require.NoError(t, mcms.Validate())
+
+			// transfer ownership of all contracts to the MCMS
+			env, err = commonchangeset.ApplyChangesets(t, env, map[uint64]*gethwrappers.RBACTimelock{sel: mcms.Timelock}, []commonchangeset.ChangesetApplication{
+				{
+					Changeset: commonchangeset.WrapChangeSet(kschangeset.AcceptAllOwnershipsProposal),
+					Config: &kschangeset.AcceptAllOwnershipRequest{
+						ChainSelector: sel,
+						MinDelay:      0,
+					},
 				},
-			},
-		})
-		require.NoError(t, err)
-		// ensure the MCMS is deployed
-		req = &keystone.GetContractSetsRequest{
-			Chains:      env.Chains,
-			AddressBook: env.ExistingAddresses,
+			})
+			require.NoError(t, err)
 		}
-		contractSetsResp, err = keystone.GetContractSets(lggr, req)
-		require.NoError(t, err)
-		require.Len(t, contractSetsResp.ContractSets, len(env.Chains))
-		// check the mcms contract on registry chain
-		gotMCMS := contractSetsResp.ContractSets[registryChainSel].MCMSWithTimelockState
-		require.NoError(t, gotMCMS.Validate())
 	}
 	return TestEnv{
+		t:                t,
 		Env:              env,
 		RegistrySelector: registryChainSel,
 		WFNodes:          wfNodes,