refactor update capabilities to support mcms

deployment/common/proposalutils/propose.go

@@ -34,7 +34,7 @@ func buildProposalMetadata(
 	return metaDataPerChain, nil
 }
 
-// BuildProposalFromBatches Given batches of operations, we build the metadata and timelock addresses of those opartions
+// BuildProposalFromBatches Given batches of operations, we build the metadata and timelock addresses of those operations
 // We then return a proposal that can be executed and signed
 func BuildProposalFromBatches(
 	timelocksPerChain map[uint64]common.Address,

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/append_node_capbilities.go

@@ -37,7 +37,7 @@ func (req *AppendNodeCapabilitiesRequest) convert(e deployment.Environment) (*in
 	}
 	contracts, 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)

deployment/keystone/changeset/internal/append_node_capabilities.go

@@ -14,6 +14,7 @@ type AppendNodeCapabilitiesRequest struct {
 	Chain    deployment.Chain
 	Registry *kcr.CapabilitiesRegistry
 
+	ContractSet       *kslib.ContractSet
 	P2pToCapabilities map[p2pkey.PeerID][]kcr.CapabilitiesRegistryCapability
 	UseMCMS           bool
 }
@@ -37,7 +38,7 @@ func AppendNodeCapabilitiesImpl(lggr logger.Logger, req *AppendNodeCapabilitiesR
 	for _, cap := range req.P2pToCapabilities {
 		capabilities = append(capabilities, cap...)
 	}
-	proposals, err := kslib.AddCapabilities(lggr, req.Registry, req.Chain, capabilities, req.UseMCMS)
+	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)
 	}
@@ -55,13 +56,14 @@ func AppendNodeCapabilitiesImpl(lggr logger.Logger, req *AppendNodeCapabilitiesR
 	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/update_node_capabilities.go

@@ -11,9 +11,9 @@ import (
 )
 
 type UpdateNodeCapabilitiesImplRequest struct {
-	Chain    deployment.Chain
-	Registry *kcr.CapabilitiesRegistry
-
+	Chain             deployment.Chain
+	Registry          *kcr.CapabilitiesRegistry
+	ContractSet       *kslib.ContractSet
 	P2pToCapabilities map[p2pkey.PeerID][]kcr.CapabilitiesRegistryCapability
 
 	UseMCMS bool
@@ -39,7 +39,7 @@ func UpdateNodeCapabilitiesImpl(lggr logger.Logger, req *UpdateNodeCapabilitiesI
 	for _, cap := range req.P2pToCapabilities {
 		capabilities = append(capabilities, cap...)
 	}
-	proposals, err := kslib.AddCapabilities(lggr, req.Registry, req.Chain, capabilities, req.UseMCMS)
+	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)
 	}
@@ -53,12 +53,13 @@ func UpdateNodeCapabilitiesImpl(lggr logger.Logger, req *UpdateNodeCapabilitiesI
 		Chain:        req.Chain,
 		Registry:     req.Registry,
 		P2pToUpdates: p2pToUpdates,
+		ContractSet:  req.ContractSet,
+		Ops:          op,
 		UseMCMS:      req.UseMCMS,
 	}
 	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/update_nodes.go

@@ -9,17 +9,14 @@ import (
 	"sort"
 
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
-	"github.com/ethereum/go-ethereum/common"
 
-	"github.com/smartcontractkit/ccip-owner-contracts/pkg/gethwrappers"
 	"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"
 	kcr "github.com/smartcontractkit/chainlink/v2/core/gethwrappers/keystone/generated/capabilities_registry"
 	"github.com/smartcontractkit/chainlink/v2/core/services/keystore/keys/p2pkey"
 
 	"github.com/smartcontractkit/chainlink/deployment"
-	"github.com/smartcontractkit/chainlink/deployment/common/proposalutils"
 	kslib "github.com/smartcontractkit/chainlink/deployment/keystone"
 )
 
@@ -37,7 +34,8 @@ type UpdateNodesRequest struct {
 
 	P2pToUpdates map[p2pkey.PeerID]NodeUpdate
 
-	ContractSet kslib.ContractSet // contract set for the given chain
+	ContractSet *kslib.ContractSet // contract set for the given chain
+	Ops         *timelock.BatchChainOperation
 	UseMCMS     bool
 }
 
@@ -89,7 +87,8 @@ func (req *UpdateNodesRequest) Validate() error {
 
 type UpdateNodesResponse struct {
 	NodeParams []kcr.CapabilitiesRegistryNodeParams
-	Proposals  []timelock.MCMSWithTimelockProposal
+	//Proposals  []timelock.MCMSWithTimelockProposal
+	Ops *timelock.BatchChainOperation
 }
 
 // UpdateNodes updates the nodes in the registry
@@ -115,44 +114,32 @@ func UpdateNodes(lggr logger.Logger, req *UpdateNodesRequest) (*UpdateNodesRespo
 		return nil, fmt.Errorf("failed to call UpdateNodes: %w", err)
 	}
 
-	var proposals []timelock.MCMSWithTimelockProposal
+	ops := req.Ops
 	if !req.UseMCMS {
 		_, err = req.Chain.Confirm(tx)
 		if err != nil {
 			return nil, fmt.Errorf("failed to confirm UpdateNodes confirm transaction %s: %w", tx.Hash().String(), err)
 		}
 	} else {
-		ops := timelock.BatchChainOperation{
-			ChainIdentifier: mcms.ChainIdentifier(req.Chain.Selector),
-			Batch: []mcms.Operation{
-				{
-					To:    req.Registry.Address(),
-					Data:  tx.Data(),
-					Value: big.NewInt(0),
-				},
-			},
-		}
-		timelocksPerChain := map[uint64]common.Address{
-			req.Chain.Selector: req.ContractSet.Timelock.Address(),
-		}
-		proposerMCMSes := map[uint64]*gethwrappers.ManyChainMultiSig{
-			req.Chain.Selector: req.ContractSet.ProposerMcm,
+		op := mcms.Operation{
+			To:    req.Registry.Address(),
+			Data:  tx.Data(),
+			Value: big.NewInt(0),
 		}
 
-		proposal, err := proposalutils.BuildProposalFromBatches(
-			timelocksPerChain,
-			proposerMCMSes,
-			[]timelock.BatchChainOperation{ops},
-			"proposal to set update nodes",
-			0,
-		)
-		if err != nil {
-			return nil, fmt.Errorf("failed to build proposal: %w", err)
+		if ops == nil {
+			ops = &timelock.BatchChainOperation{
+				ChainIdentifier: mcms.ChainIdentifier(req.Chain.Selector),
+				Batch: []mcms.Operation{
+					op,
+				},
+			}
+		} else {
+			ops.Batch = append(ops.Batch, op)
 		}
-		proposals = append(proposals, *proposal)
 	}
 
-	return &UpdateNodesResponse{NodeParams: params, Proposals: proposals}, nil
+	return &UpdateNodesResponse{NodeParams: params, Ops: ops}, nil
 }
 
 // AppendCapabilities appends the capabilities to the existing capabilities of the nodes listed in p2pIds in the registry

deployment/keystone/changeset/update_node_capabilities.go

@@ -1,13 +1,16 @@
 package changeset
 
 import (
-	"encoding/json"
 	"fmt"
 	"strconv"
 
+	"github.com/ethereum/go-ethereum/common"
 	chainsel "github.com/smartcontractkit/chain-selectors"
 
+	"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"
 
@@ -50,17 +53,13 @@ type UpdateNodeCapabilitiesRequest = MutateNodeCapabilitiesRequest
 
 // MutateNodeCapabilitiesRequest is a request to change the capabilities of nodes in the registry
 type MutateNodeCapabilitiesRequest struct {
-	AddressBook      deployment.AddressBook
 	RegistryChainSel uint64
 
 	P2pToCapabilities map[p2pkey.PeerID][]kcr.CapabilitiesRegistryCapability
 	UseMCMS           bool
 }
 
 func (req *MutateNodeCapabilitiesRequest) Validate() error {
-	if req.AddressBook == nil {
-		return fmt.Errorf("address book is nil")
-	}
 	if len(req.P2pToCapabilities) == 0 {
 		return fmt.Errorf("p2pToCapabilities is empty")
 	}
@@ -80,39 +79,62 @@ func (req *MutateNodeCapabilitiesRequest) updateNodeCapabilitiesImplRequest(e de
 	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)
+	contractSet, exists := resp.ContractSets[req.RegistryChainSel]
+	if !exists {
+		return nil, fmt.Errorf("contract set not found for chain %d", req.RegistryChainSel)
 	}
 
 	return &internal.UpdateNodeCapabilitiesImplRequest{
 		Chain:             registryChain,
-		Registry:          registry,
+		Registry:          contractSet.CapabilitiesRegistry,
 		P2pToCapabilities: req.P2pToCapabilities,
+		ContractSet:       &contractSet,
 		UseMCMS:           req.UseMCMS,
 	}, nil
 }
 
 // UpdateNodeCapabilities updates the capabilities of nodes in the registry
-func UpdateNodeCapabilities(env deployment.Environment, req *MutateNodeCapabilitiesRequest) (deployment.ChangesetOutput, error) {
+func UpdateNodeCapabilities(env deployment.Environment, req *UpdateNodeCapabilitiesRequest) (deployment.ChangesetOutput, error) {
 	c, err := req.updateNodeCapabilitiesImplRequest(env)
 	if err != nil {
 		return deployment.ChangesetOutput{}, fmt.Errorf("failed to convert request: %w", err)
 	}
 
 	r, err := internal.UpdateNodeCapabilitiesImpl(env.Logger, c)
-	if err == nil {
-		b, err2 := json.Marshal(r)
-		if err2 != nil {
-			env.Logger.Debugf("Updated node capabilities '%s'", b)
+	if err != nil {
+		return deployment.ChangesetOutput{}, fmt.Errorf("failed to update nodes: %w", err)
+	}
+
+	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 deployment.ChangesetOutput{}, err
+	return out, nil
 }