[CCIP-3522] address book remove feature

.changeset/breezy-suits-float.md

@@ -0,0 +1,5 @@
+---
+"chainlink": minor
+---
+
+#added address book remove feature

deployment/address_book.go

@@ -3,6 +3,9 @@ package deployment
 import (
 	"fmt"
 	"strings"
+	"sync"
+
+	"golang.org/x/exp/maps"
 
 	"github.com/Masterminds/semver/v3"
 	"github.com/ethereum/go-ethereum/common"
@@ -82,10 +85,12 @@ type AddressBook interface {
 	AddressesForChain(chain uint64) (map[string]TypeAndVersion, error)
 	// Allows for merging address books (e.g. new deployments with existing ones)
 	Merge(other AddressBook) error
+	Remove(ab AddressBook) error
 }
 
 type AddressBookMap struct {
 	AddressesByChain map[uint64]map[string]TypeAndVersion
+	mtx              sync.RWMutex
 }
 
 // Save will save an address for a given chain selector. It will error if there is a conflicting existing address.
@@ -106,6 +111,10 @@ func (m *AddressBookMap) Save(chainSelector uint64, address string, typeAndVersi
 	if typeAndVersion.Type == "" {
 		return fmt.Errorf("type cannot be empty")
 	}
+
+	m.mtx.Lock()
+	defer m.mtx.Unlock()
+
 	if _, exists := m.AddressesByChain[chainSelector]; !exists {
 		// First time chain add, create map
 		m.AddressesByChain[chainSelector] = make(map[string]TypeAndVersion)
@@ -118,18 +127,32 @@ func (m *AddressBookMap) Save(chainSelector uint64, address string, typeAndVersi
 }
 
 func (m *AddressBookMap) Addresses() (map[uint64]map[string]TypeAndVersion, error) {
-	return m.AddressesByChain, nil
+	m.mtx.RLock()
+	defer m.mtx.RUnlock()
+
+	// maps are mutable and pass via a pointer
+	// creating a copy of the map to prevent concurrency
+	// read and changes outside object-bound
+	return m.cloneAddresses(m.AddressesByChain), nil
 }
 
 func (m *AddressBookMap) AddressesForChain(chainSelector uint64) (map[string]TypeAndVersion, error) {
 	_, exists := chainsel.ChainBySelector(chainSelector)
 	if !exists {
 		return nil, errors.Wrapf(ErrInvalidChainSelector, "chain selector %d", chainSelector)
 	}
+
+	m.mtx.RLock()
+	defer m.mtx.RUnlock()
+
 	if _, exists := m.AddressesByChain[chainSelector]; !exists {
 		return nil, errors.Wrapf(ErrChainNotFound, "chain selector %d", chainSelector)
 	}
-	return m.AddressesByChain[chainSelector], nil
+
+	// maps are mutable and pass via a pointer
+	// creating a copy of the map to prevent concurrency
+	// read and changes outside object-bound
+	return maps.Clone(m.AddressesByChain[chainSelector]), nil
 }
 
 // Merge will merge the addresses from another address book into this one.
@@ -139,16 +162,61 @@ func (m *AddressBookMap) Merge(ab AddressBook) error {
 	if err != nil {
 		return err
 	}
-	for chain, chainAddresses := range addresses {
-		for address, typeAndVersions := range chainAddresses {
-			if err := m.Save(chain, address, typeAndVersions); err != nil {
-				return err
+
+	m.mtx.Lock()
+	defer m.mtx.Unlock()
+
+	for chainSelector, chainAddresses := range addresses {
+		for address, typeAndVersion := range chainAddresses {
+			if _, exists := m.AddressesByChain[chainSelector]; !exists {
+				m.AddressesByChain[chainSelector] = make(map[string]TypeAndVersion)
+			}
+			if _, exists := m.AddressesByChain[chainSelector][address]; exists {
+				return fmt.Errorf("address %s already exists for chain %d", address, chainSelector)
+			}
+			m.AddressesByChain[chainSelector][address] = typeAndVersion
+		}
+	}
+	return nil
+}
+
+// Remove removes the address book addresses specified via the argument from the AddressBookMap.
+// Errors if all the addresses in the given address book are not contained in the AddressBookMap.
+func (m *AddressBookMap) Remove(ab AddressBook) error {
+	addresses, err := ab.Addresses()
+	if err != nil {
+		return err
+	}
+
+	m.mtx.Lock()
+	defer m.mtx.Unlock()
+
+	addressNotFound := true
+	for chainSelector, chainAddresses := range addresses {
+		for address, _ := range chainAddresses {
+			if _, exists := m.AddressesByChain[chainSelector][address]; exists {
+				addressNotFound = false
+				delete(m.AddressesByChain[chainSelector], address)
 			}
 		}
 	}
+
+	if addressNotFound {
+		return errors.New("AddressBookMap does not contain any address from the given address book")
+	}
+
 	return nil
 }
 
+// cloneAddresses creates a deep copy of map[uint64]map[string]TypeAndVersion object
+func (m *AddressBookMap) cloneAddresses(input map[uint64]map[string]TypeAndVersion) map[uint64]map[string]TypeAndVersion {
+	result := make(map[uint64]map[string]TypeAndVersion)
+	for chainSelector, chainAddresses := range input {
+		result[chainSelector] = maps.Clone(chainAddresses)
+	}
+	return result
+}
+
 // TODO: Maybe could add an environment argument
 // which would ensure only mainnet/testnet chain selectors are used
 // for further safety?

deployment/address_book_test.go

@@ -2,6 +2,8 @@ package deployment
 
 import (
 	"errors"
+	"math/big"
+	"sync"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -118,3 +120,114 @@ func TestAddressBook_Merge(t *testing.T) {
 		},
 	})
 }
+
+func TestAddressBook_Remove(t *testing.T) {
+	onRamp100 := NewTypeAndVersion("OnRamp", Version1_0_0)
+	onRamp110 := NewTypeAndVersion("OnRamp", Version1_1_0)
+	addr1 := common.HexToAddress("0x1").String()
+	addr2 := common.HexToAddress("0x2").String()
+	addr3 := common.HexToAddress("0x3").String()
+
+	baseAB := NewMemoryAddressBookFromMap(map[uint64]map[string]TypeAndVersion{
+		chainsel.TEST_90000001.Selector: {
+			addr1: onRamp100,
+			addr2: onRamp100,
+		},
+		chainsel.TEST_90000002.Selector: {
+			addr1: onRamp110,
+			addr3: onRamp110,
+		},
+	})
+
+	testAB := NewMemoryAddressBookFromMap(map[uint64]map[string]TypeAndVersion{
+		chainsel.TEST_90000001.Selector: {
+			addr1: onRamp100, // should be removed
+			addr3: onRamp100, // doesn't exist in TEST_90000001.Selector
+		},
+	})
+
+	expectingAB := NewMemoryAddressBookFromMap(map[uint64]map[string]TypeAndVersion{
+		chainsel.TEST_90000001.Selector: {
+			addr2: onRamp100,
+		},
+		chainsel.TEST_90000002.Selector: {
+			addr1: onRamp110,
+			addr3: onRamp110,
+		},
+	})
+
+	// first attempt,
+	require.NoError(t, baseAB.Remove(testAB))
+	require.EqualValues(t, baseAB, expectingAB)
+
+	// Remove method should error in case no address was matched with the given address book.
+	require.Error(t, baseAB.Remove(testAB))
+	require.EqualValues(t, baseAB, expectingAB)
+}
+
+func TestAddressBook_ConcurrencyAndDeadlock(t *testing.T) {
+	onRamp100 := NewTypeAndVersion("OnRamp", Version1_0_0)
+	onRamp110 := NewTypeAndVersion("OnRamp", Version1_1_0)
+
+	baseAB := NewMemoryAddressBookFromMap(map[uint64]map[string]TypeAndVersion{
+		chainsel.TEST_90000001.Selector: {
+			common.BigToAddress(big.NewInt(1)).String(): onRamp100,
+		},
+	})
+
+	// concurrent writes
+	var i int64
+	wg := sync.WaitGroup{}
+	for i = 2; i < 1000; i++ {
+		wg.Add(1)
+		go func(input int64) {
+			require.NoError(t, baseAB.Save(
+				chainsel.TEST_90000001.Selector,
+				common.BigToAddress(big.NewInt(input)).String(),
+				onRamp100,
+			))
+			wg.Done()
+		}(i)
+	}
+
+	// concurrent reads
+	for i = 0; i < 100; i++ {
+		wg.Add(1)
+		go func(input int64) {
+			addresses, err := baseAB.Addresses()
+			require.NoError(t, err)
+			for chainSelector, chainAddresses := range addresses {
+				// concurrent read chainAddresses from Addresses() method
+				for address, _ := range chainAddresses {
+					addresses[chainSelector][address] = onRamp110
+				}
+
+				// concurrent read chainAddresses from AddressesForChain() method
+				chainAddresses, err = baseAB.AddressesForChain(chainSelector)
+				require.NoError(t, err)
+				for address, _ := range chainAddresses {
+					_ = addresses[chainSelector][address]
+				}
+			}
+			require.NoError(t, err)
+			wg.Done()
+		}(i)
+	}
+
+	// concurrent merges, starts from 1001 to avoid address conflicts
+	for i = 1001; i < 1100; i++ {
+		wg.Add(1)
+		go func(input int64) {
+			// concurrent merge
+			additionalAB := NewMemoryAddressBookFromMap(map[uint64]map[string]TypeAndVersion{
+				chainsel.TEST_90000002.Selector: {
+					common.BigToAddress(big.NewInt(input)).String(): onRamp100,
+				},
+			})
+			require.NoError(t, baseAB.Merge(additionalAB))
+			wg.Done()
+		}(i)
+	}
+
+	wg.Wait()
+}

deployment/go.mod

@@ -31,6 +31,7 @@ require (
 	github.com/testcontainers/testcontainers-go v0.34.0
 	go.uber.org/multierr v1.11.0
 	go.uber.org/zap v1.27.0
+	golang.org/x/exp v0.0.0-20241009180824-f66d83c29e7c
 	golang.org/x/sync v0.8.0
 	google.golang.org/grpc v1.67.1
 	google.golang.org/protobuf v1.35.1
@@ -475,7 +476,6 @@ require (
 	go4.org/netipx v0.0.0-20230125063823-8449b0a6169f // indirect
 	golang.org/x/arch v0.11.0 // indirect
 	golang.org/x/crypto v0.28.0 // indirect
-	golang.org/x/exp v0.0.0-20241009180824-f66d83c29e7c // indirect
 	golang.org/x/mod v0.21.0 // indirect
 	golang.org/x/net v0.30.0 // indirect
 	golang.org/x/oauth2 v0.23.0 // indirect