Fix DS MCMS contracts loading issue for legacy contracts

deployment/common/changeset/state/evm.go

@@ -161,17 +161,19 @@ func MaybeLoadMCMSWithTimelockStateDataStoreWithQualifier(env cldf.Environment,
 }
 
 // GetMCMSWithTimelockState loads the MCMSWithTimelockState for a specific chain and qualifier from the DataStore.
+// It filters AddressRefs to avoid key collisions that occur when multiple contract types share the same address (e.g. bypasser and canceller).
 func GetMCMSWithTimelockState(store datastore.AddressRefStore, chain cldf_evm.Chain, qualifier string) (*MCMSWithTimelockState, error) {
-	addressesChain, err := GetAddressTypeVersionByQualifier(store, chain.Selector, qualifier)
-	if err != nil {
-		return nil, fmt.Errorf("failed to load addresses from DataStore for chain %d, qualifier %s: %w", chain.Selector, qualifier, err)
+	filters := []datastore.FilterFunc[datastore.AddressRefKey, datastore.AddressRef]{datastore.AddressRefByChainSelector(chain.Selector)}
+	if qualifier != "" {
+		filters = append(filters, datastore.AddressRefByQualifier(qualifier))
 	}
 
-	state, err := MaybeLoadMCMSWithTimelockChainState(chain, addressesChain)
-	if err != nil {
-		return nil, fmt.Errorf("failed to load MCMSWithTimelock state for chain %d: %w", chain.Selector, err)
+	refs := store.Filter(filters...)
+	if len(refs) == 0 {
+		return nil, fmt.Errorf("no addresses found for chain %d", chain.Selector)
 	}
-	return state, nil
+
+	return MaybeLoadMCMSWithTimelockChainStateFromRefs(chain, refs)
 }
 
 // LoadAddressesFromDataStore loads addresses from DataStore with optional qualifier.
@@ -310,6 +312,113 @@ func MaybeLoadMCMSWithTimelockChainState(chain cldf_evm.Chain, addresses map[str
 	return &state, nil
 }
 
+// MaybeLoadMCMSWithTimelockChainStateFromRefs is the DataStore-native equivalent of MaybeLoadMCMSWithTimelockChainState.
+// It accepts []datastore.AddressRef directly, to preserve entries when multiple contract types share the same address (e.g. bypasser and canceller).
+func MaybeLoadMCMSWithTimelockChainStateFromRefs(chain cldf_evm.Chain, refs []datastore.AddressRef) (*MCMSWithTimelockState, error) {
+	state := MCMSWithTimelockState{}
+	var (
+		// We expect one of each contract on the chain.
+		timelock  = cldf.NewTypeAndVersion(types.RBACTimelock, deployment.Version1_0_0)
+		callProxy = cldf.NewTypeAndVersion(types.CallProxy, deployment.Version1_0_0)
+		proposer  = cldf.NewTypeAndVersion(types.ProposerManyChainMultisig, deployment.Version1_0_0)
+		canceller = cldf.NewTypeAndVersion(types.CancellerManyChainMultisig, deployment.Version1_0_0)
+		bypasser  = cldf.NewTypeAndVersion(types.BypasserManyChainMultisig, deployment.Version1_0_0)
+
+		// the same contract can have different roles
+		multichain    = cldf.NewTypeAndVersion(types.ManyChainMultisig, deployment.Version1_0_0)
+		proposerMCMS  = cldf.NewTypeAndVersion(types.ManyChainMultisig, deployment.Version1_0_0)
+		bypasserMCMS  = cldf.NewTypeAndVersion(types.ManyChainMultisig, deployment.Version1_0_0)
+		cancellerMCMS = cldf.NewTypeAndVersion(types.ManyChainMultisig, deployment.Version1_0_0)
+	)
+
+	proposerMCMS.Labels.Add(types.ProposerRole.String())
+	bypasserMCMS.Labels.Add(types.BypasserRole.String())
+	cancellerMCMS.Labels.Add(types.CancellerRole.String())
+	wantTypes := []cldf.TypeAndVersion{timelock, proposer, canceller, bypasser, callProxy,
+		proposerMCMS, bypasserMCMS, cancellerMCMS,
+	}
+
+	dedupMap := make(map[string]cldf.TypeAndVersion, len(refs))
+	for _, ref := range refs {
+		tv := cldf.TypeAndVersion{
+			Type:    cldf.ContractType(ref.Type),
+			Version: *ref.Version,
+		}
+		if !ref.Labels.IsEmpty() {
+			tv.Labels = cldf.NewLabelSet(ref.Labels.List()...)
+		}
+		dedupMap[ref.Key().String()] = tv
+	}
+
+	// Ensure we either have the bundle or not.
+	_, err := cldf.EnsureDeduped(dedupMap, wantTypes)
+	if err != nil {
+		return nil, fmt.Errorf("unable to check MCMS contracts on chain %s error: %w", chain.Name(), err)
+	}
+
+	for _, ref := range refs {
+		addr := common.HexToAddress(ref.Address)
+		tv := cldf.TypeAndVersion{
+			Type:    cldf.ContractType(ref.Type),
+			Version: *ref.Version,
+		}
+		if !ref.Labels.IsEmpty() {
+			tv.Labels = cldf.NewLabelSet(ref.Labels.List()...)
+		}
+
+		switch {
+		case tv.Type == timelock.Type && tv.Version.String() == timelock.Version.String():
+			tl, err := bindings.NewRBACTimelock(addr, chain.Client)
+			if err != nil {
+				return nil, err
+			}
+			state.Timelock = tl
+		case tv.Type == callProxy.Type && tv.Version.String() == callProxy.Version.String():
+			cp, err := bindings.NewCallProxy(addr, chain.Client)
+			if err != nil {
+				return nil, err
+			}
+			state.CallProxy = cp
+		case tv.Type == proposer.Type && tv.Version.String() == proposer.Version.String():
+			mcms, err := bindings.NewManyChainMultiSig(addr, chain.Client)
+			if err != nil {
+				return nil, err
+			}
+			state.ProposerMcm = mcms
+		case tv.Type == bypasser.Type && tv.Version.String() == bypasser.Version.String():
+			mcms, err := bindings.NewManyChainMultiSig(addr, chain.Client)
+			if err != nil {
+				return nil, err
+			}
+			state.BypasserMcm = mcms
+		case tv.Type == canceller.Type && tv.Version.String() == canceller.Version.String():
+			mcms, err := bindings.NewManyChainMultiSig(addr, chain.Client)
+			if err != nil {
+				return nil, err
+			}
+			state.CancellerMcm = mcms
+		case tv.Type == multichain.Type && tv.Version.String() == multichain.Version.String():
+			// Contract of type ManyChainMultiSig must be labeled to assign to the proper state
+			// field.  If a specifically typed contract already occupies the field, then this
+			// contract will be ignored.
+			mcms, err := bindings.NewManyChainMultiSig(addr, chain.Client)
+			if err != nil {
+				return nil, err
+			}
+			if tv.Labels.Contains(types.ProposerRole.String()) && state.ProposerMcm == nil {
+				state.ProposerMcm = mcms
+			}
+			if tv.Labels.Contains(types.BypasserRole.String()) && state.BypasserMcm == nil {
+				state.BypasserMcm = mcms
+			}
+			if tv.Labels.Contains(types.CancellerRole.String()) && state.CancellerMcm == nil {
+				state.CancellerMcm = mcms
+			}
+		}
+	}
+	return &state, nil
+}
+
 type LinkTokenState struct {
 	LinkToken *link_token.LinkToken
 }

deployment/common/changeset/state/evm_test.go

@@ -393,6 +393,90 @@ func TestAddressesForChain(t *testing.T) {
 	})
 }
 
+func TestSharedAddressWithDifferentLabels(t *testing.T) {
+	selector := chain_selectors.TEST_90000001.Selector
+	env, err := environment.New(t.Context(),
+		environment.WithEVMSimulated(t, []uint64{selector}),
+	)
+	require.NoError(t, err)
+
+	chain := env.BlockChains.EVMChains()[selector]
+
+	sharedMcm := deployMCMEvm(t, chain, &mcmstypes.Config{Quorum: 1, Signers: []common.Address{
+		common.HexToAddress("0x0000000000000000000000000000000000000001"),
+	}})
+	sharedAddress := strings.ToLower(sharedMcm.Address().Hex())
+
+	timelock := deployTimelockEvm(t, chain, big.NewInt(1),
+		common.HexToAddress("0x0000000000000000000000000000000000000004"),
+		[]common.Address{common.HexToAddress("0x0000000000000000000000000000000000000005")},
+		[]common.Address{common.HexToAddress("0x0000000000000000000000000000000000000006")},
+		[]common.Address{common.HexToAddress("0x0000000000000000000000000000000000000007")},
+		[]common.Address{common.HexToAddress("0x0000000000000000000000000000000000000008")},
+	)
+	callProxy := deployCallProxyEvm(t, chain,
+		common.HexToAddress("0x0000000000000000000000000000000000000009"))
+	proposerMcm := deployMCMEvm(t, chain, &mcmstypes.Config{Quorum: 1, Signers: []common.Address{
+		common.HexToAddress("0x0000000000000000000000000000000000000002"),
+	}})
+
+	// Populate DataStore - bypasser and canceller share the same address with different labels
+	store := datastore.NewMemoryDataStore()
+	err = store.Addresses().Add(datastore.AddressRef{
+		ChainSelector: selector,
+		Address:       sharedAddress,
+		Type:          datastore.ContractType(types.ManyChainMultisig),
+		Version:       &deployment.Version1_0_0,
+		Qualifier:     "bypasser",
+		Labels:        datastore.NewLabelSet(types.BypasserRole.String()),
+	})
+	require.NoError(t, err)
+	err = store.Addresses().Add(datastore.AddressRef{
+		ChainSelector: selector,
+		Address:       sharedAddress,
+		Type:          datastore.ContractType(types.ManyChainMultisig),
+		Version:       &deployment.Version1_0_0,
+		Qualifier:     "canceller",
+		Labels:        datastore.NewLabelSet(types.CancellerRole.String()),
+	})
+	require.NoError(t, err)
+	err = store.Addresses().Add(datastore.AddressRef{
+		ChainSelector: selector,
+		Address:       strings.ToLower(timelock.Address().Hex()),
+		Type:          datastore.ContractType(types.RBACTimelock),
+		Version:       &deployment.Version1_0_0,
+	})
+	require.NoError(t, err)
+	err = store.Addresses().Add(datastore.AddressRef{
+		ChainSelector: selector,
+		Address:       strings.ToLower(callProxy.Address().Hex()),
+		Type:          datastore.ContractType(types.CallProxy),
+		Version:       &deployment.Version1_0_0,
+	})
+	require.NoError(t, err)
+	err = store.Addresses().Add(datastore.AddressRef{
+		ChainSelector: selector,
+		Address:       strings.ToLower(proposerMcm.Address().Hex()),
+		Type:          datastore.ContractType(types.ProposerManyChainMultisig),
+		Version:       &deployment.Version1_0_0,
+	})
+	require.NoError(t, err)
+
+	// this should correctly load both bypasser and canceller even though they share the same address
+	state, err := GetMCMSWithTimelockState(store.Seal().Addresses(), chain, "")
+	require.NoError(t, err)
+
+	require.NotNil(t, state.Timelock, "timelock should be loaded")
+	require.NotNil(t, state.CallProxy, "call proxy should be loaded")
+	require.NotNil(t, state.ProposerMcm, "proposer should be loaded")
+	require.NotNil(t, state.BypasserMcm, "bypasser should be loaded despite shared address")
+	require.NotNil(t, state.CancellerMcm, "canceller should be loaded despite shared address")
+
+	// validate that they are the same addresses again
+	require.Equal(t, sharedMcm.Address(), state.BypasserMcm.Address())
+	require.Equal(t, sharedMcm.Address(), state.CancellerMcm.Address())
+}
+
 // ----- helpers -----
 
 func toJSON[T any](t *testing.T, value T) string {