feat(keystone): support multiple OCR3 contracts

deployment/keystone/changeset/deploy_ocr3.go

@@ -6,6 +6,7 @@ import (
 	"io"
 
 	"github.com/ethereum/go-ethereum/common"
+
 	"github.com/smartcontractkit/ccip-owner-contracts/pkg/gethwrappers"
 	"github.com/smartcontractkit/ccip-owner-contracts/pkg/proposal/timelock"
 
@@ -37,6 +38,7 @@ var _ deployment.ChangeSet[ConfigureOCR3Config] = ConfigureOCR3Contract
 type ConfigureOCR3Config struct {
 	ChainSel             uint64
 	NodeIDs              []string
+	Address              *common.Address // address of the OCR3 contract to configure
 	OCR3Config           *kslib.OracleConfig
 	DryRun               bool
 	WriteGeneratedConfig io.Writer // if not nil, write the generated config to this writer as JSON [OCR2OracleConfig]
@@ -54,6 +56,7 @@ func ConfigureOCR3Contract(env deployment.Environment, cfg ConfigureOCR3Config)
 		ChainSel:   cfg.ChainSel,
 		NodeIDs:    cfg.NodeIDs,
 		OCR3Config: cfg.OCR3Config,
+		Address:    cfg.Address,
 		DryRun:     cfg.DryRun,
 		UseMCMS:    cfg.UseMCMS(),
 	})

deployment/keystone/changeset/deploy_ocr3_test.go

@@ -7,6 +7,7 @@
 
 	"go.uber.org/zap/zapcore"
 
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -85,6 +86,162 @@
 		assert.Nil(t, csOut.Proposals)
 	})
 
+	t.Run("success multiple OCR3 contracts", func(t *testing.T) {
+		te := test.SetupTestEnv(t, test.TestConfig{
+			WFDonConfig:     test.DonConfig{N: 4},
+			AssetDonConfig:  test.DonConfig{N: 4},
+			WriterDonConfig: test.DonConfig{N: 4},
+			NumChains:       1,
+		})
+
+		registrySel := te.Env.AllChainSelectors()[0]
+
+		existingContracts, err := te.Env.ExistingAddresses.AddressesForChain(registrySel)
+		require.NoError(t, err)
+		require.Len(t, existingContracts, 4)
+
+		// Find existing OCR3 contract
+		var existingOCR3Addr string
+		for addr, tv := range existingContracts {
+			if tv.Type == internal.OCR3Capability {
+				existingOCR3Addr = addr
+				break
+			}
+		}
+
+		// Deploy a new OCR3 contract
+		resp, err := changeset.DeployOCR3(te.Env, registrySel)
+		require.NoError(t, err)
+		require.NotNil(t, resp)
+		require.NoError(t, te.Env.ExistingAddresses.Merge(resp.AddressBook))
+
+		// Verify after merge there are three original contracts plus one new one
+		addrs, err := te.Env.ExistingAddresses.AddressesForChain(registrySel)
+		require.NoError(t, err)
+		require.Len(t, addrs, 5)
+
+		// Find new OCR3 contract
+		var newOCR3Addr string
+		for addr, tv := range addrs {
+			if tv.Type == internal.OCR3Capability && addr != existingOCR3Addr {
+				newOCR3Addr = addr
+				break
+			}
+		}
+
+		var wfNodes []string
+		for id, _ := range te.WFNodes {
+			wfNodes = append(wfNodes, id)
+		}
+
+		na := common.HexToAddress(newOCR3Addr)
+		w := &bytes.Buffer{}
+		cfg := changeset.ConfigureOCR3Config{
+			ChainSel:             te.RegistrySelector,
+			NodeIDs:              wfNodes,
+			Address:              &na, // Use the new OCR3 contract to configure
+			OCR3Config:           &c,
+			WriteGeneratedConfig: w,
+		}
+
+		csOut, err := changeset.ConfigureOCR3Contract(te.Env, cfg)
+		require.NoError(t, err)
+		var got internal.OCR2OracleConfig
+		err = json.Unmarshal(w.Bytes(), &got)
+		require.NoError(t, err)
+		assert.Len(t, got.Signers, 4)
+		assert.Len(t, got.Transmitters, 4)
+		assert.Nil(t, csOut.Proposals)
+	})
+
+	t.Run("fails multiple OCR3 contracts but unspecified address", func(t *testing.T) {
+		te := test.SetupTestEnv(t, test.TestConfig{
+			WFDonConfig:     test.DonConfig{N: 4},
+			AssetDonConfig:  test.DonConfig{N: 4},
+			WriterDonConfig: test.DonConfig{N: 4},
+			NumChains:       1,
+		})
+
+		registrySel := te.Env.AllChainSelectors()[0]
+
+		existingContracts, err := te.Env.ExistingAddresses.AddressesForChain(registrySel)
+		require.NoError(t, err)
+		require.Len(t, existingContracts, 4)
+
+		// Deploy a new OCR3 contract
+		resp, err := changeset.DeployOCR3(te.Env, registrySel)
+		require.NoError(t, err)
+		require.NotNil(t, resp)
+		require.NoError(t, te.Env.ExistingAddresses.Merge(resp.AddressBook))
+
+		// Verify after merge there are original contracts plus one new one
+		addrs, err := te.Env.ExistingAddresses.AddressesForChain(registrySel)
+		require.NoError(t, err)
+		require.Len(t, addrs, 5)
+
+		var wfNodes []string
+		for id, _ := range te.WFNodes {
+			wfNodes = append(wfNodes, id)
+		}
+
+		w := &bytes.Buffer{}
+		cfg := changeset.ConfigureOCR3Config{
+			ChainSel:             te.RegistrySelector,
+			NodeIDs:              wfNodes,
+			OCR3Config:           &c,
+			WriteGeneratedConfig: w,
+		}
+
+		_, err = changeset.ConfigureOCR3Contract(te.Env, cfg)
+		require.Error(t, err)
+		require.ErrorContains(t, err, "OCR contract address is unspecified")
+	})
+
+	t.Run("fails multiple OCR3 contracts but address not found", func(t *testing.T) {
+		te := test.SetupTestEnv(t, test.TestConfig{
+			WFDonConfig:     test.DonConfig{N: 4},
+			AssetDonConfig:  test.DonConfig{N: 4},
+			WriterDonConfig: test.DonConfig{N: 4},
+			NumChains:       1,
+		})
+
+		registrySel := te.Env.AllChainSelectors()[0]
+
+		existingContracts, err := te.Env.ExistingAddresses.AddressesForChain(registrySel)
+		require.NoError(t, err)
+		require.Len(t, existingContracts, 4)
+
+		// Deploy a new OCR3 contract
+		resp, err := changeset.DeployOCR3(te.Env, registrySel)
+		require.NoError(t, err)
+		require.NotNil(t, resp)
+		require.NoError(t, te.Env.ExistingAddresses.Merge(resp.AddressBook))
+
+		// Verify after merge there are original contracts plus one new one
+		addrs, err := te.Env.ExistingAddresses.AddressesForChain(registrySel)
+		require.NoError(t, err)
+		require.Len(t, addrs, 5)
+
+		var wfNodes []string
+		for id, _ := range te.WFNodes {
+			wfNodes = append(wfNodes, id)
+		}
+
+		nfa := common.HexToAddress("0x1234567890123456789012345678901234567890")
+		w := &bytes.Buffer{}
+		cfg := changeset.ConfigureOCR3Config{
+			ChainSel:             te.RegistrySelector,
+			NodeIDs:              wfNodes,
+			OCR3Config:           &c,
+			Address:              &nfa,
+			WriteGeneratedConfig: w,
+		}
+
+		_, err = changeset.ConfigureOCR3Contract(te.Env, cfg)
+		require.Error(t, err)
+		require.ErrorContains(t, err, "not found in contract set")
+	})
+
 	t.Run("mcms", func(t *testing.T) {
 		te := test.SetupTestEnv(t, test.TestConfig{
 			WFDonConfig:     test.DonConfig{N: 4},

deployment/keystone/changeset/internal/deploy.go

@@ -14,6 +14,7 @@ import (
 	"time"
 
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
+	"github.com/ethereum/go-ethereum/common"
 	"golang.org/x/exp/maps"
 
 	"github.com/smartcontractkit/ccip-owner-contracts/pkg/proposal/mcms"
@@ -30,6 +31,7 @@ import (
 
 	capabilities_registry "github.com/smartcontractkit/chainlink/v2/core/gethwrappers/keystone/generated/capabilities_registry_1_1_0"
 	kf "github.com/smartcontractkit/chainlink/v2/core/gethwrappers/keystone/generated/forwarder_1_0_0"
+	ocr3_capability "github.com/smartcontractkit/chainlink/v2/core/gethwrappers/keystone/generated/ocr3_capability_1_0_0"
 
 	"github.com/smartcontractkit/chainlink-common/pkg/logger"
 )
@@ -321,12 +323,13 @@ func ConfigureOCR3Contract(env *deployment.Environment, chainSel uint64, dons []
 		if !ok {
 			return fmt.Errorf("failed to get contract set for chain %d", chainSel)
 		}
-		contract := contracts.OCR3
-		if contract == nil {
-			return fmt.Errorf("no ocr3 contract found for chain %d", chainSel)
+
+		contract, err := getOCR3Contract(contracts.OCR3, nil)
+		if err != nil {
+			return fmt.Errorf("failed to get OCR3 contract: %w", err)
 		}
 
-		_, err := configureOCR3contract(configureOCR3Request{
+		_, err = configureOCR3contract(configureOCR3Request{
 			cfg:         cfg,
 			chain:       registryChain,
 			contract:    contract,
@@ -349,6 +352,7 @@ type ConfigureOCR3Resp struct {
 type ConfigureOCR3Config struct {
 	ChainSel   uint64
 	NodeIDs    []string
+	Address    *common.Address // address of the OCR3 contract to configure
 	OCR3Config *OracleConfig
 	DryRun     bool
 
@@ -377,10 +381,12 @@ func ConfigureOCR3ContractFromJD(env *deployment.Environment, cfg ConfigureOCR3C
 	if !ok {
 		return nil, fmt.Errorf("failed to get contract set for chain %d", cfg.ChainSel)
 	}
-	contract := contracts.OCR3
-	if contract == nil {
-		return nil, fmt.Errorf("no ocr3 contract found for chain %d", cfg.ChainSel)
+
+	contract, err := getOCR3Contract(contracts.OCR3, cfg.Address)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get OCR3 contract: %w", err)
 	}
+
 	nodes, err := deployment.NodeInfo(cfg.NodeIDs, env.Offchain)
 	if err != nil {
 		return nil, err
@@ -963,3 +969,30 @@ func configureForwarder(lggr logger.Logger, chain deployment.Chain, contractSet
 	}
 	return opMap, nil
 }
+
+// getOCR3Contract returns the OCR3 contract from the contract set.  By default, it returns the only
+// contract in the set if there is no address specified.  If an address is specified, it returns the
+// contract with that address.  If the address is specified but not found in the contract set, it returns
+// an error.
+func getOCR3Contract(contracts map[common.Address]*ocr3_capability.OCR3Capability, addr *common.Address) (*ocr3_capability.OCR3Capability, error) {
+	// Fail if the OCR3 contract address is unspecified and there are multiple OCR3 contracts
+	if addr == nil && len(contracts) > 1 {
+		return nil, errors.New("OCR contract address is unspecified")
+	}
+
+	// Use the first OCR3 contract if the address is unspecified
+	if addr == nil && len(contracts) == 1 {
+		// use the first OCR3 contract
+		for _, c := range contracts {
+			return c, nil
+		}
+	}
+
+	// Select the OCR3 contract by address
+	if contract, ok := contracts[*addr]; ok {
+		return contract, nil
+	}
+
+	// Fail if the OCR3 contract address is specified but not found in the contract set
+	return nil, fmt.Errorf("OCR3 contract address %s not found in contract set", *addr)
+}

deployment/keystone/changeset/internal/state.go

@@ -28,7 +28,7 @@ type GetContractSetsResponse struct {
 
 type ContractSet struct {
 	commonchangeset.MCMSWithTimelockState
-	OCR3                 *ocr3_capability.OCR3Capability
+	OCR3                 map[common.Address]*ocr3_capability.OCR3Capability
 	Forwarder            *forwarder.KeystoneForwarder
 	CapabilitiesRegistry *capabilities_registry.CapabilitiesRegistry
 	WorkflowRegistry     *workflow_registry.WorkflowRegistry
@@ -37,7 +37,9 @@ type ContractSet struct {
 func (cs ContractSet) TransferableContracts() []common.Address {
 	var out []common.Address
 	if cs.OCR3 != nil {
-		out = append(out, cs.OCR3.Address())
+		for _, ocr := range cs.OCR3 {
+			out = append(out, ocr.Address())
+		}
 	}
 	if cs.Forwarder != nil {
 		out = append(out, cs.Forwarder.Address())
@@ -109,7 +111,10 @@ func loadContractSet(lggr logger.Logger, chain deployment.Chain, addresses map[s
 			if err != nil {
 				return nil, fmt.Errorf("failed to create OCR3Capability contract from address %s: %w", addr, err)
 			}
-			out.OCR3 = c
+			if out.OCR3 == nil {
+				out.OCR3 = make(map[common.Address]*ocr3_capability.OCR3Capability)
+			}
+			out.OCR3[common.HexToAddress(addr)] = c
 		case WorkflowRegistry:
 			c, err := workflow_registry.NewWorkflowRegistry(common.HexToAddress(addr), chain.Client)
 			if err != nil {