refactor update & append capabilities

deployment/keystone/capability_management.go

@@ -2,20 +2,21 @@ package keystone
 
 import (
 	"fmt"
+	"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, useMCMS bool) ([]timelock.MCMSWithTimelockProposal, 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, nil
 	}
+	registry := contractSet.CapabilitiesRegistry
 	deduped, err := dedupCapabilities(registry, capabilities)
 	if err != nil {
 		return nil, fmt.Errorf("failed to dedup capabilities: %w", err)
@@ -29,17 +30,26 @@ func AddCapabilities(lggr logger.Logger, registry *kcr.CapabilitiesRegistry, cha
 		err = DecodeErr(kcr.CapabilitiesRegistryABI, err)
 		return nil, fmt.Errorf("failed to add capabilities: %w", err)
 	}
-	var proposals []timelock.MCMSWithTimelockProposal
+	var batch *timelock.BatchChainOperation
 	if !useMCMS {
 		_, err = chain.Confirm(tx)
 		if err != nil {
 			return nil, fmt.Errorf("failed to confirm AddCapabilities confirm transaction %s: %w", tx.Hash().String(), err)
 		}
 		lggr.Info("registered capabilities", "capabilities", deduped)
 	} else {
-		// TODO
+		batch = &timelock.BatchChainOperation{
+			ChainIdentifier: mcms.ChainIdentifier(chain.Selector),
+			Batch: []mcms.Operation{
+				{
+					To:    registry.Address(),
+					Data:  tx.Data(),
+					Value: big.NewInt(0),
+				},
+			},
+		}
 	}
-	return proposals, nil
+	return batch, nil
 }
 
 // CapabilityID returns a unique id for the capability

deployment/keystone/changeset/accept_ownership_test.go

@@ -37,6 +37,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 → deployment/keystone/changeset/append_node_capabilities.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,19 @@ 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,
+		ContractSet:       &contracts,
 		P2pToCapabilities: req.P2pToCapabilities,
+		UseMCMS:           req.UseMCMS,
 	}, nil
 }

deployment/keystone/changeset/append_node_capabilities_test.go

@@ -0,0 +1,132 @@
+package changeset_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/exp/maps"
+
+	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]
+		timelockContracts := map[uint64]*commonchangeset.TimelockExecutionContracts{
+			te.RegistrySelector: {
+				Timelock:  contracts.Timelock,
+				CallProxy: contracts.CallProxy,
+			},
+		}
+
+		_, err = commonchangeset.ApplyChangesets(t, te.Env, timelockContracts, []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/internal/append_node_capabilities.go

@@ -11,8 +11,8 @@ import (
 )
 
 type AppendNodeCapabilitiesRequest struct {
-	Chain    deployment.Chain
-	Registry *kcr.CapabilitiesRegistry
+	Chain       deployment.Chain
+	ContractSet *kslib.ContractSet
 
 	P2pToCapabilities map[p2pkey.PeerID][]kcr.CapabilitiesRegistryCapability
 	UseMCMS           bool
@@ -22,7 +22,7 @@ func (req *AppendNodeCapabilitiesRequest) Validate() error {
 	if len(req.P2pToCapabilities) == 0 {
 		return fmt.Errorf("p2pToCapabilities is empty")
 	}
-	if req.Registry == nil {
+	if req.ContractSet.CapabilitiesRegistry == nil {
 		return fmt.Errorf("registry is nil")
 	}
 	return nil
@@ -32,36 +32,37 @@ func AppendNodeCapabilitiesImpl(lggr logger.Logger, req *AppendNodeCapabilitiesR
 	if err := req.Validate(); err != nil {
 		return nil, fmt.Errorf("failed to validate request: %w", err)
 	}
-	// collect all the capabilities and add them to the registry
-	var capabilities []kcr.CapabilitiesRegistryCapability
-	for _, cap := range req.P2pToCapabilities {
-		capabilities = append(capabilities, cap...)
-	}
-	proposals, err := kslib.AddCapabilities(lggr, req.Registry, req.Chain, capabilities, req.UseMCMS)
-	if err != nil {
-		return nil, fmt.Errorf("failed to add capabilities: %w", err)
-	}
 
 	// for each node, merge the new capabilities with the existing ones and update the node
 	updatesByPeer := make(map[p2pkey.PeerID]NodeUpdate)
 	for p2pID, caps := range req.P2pToCapabilities {
-		caps, err := AppendCapabilities(lggr, req.Registry, req.Chain, []p2pkey.PeerID{p2pID}, caps)
+		caps, err := AppendCapabilities(lggr, req.ContractSet.CapabilitiesRegistry, req.Chain, []p2pkey.PeerID{p2pID}, caps)
 		if err != nil {
 			return nil, fmt.Errorf("failed to append capabilities for p2p %s: %w", p2pID, err)
 		}
 		updatesByPeer[p2pID] = NodeUpdate{Capabilities: caps[p2pID]}
 	}
 
+	// collect all the capabilities and add them to the registry
+	var capabilities []kcr.CapabilitiesRegistryCapability
+	for _, cap := range req.P2pToCapabilities {
+		capabilities = append(capabilities, cap...)
+	}
+	op, err := kslib.AddCapabilities(lggr, req.ContractSet, req.Chain, capabilities, req.UseMCMS)
+	if err != nil {
+		return nil, fmt.Errorf("failed to add capabilities: %w", err)
+	}
+
 	updateNodesReq := &UpdateNodesRequest{
 		Chain:        req.Chain,
-		Registry:     req.Registry,
+		ContractSet:  req.ContractSet,
 		P2pToUpdates: updatesByPeer,
 		UseMCMS:      req.UseMCMS,
+		Ops:          op,
 	}
 	resp, err := UpdateNodes(lggr, updateNodesReq)
 	if err != nil {
 		return nil, fmt.Errorf("failed to update nodes: %w", err)
 	}
-	resp.Proposals = append(proposals, resp.Proposals...)
 	return resp, nil
 }

deployment/keystone/changeset/internal/append_node_capabilities_test.go

@@ -94,8 +94,8 @@ func TestAppendNodeCapabilities(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			setupResp := kstest.SetupTestRegistry(t, lggr, tt.args.initialState)
 
-			tt.args.req.Registry = setupResp.Registry
 			tt.args.req.Chain = setupResp.Chain
+			tt.args.req.ContractSet = setupResp.ContractSet
 
 			got, err := internal.AppendNodeCapabilitiesImpl(tt.args.lggr, tt.args.req)
 			if (err != nil) != tt.wantErr {

deployment/keystone/changeset/internal/test/utils.go

@@ -40,6 +40,7 @@ type SetupTestRegistryResponse struct {
 	Registry         *kcr.CapabilitiesRegistry
 	Chain            deployment.Chain
 	RegistrySelector uint64
+	ContractSet      *kslib.ContractSet
 }
 
 func SetupTestRegistry(t *testing.T, lggr logger.Logger, req *SetupTestRegistryRequest) *SetupTestRegistryResponse {
@@ -100,6 +101,9 @@ func SetupTestRegistry(t *testing.T, lggr logger.Logger, req *SetupTestRegistryR
 		Registry:         registry,
 		Chain:            chain,
 		RegistrySelector: chain.Selector,
+		ContractSet: &kslib.ContractSet{
+			CapabilitiesRegistry: registry,
+		},
 	}
 }