refactor update & append capabilities

deployment/keystone/capability_management.go

@@ -2,69 +2,91 @@ package keystone
 
 import (
 	"fmt"
-	"strings"
+	"math/big"
 
+	"github.com/smartcontractkit/ccip-owner-contracts/pkg/proposal/mcms"
+	"github.com/smartcontractkit/ccip-owner-contracts/pkg/proposal/timelock"
 	"github.com/smartcontractkit/chainlink-common/pkg/logger"
 	"github.com/smartcontractkit/chainlink/deployment"
 	kcr "github.com/smartcontractkit/chainlink/v2/core/gethwrappers/keystone/generated/capabilities_registry"
 )
 
 // AddCapabilities adds the capabilities to the registry
-// it tries to add all capabilities in one go, if that fails, it falls back to adding them one by one
-func AddCapabilities(lggr logger.Logger, registry *kcr.CapabilitiesRegistry, chain deployment.Chain, capabilities []kcr.CapabilitiesRegistryCapability) error {
+func AddCapabilities(lggr logger.Logger, contractSet *ContractSet, chain deployment.Chain, capabilities []kcr.CapabilitiesRegistryCapability, useMCMS bool) (*timelock.BatchChainOperation, error) {
 	if len(capabilities) == 0 {
-		return nil
+		return nil, nil
 	}
-	// dedup capabilities
-	var deduped []kcr.CapabilitiesRegistryCapability
-	seen := make(map[string]struct{})
-	for _, cap := range capabilities {
-		if _, ok := seen[CapabilityID(cap)]; !ok {
-			seen[CapabilityID(cap)] = struct{}{}
-			deduped = append(deduped, cap)
-		}
+	registry := contractSet.CapabilitiesRegistry
+	deduped, err := dedupCapabilities(registry, capabilities)
+	if err != nil {
+		return nil, fmt.Errorf("failed to dedup capabilities: %w", err)
 	}
-
-	tx, err := registry.AddCapabilities(chain.DeployerKey, deduped)
+	txOpts := chain.DeployerKey
+	if useMCMS {
+		txOpts = deployment.SimTransactOpts()
+	}
+	tx, err := registry.AddCapabilities(txOpts, deduped)
 	if err != nil {
 		err = DecodeErr(kcr.CapabilitiesRegistryABI, err)
-		// no typed errors in the abi, so we have to do string matching
-		// try to add all capabilities in one go, if that fails, fall back to 1-by-1
-		if !strings.Contains(err.Error(), "CapabilityAlreadyExists") {
-			return fmt.Errorf("failed to call AddCapabilities: %w", err)
-		}
-		lggr.Warnw("capabilities already exist, falling back to 1-by-1", "capabilities", deduped)
-		for _, cap := range deduped {
-			tx, err = registry.AddCapabilities(chain.DeployerKey, []kcr.CapabilitiesRegistryCapability{cap})
-			if err != nil {
-				err = DecodeErr(kcr.CapabilitiesRegistryABI, err)
-				if strings.Contains(err.Error(), "CapabilityAlreadyExists") {
-					lggr.Warnw("capability already exists, skipping", "capability", cap)
-					continue
-				}
-				return fmt.Errorf("failed to call AddCapabilities for capability %v: %w", cap, err)
-			}
-			// 1-by-1 tx is pending and we need to wait for it to be mined
-			_, err = chain.Confirm(tx)
-			if err != nil {
-				return fmt.Errorf("failed to confirm AddCapabilities confirm transaction %s: %w", tx.Hash().String(), err)
-			}
-			lggr.Debugw("registered capability", "capability", cap)
-
-		}
-	} else {
-		// the bulk add tx is pending and we need to wait for it to be mined
+		return nil, fmt.Errorf("failed to add capabilities: %w", err)
+	}
+	var batch *timelock.BatchChainOperation
+	if !useMCMS {
 		_, err = chain.Confirm(tx)
 		if err != nil {
-			return fmt.Errorf("failed to confirm AddCapabilities confirm transaction %s: %w", tx.Hash().String(), err)
+			return nil, fmt.Errorf("failed to confirm AddCapabilities confirm transaction %s: %w", tx.Hash().String(), err)
 		}
 		lggr.Info("registered capabilities", "capabilities", deduped)
+	} else {
+		batch = &timelock.BatchChainOperation{
+			ChainIdentifier: mcms.ChainIdentifier(chain.Selector),
+			Batch: []mcms.Operation{
+				{
+					To:    registry.Address(),
+					Data:  tx.Data(),
+					Value: big.NewInt(0),
+				},
+			},
+		}
 	}
-	return nil
+	return batch, nil
 }
 
 // CapabilityID returns a unique id for the capability
 // TODO: mv to chainlink-common? ref https://github.com/smartcontractkit/chainlink/blob/4fb06b4525f03c169c121a68defa9b13677f5f20/contracts/src/v0.8/keystone/CapabilitiesRegistry.sol#L170
 func CapabilityID(c kcr.CapabilitiesRegistryCapability) string {
 	return fmt.Sprintf("%s@%s", c.LabelledName, c.Version)
 }
+
+// dedupCapabilities deduplicates the capabilities
+// dedup capabilities with respect to the registry
+// contract reverts on adding the same capability twice and that would cause the whole transaction to revert
+// which is very bad for us for mcms
+func dedupCapabilities(registry *kcr.CapabilitiesRegistry, capabilities []kcr.CapabilitiesRegistryCapability) ([]kcr.CapabilitiesRegistryCapability, error) {
+	var out []kcr.CapabilitiesRegistryCapability
+	existing, err := registry.GetCapabilities(nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to call GetCapabilities: %w", err)
+	}
+	existingByID := make(map[[32]byte]struct{})
+	for _, cap := range existing {
+		existingByID[cap.HashedId] = struct{}{}
+	}
+	seen := make(map[string]struct{})
+	for _, candidate := range capabilities {
+		h, err := registry.GetHashedCapabilityId(nil, candidate.LabelledName, candidate.Version)
+		if err != nil {
+			return nil, fmt.Errorf("failed to call GetHashedCapabilityId: %w", err)
+		}
+		// dedup input capabilities
+		if _, exists := seen[CapabilityID(candidate)]; exists {
+			continue
+		}
+		seen[CapabilityID(candidate)] = struct{}{}
+		// dedup with respect to the registry
+		if _, exists := existingByID[h]; !exists {
+			out = append(out, candidate)
+		}
+	}
+	return out, nil
+}

deployment/keystone/changeset/accept_ownership_test.go

@@ -38,6 +38,10 @@ func TestAcceptAllOwnership(t *testing.T) {
 			Changeset: commonchangeset.WrapChangeSet(changeset.DeployForwarder),
 			Config:    registrySel,
 		},
+		{
+			Changeset: commonchangeset.WrapChangeSet(changeset.DeployFeedsConsumer),
+			Config:    &changeset.DeployFeedsConsumerRequest{ChainSelector: registrySel},
+		},
 		{
 			Changeset: commonchangeset.WrapChangeSet(commonchangeset.DeployMCMSWithTimelock),
 			Config: map[uint64]types.MCMSWithTimelockConfig{

deployment/keystone/changeset/append_node_capbilities.go

@@ -3,7 +3,11 @@ package changeset
 import (
 	"fmt"
 
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/smartcontractkit/ccip-owner-contracts/pkg/gethwrappers"
+	"github.com/smartcontractkit/ccip-owner-contracts/pkg/proposal/timelock"
 	"github.com/smartcontractkit/chainlink/deployment"
+	"github.com/smartcontractkit/chainlink/deployment/common/proposalutils"
 	kslib "github.com/smartcontractkit/chainlink/deployment/keystone"
 	"github.com/smartcontractkit/chainlink/deployment/keystone/changeset/internal"
 )
@@ -16,15 +20,39 @@ type AppendNodeCapabilitiesRequest = MutateNodeCapabilitiesRequest
 // AppendNodeCapabilities adds any new capabilities to the registry, merges the new capabilities with the existing capabilities
 // of the node, and updates the nodes in the registry host the union of the new and existing capabilities.
 func AppendNodeCapabilities(env deployment.Environment, req *AppendNodeCapabilitiesRequest) (deployment.ChangesetOutput, error) {
-	cfg, err := req.convert(env)
+	c, err := req.convert(env)
 	if err != nil {
 		return deployment.ChangesetOutput{}, err
 	}
-	_, err = internal.AppendNodeCapabilitiesImpl(env.Logger, cfg)
+	r, err := internal.AppendNodeCapabilitiesImpl(env.Logger, c)
 	if err != nil {
 		return deployment.ChangesetOutput{}, err
 	}
-	return deployment.ChangesetOutput{}, nil
+	out := deployment.ChangesetOutput{}
+	if req.UseMCMS {
+		if r.Ops == nil {
+			return out, fmt.Errorf("expected MCMS operation to be non-nil")
+		}
+		timelocksPerChain := map[uint64]common.Address{
+			c.Chain.Selector: c.ContractSet.Timelock.Address(),
+		}
+		proposerMCMSes := map[uint64]*gethwrappers.ManyChainMultiSig{
+			c.Chain.Selector: c.ContractSet.ProposerMcm,
+		}
+
+		proposal, err := proposalutils.BuildProposalFromBatches(
+			timelocksPerChain,
+			proposerMCMSes,
+			[]timelock.BatchChainOperation{*r.Ops},
+			"proposal to set update node capabilities",
+			0,
+		)
+		if err != nil {
+			return out, fmt.Errorf("failed to build proposal: %w", err)
+		}
+		out.Proposals = []timelock.MCMSWithTimelockProposal{*proposal}
+	}
+	return out, nil
 }
 
 func (req *AppendNodeCapabilitiesRequest) convert(e deployment.Environment) (*internal.AppendNodeCapabilitiesRequest, error) {
@@ -35,21 +63,20 @@ func (req *AppendNodeCapabilitiesRequest) convert(e deployment.Environment) (*in
 	if !ok {
 		return nil, fmt.Errorf("registry chain selector %d does not exist in environment", req.RegistryChainSel)
 	}
-	contracts, err := kslib.GetContractSets(e.Logger, &kslib.GetContractSetsRequest{
+	resp, err := kslib.GetContractSets(e.Logger, &kslib.GetContractSetsRequest{
 		Chains:      map[uint64]deployment.Chain{req.RegistryChainSel: registryChain},
-		AddressBook: req.AddressBook,
+		AddressBook: e.ExistingAddresses,
 	})
 	if err != nil {
 		return nil, fmt.Errorf("failed to get contract sets: %w", err)
 	}
-	registry := contracts.ContractSets[req.RegistryChainSel].CapabilitiesRegistry
-	if registry == nil {
-		return nil, fmt.Errorf("capabilities registry not found for chain %d", req.RegistryChainSel)
-	}
+	contracts := resp.ContractSets[req.RegistryChainSel]
 
 	return &internal.AppendNodeCapabilitiesRequest{
 		Chain:             registryChain,
-		Registry:          registry,
+		Registry:          contracts.CapabilitiesRegistry,
+		ContractSet:       &contracts,
 		P2pToCapabilities: req.P2pToCapabilities,
+		UseMCMS:           req.UseMCMS,
 	}, nil
 }

deployment/keystone/changeset/append_node_capbilities_test.go

@@ -0,0 +1,129 @@
+package changeset_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/exp/maps"
+
+	"github.com/smartcontractkit/ccip-owner-contracts/pkg/gethwrappers"
+	commonchangeset "github.com/smartcontractkit/chainlink/deployment/common/changeset"
+	"github.com/smartcontractkit/chainlink/deployment/keystone/changeset"
+	kcr "github.com/smartcontractkit/chainlink/v2/core/gethwrappers/keystone/generated/capabilities_registry"
+	"github.com/smartcontractkit/chainlink/v2/core/services/keystore/keys/p2pkey"
+)
+
+func TestAppendNodeCapabilities(t *testing.T) {
+	t.Parallel()
+
+	var (
+		capA = kcr.CapabilitiesRegistryCapability{
+			LabelledName: "capA",
+			Version:      "0.4.2",
+		}
+		capB = kcr.CapabilitiesRegistryCapability{
+			LabelledName: "capB",
+			Version:      "3.16.0",
+		}
+		caps = []kcr.CapabilitiesRegistryCapability{capA, capB}
+	)
+	t.Run("no mcms", func(t *testing.T) {
+		te := SetupTestEnv(t, TestConfig{
+			WFDonConfig:     DonConfig{N: 4},
+			AssetDonConfig:  DonConfig{N: 4},
+			WriterDonConfig: DonConfig{N: 4},
+			NumChains:       1,
+		})
+
+		newCapabilities := make(map[p2pkey.PeerID][]kcr.CapabilitiesRegistryCapability)
+		for id, _ := range te.WFNodes {
+			k, err := p2pkey.MakePeerID(id)
+			require.NoError(t, err)
+			newCapabilities[k] = caps
+		}
+
+		t.Run("succeeds if existing capabilities not explicit", func(t *testing.T) {
+			cfg := changeset.AppendNodeCapabilitiesRequest{
+				RegistryChainSel:  te.RegistrySelector,
+				P2pToCapabilities: newCapabilities,
+			}
+
+			csOut, err := changeset.AppendNodeCapabilities(te.Env, &cfg)
+			require.NoError(t, err)
+			require.Len(t, csOut.Proposals, 0)
+			require.Nil(t, csOut.AddressBook)
+
+			validateCapabilityAppends(t, te, newCapabilities)
+		})
+	})
+	t.Run("with mcms", func(t *testing.T) {
+		te := SetupTestEnv(t, TestConfig{
+			WFDonConfig:     DonConfig{N: 4},
+			AssetDonConfig:  DonConfig{N: 4},
+			WriterDonConfig: DonConfig{N: 4},
+			NumChains:       1,
+			UseMCMS:         true,
+		})
+
+		newCapabilities := make(map[p2pkey.PeerID][]kcr.CapabilitiesRegistryCapability)
+		for id, _ := range te.WFNodes {
+			k, err := p2pkey.MakePeerID(id)
+			require.NoError(t, err)
+			newCapabilities[k] = caps
+		}
+
+		cfg := changeset.AppendNodeCapabilitiesRequest{
+			RegistryChainSel:  te.RegistrySelector,
+			P2pToCapabilities: newCapabilities,
+			UseMCMS:           true,
+		}
+
+		csOut, err := changeset.AppendNodeCapabilities(te.Env, &cfg)
+		require.NoError(t, err)
+		require.Len(t, csOut.Proposals, 1)
+		require.Len(t, csOut.Proposals[0].Transactions, 1)
+		require.Len(t, csOut.Proposals[0].Transactions[0].Batch, 2) // add capabilities, update nodes
+		require.Nil(t, csOut.AddressBook)
+
+		// now apply the changeset such that the proposal is signed and execed
+		contracts := te.ContractSets()[te.RegistrySelector]
+		timelocks := map[uint64]*gethwrappers.RBACTimelock{
+			te.RegistrySelector: contracts.Timelock,
+		}
+		_, err = commonchangeset.ApplyChangesets(t, te.Env, timelocks, []commonchangeset.ChangesetApplication{
+			{
+				Changeset: commonchangeset.WrapChangeSet(changeset.AppendNodeCapabilities),
+				Config:    &cfg,
+			},
+		})
+		require.NoError(t, err)
+		validateCapabilityAppends(t, te, newCapabilities)
+	})
+
+}
+
+// validateUpdate checks reads nodes from the registry and checks they have the expected updates
+func validateCapabilityAppends(t *testing.T, te TestEnv, appended map[p2pkey.PeerID][]kcr.CapabilitiesRegistryCapability) {
+	registry := te.ContractSets()[te.RegistrySelector].CapabilitiesRegistry
+	wfP2PIDs := p2pIDs(t, maps.Keys(te.WFNodes))
+	nodes, err := registry.GetNodesByP2PIds(nil, wfP2PIDs)
+	require.NoError(t, err)
+	require.Len(t, nodes, len(wfP2PIDs))
+	for _, node := range nodes {
+		want := appended[node.P2pId]
+		require.NotNil(t, want)
+		assertContainsCapabilities(t, registry, want, node)
+	}
+}
+
+func assertContainsCapabilities(t *testing.T, registry *kcr.CapabilitiesRegistry, want []kcr.CapabilitiesRegistryCapability, got kcr.INodeInfoProviderNodeInfo) {
+	wantHashes := make([][32]byte, len(want))
+	for i, c := range want {
+		h, err := registry.GetHashedCapabilityId(nil, c.LabelledName, c.Version)
+		require.NoError(t, err)
+		wantHashes[i] = h
+		assert.Contains(t, got.HashedCapabilityIds, h, "missing capability %v", c)
+	}
+	assert.LessOrEqual(t, len(want), len(got.HashedCapabilityIds))
+}

deployment/keystone/changeset/deploy_ocr3_test.go

@@ -46,7 +46,6 @@ func TestDeployOCR3(t *testing.T) {
 
 func TestConfigureOCR3(t *testing.T) {
 	t.Parallel()
-	lggr := logger.Test(t)
 
 	c := kslib.OracleConfigWithSecrets{
 		OracleConfig: kslib.OracleConfig{
@@ -122,13 +121,9 @@ func TestConfigureOCR3(t *testing.T) {
 		assert.NotNil(t, csOut.Proposals)
 		t.Logf("got: %v", csOut.Proposals[0])
 
-		contractSetsResp, err := kslib.GetContractSets(lggr, &kslib.GetContractSetsRequest{
-			Chains:      te.Env.Chains,
-			AddressBook: te.Env.ExistingAddresses,
-		})
-		require.NoError(t, err)
+		contracts := te.ContractSets()[te.RegistrySelector]
 		var timelocks = map[uint64]*gethwrappers.RBACTimelock{
-			te.RegistrySelector: contractSetsResp.ContractSets[te.RegistrySelector].Timelock,
+			te.RegistrySelector: contracts.Timelock,
 		}
 		// now apply the changeset such that the proposal is signed and execed
 		w2 := &bytes.Buffer{}