From bfec5730b2d6d628371ff03e0c198cafdd06ed97 Mon Sep 17 00:00:00 2001 From: connorwstein Date: Tue, 6 Aug 2024 17:14:17 -0400 Subject: [PATCH] More flexible address book --- core/capabilities/ccip/deployment/deploy.go | 21 +++-- .../ccip/deployment/deploy_test.go | 7 +- core/capabilities/ccip/deployment/state.go | 12 ++- core/environment/address_book.go | 47 ++++++++++ core/environment/address_book_test.go | 1 + core/environment/environment.go | 28 +++--- core/environment/memory/address_book.go | 25 ----- core/environment/persistent/address_book.go | 94 ------------------- .../persistent/address_book_test.go | 44 --------- 9 files changed, 89 insertions(+), 190 deletions(-) create mode 100644 core/environment/address_book.go create mode 100644 core/environment/address_book_test.go delete mode 100644 core/environment/memory/address_book.go delete mode 100644 core/environment/persistent/address_book.go delete mode 100644 core/environment/persistent/address_book_test.go diff --git a/core/capabilities/ccip/deployment/deploy.go b/core/capabilities/ccip/deployment/deploy.go index a11354e67d..95b569eca7 100644 --- a/core/capabilities/ccip/deployment/deploy.go +++ b/core/capabilities/ccip/deployment/deploy.go @@ -57,13 +57,20 @@ func GenerateJobSpecs(capReg common.Address) CCIPSpec { return CCIPSpec{} } +type DeployCCIPContractConfig struct { + Weth9s map[uint64]common.Address + // TODO: More params as needed +} + // TODO: Likely we'll want to further parameterize the deployment // For example a list of contracts to skip deploying if they already exist. // Or mock vs real RMN. -func DeployCCIPContracts(e environment.Environment) error { +// Deployment produces an address book of everything it deployed. +func DeployCCIPContracts(e environment.Environment, c DeployCCIPContractConfig) (environment.AddressBook, error) { + ab := environment.NewMemoryAddressBook() for _, chain := range e.Chains { saveToChain := func(addr common.Address) error { - return e.AddressBook.Save(chain.Selector, addr.String()) + return ab.Save(chain.Selector, addr.String()) } // TODO: Still waiting for RMNRemote/RMNHome contracts etc. @@ -77,7 +84,7 @@ func DeployCCIPContracts(e environment.Environment) error { }, chain.Confirm, saveToChain) if err != nil { e.Logger.Errorw("Failed to deploy mockARM", "err", err) - return err + return ab, err } e.Logger.Infow("deployed mockARM", "addr", mockARM) @@ -92,7 +99,7 @@ func DeployCCIPContracts(e environment.Environment) error { }, chain.Confirm, saveToChain) if err != nil { e.Logger.Errorw("Failed to deploy armProxy", "err", err) - return err + return ab, err } e.Logger.Infow("deployed armProxy", "addr", armProxy) @@ -121,7 +128,7 @@ func DeployCCIPContracts(e environment.Environment) error { }, chain.Confirm, saveToChain) if err != nil { e.Logger.Errorw("Failed to deploy router", "err", err) - return err + return ab, err } e.Logger.Infow("deployed router", "addr", routerAddr) @@ -134,9 +141,9 @@ func DeployCCIPContracts(e environment.Environment) error { }, chain.Confirm, saveToChain) if err != nil { e.Logger.Errorw("Failed to deploy token admin registry", "err", err) - return err + return ab, err } e.Logger.Infow("deployed tokenAdminRegistry", "addr", tokenAdminRegistry) } - return nil + return ab, nil } diff --git a/core/capabilities/ccip/deployment/deploy_test.go b/core/capabilities/ccip/deployment/deploy_test.go index b2101cd5d4..74c53373ec 100644 --- a/core/capabilities/ccip/deployment/deploy_test.go +++ b/core/capabilities/ccip/deployment/deploy_test.go @@ -16,15 +16,16 @@ func TestDeployCCIPContracts(t *testing.T) { Nodes: 1, }) // Deploy all the CCIP contracts. - require.NoError(t, DeployCCIPContracts(e)) - state, err := GenerateOnchainState(e) + ab, err := DeployCCIPContracts(e, DeployCCIPContractConfig{}) + require.NoError(t, err) + state, err := GenerateOnchainState(e, ab) require.NoError(t, err) snap, err := state.Snapshot(e.AllChainSelectors()) require.NoError(t, err) // Assert expect every deployed address to be in the address book. for name, chain := range snap.Chains { - addrs, err := e.AddressBook.Addresses() + addrs, err := ab.Addresses() require.NoError(t, err) evmChainID, _ := chainsel.ChainIdFromName(name) sel, _ := chainsel.SelectorFromChainId(evmChainID) diff --git a/core/capabilities/ccip/deployment/state.go b/core/capabilities/ccip/deployment/state.go index 8fe98fcdfe..63ab995823 100644 --- a/core/capabilities/ccip/deployment/state.go +++ b/core/capabilities/ccip/deployment/state.go @@ -70,7 +70,15 @@ func (s CCIPOnChainState) Snapshot(chains []uint64) (CCIPSnapShot, error) { return snapshot, nil } -func GenerateOnchainState(e environment.Environment) (CCIPOnChainState, error) { +func SnapshotState(e environment.Environment, ab environment.AddressBook) (CCIPSnapShot, error) { + state, err := GenerateOnchainState(e, ab) + if err != nil { + return CCIPSnapShot{}, err + } + return state.Snapshot(e.AllChainSelectors()) +} + +func GenerateOnchainState(e environment.Environment, ab environment.AddressBook) (CCIPOnChainState, error) { state := CCIPOnChainState{ EvmOnRampsV160: make(map[uint64]*evm_2_evm_multi_onramp.EVM2EVMMultiOnRamp), EvmOffRampsV160: make(map[uint64]*evm_2_evm_multi_offramp.EVM2EVMMultiOffRamp), @@ -82,7 +90,7 @@ func GenerateOnchainState(e environment.Environment) (CCIPOnChainState, error) { Weth9s: make(map[uint64]*weth9.WETH9), } // Get all the onchain state - addresses, err := e.AddressBook.Addresses() + addresses, err := ab.Addresses() if err != nil { return state, errors.Wrap(err, "could not get addresses") } diff --git a/core/environment/address_book.go b/core/environment/address_book.go new file mode 100644 index 0000000000..9ff5e0d9da --- /dev/null +++ b/core/environment/address_book.go @@ -0,0 +1,47 @@ +package environment + +type AddressBook interface { + // TODO: Need manualTV override + Save(chainSelector uint64, address string) error + Addresses() (map[uint64]map[string]struct{}, error) + // Allows for merging address books + Merge(other AddressBook) error +} + +type AddressBookMap struct { + AddressesByChain map[uint64]map[string]struct{} +} + +func (m *AddressBookMap) Save(chainSelector uint64, address string) error { + if _, exists := m.AddressesByChain[chainSelector]; !exists { + m.AddressesByChain[chainSelector] = make(map[string]struct{}) + } else { + // Error? + } + m.AddressesByChain[chainSelector][address] = struct{}{} + return nil +} + +func (m *AddressBookMap) Addresses() (map[uint64]map[string]struct{}, error) { + return m.AddressesByChain, nil +} + +// Attention this will mutate existing book +func (m *AddressBookMap) Merge(ab AddressBook) error { + addresses, err := ab.Addresses() + if err != nil { + return err + } + for chain, chainAddresses := range addresses { + for address := range chainAddresses { + return m.Save(chain, address) + } + } + return nil +} + +func NewMemoryAddressBook() *AddressBookMap { + return &AddressBookMap{ + AddressesByChain: make(map[uint64]map[string]struct{}), + } +} diff --git a/core/environment/address_book_test.go b/core/environment/address_book_test.go new file mode 100644 index 0000000000..f2d9ed195f --- /dev/null +++ b/core/environment/address_book_test.go @@ -0,0 +1 @@ +package environment diff --git a/core/environment/environment.go b/core/environment/environment.go index 686face2ff..a870b39362 100644 --- a/core/environment/environment.go +++ b/core/environment/environment.go @@ -15,11 +15,9 @@ import ( "github.com/smartcontractkit/chainlink/v2/core/environment/memory" ) -type AddressBook interface { - // TODO: Need manualTV override - Save(chainSelector uint64, address string) error - Addresses() (map[uint64]map[string]struct{}, error) -} +const ( + Memory = "memory" +) type OnchainClient interface { // For EVM specifically we can use existing geth interface @@ -43,11 +41,11 @@ type Chain struct { } type Environment struct { - Chains map[uint64]Chain - NodeIds []string - AddressBook AddressBook - Offchain OffchainClient - Logger logger.Logger + Name string + Chains map[uint64]Chain + NodeIds []string + Offchain OffchainClient + Logger logger.Logger } func (e Environment) AllChainSelectors() []uint64 { @@ -92,10 +90,10 @@ func NewMemoryEnvironment(t *testing.T, config MemoryEnvironmentConfig) Environm lggr, err := logger.New() require.NoError(t, err) return Environment{ - Offchain: memory.NewMemoryJobClient(nodesByPeerID), - Chains: chains, - NodeIds: keys, - AddressBook: memory.NewMemoryAddressBook(), - Logger: lggr, + Name: Memory, + Offchain: memory.NewMemoryJobClient(nodesByPeerID), + Chains: chains, + NodeIds: keys, + Logger: lggr, } } diff --git a/core/environment/memory/address_book.go b/core/environment/memory/address_book.go deleted file mode 100644 index 771fa304ab..0000000000 --- a/core/environment/memory/address_book.go +++ /dev/null @@ -1,25 +0,0 @@ -package memory - -type AddressBook struct { - AddressesByChain map[uint64]map[string]struct{} -} - -func (m *AddressBook) Save(chainSelector uint64, address string) error { - if _, exists := m.AddressesByChain[chainSelector]; !exists { - m.AddressesByChain[chainSelector] = make(map[string]struct{}) - } else { - // Error? - } - m.AddressesByChain[chainSelector][address] = struct{}{} - return nil -} - -func (m *AddressBook) Addresses() (map[uint64]map[string]struct{}, error) { - return m.AddressesByChain, nil -} - -func NewMemoryAddressBook() *AddressBook { - return &AddressBook{ - AddressesByChain: make(map[uint64]map[string]struct{}), - } -} diff --git a/core/environment/persistent/address_book.go b/core/environment/persistent/address_book.go deleted file mode 100644 index 2b551393d3..0000000000 --- a/core/environment/persistent/address_book.go +++ /dev/null @@ -1,94 +0,0 @@ -package persistent - -import ( - "encoding/json" - "fmt" - "github.com/pkg/errors" - "io/ioutil" - "os" - "sync" -) - -type AddressBook struct { - filePath string - mu sync.Mutex -} - -// Save stores an address in the address book. -func (m *AddressBook) Save(chainSelector uint64, address string) error { - m.mu.Lock() - defer m.mu.Unlock() - - addressesByChain, err := m.loadFromFile() - if err != nil { - return err - } - - if _, exists := addressesByChain[chainSelector]; !exists { - addressesByChain[chainSelector] = make(map[string]struct{}) - } - addressesByChain[chainSelector][address] = struct{}{} - - return m.saveToFile(addressesByChain) -} - -// Addresses returns all addresses. -func (m *AddressBook) Addresses() (map[uint64]map[string]struct{}, error) { - m.mu.Lock() - defer m.mu.Unlock() - return m.loadFromFile() -} - -// saveToFile writes the address book to the file. -func (m *AddressBook) saveToFile(addressesByChain map[uint64]map[string]struct{}) error { - // Make json friendly - addressLists := make(map[uint64][]string) - for chain, addressMp := range addressesByChain { - for addr := range addressMp { - addressLists[chain] = append(addressLists[chain], addr) - } - } - data, err := json.Marshal(addressLists) - if err != nil { - return err - } - return ioutil.WriteFile(m.filePath, data, 0644) -} - -// loadFromFile loads the address book from the file. -func (m *AddressBook) loadFromFile() (map[uint64]map[string]struct{}, error) { - file, err := os.Open(m.filePath) - if err != nil { - if errors.Is(err, os.ErrNotExist) { - return make(map[uint64]map[string]struct{}), err - } - return nil, err - } - defer file.Close() - - data, err := ioutil.ReadAll(file) - if err != nil { - return nil, err - } - - addressLists := make(map[uint64][]string) - err = json.Unmarshal(data, &addressLists) - if err != nil { - return nil, errors.Wrap(err, fmt.Sprintf("failed to unmarshal address book %s %s", m.filePath, data)) - } - addressesByChain := make(map[uint64]map[string]struct{}) - for chain, addresses := range addressLists { - addressesByChain[chain] = make(map[string]struct{}) - for _, address := range addresses { - addressesByChain[chain][address] = struct{}{} - } - } - return addressesByChain, nil -} - -// NewFileBackedAddressBook creates a new AddressBook with file storage. -func NewAddressBook(filePath string) *AddressBook { - return &AddressBook{ - filePath: filePath, - } -} diff --git a/core/environment/persistent/address_book_test.go b/core/environment/persistent/address_book_test.go deleted file mode 100644 index 94cf960ec2..0000000000 --- a/core/environment/persistent/address_book_test.go +++ /dev/null @@ -1,44 +0,0 @@ -package persistent - -import ( - "github.com/stretchr/testify/require" - "gotest.tools/v3/assert" - "os" - "testing" -) - -func TestAddressBook(t *testing.T) { - fn := "blah.json" - _, err := os.Create(fn) - require.NoError(t, err) - os.WriteFile(fn, []byte("{}"), 0644) - defer func() { - require.NoError(t, os.Remove(fn)) - }() - - a := NewAddressBook(fn) - addrs, err := a.Addresses() - require.NoError(t, err) - require.Equal(t, 0, len(addrs)) - - err = a.Save(1, "0x1") - addrs, err = a.Addresses() - require.NoError(t, err) - assert.DeepEqual(t, map[uint64]map[string]struct{}{1: {"0x1": {}}}, addrs) - - err = a.Save(1, "0x2") - addrs, err = a.Addresses() - require.NoError(t, err) - assert.DeepEqual(t, map[uint64]map[string]struct{}{ - 1: {"0x1": {}, "0x2": {}}, - }, addrs) - - // TODO: Maybe prevent chain collisions? - err = a.Save(2, "0x2") - addrs, err = a.Addresses() - require.NoError(t, err) - assert.DeepEqual(t, map[uint64]map[string]struct{}{ - 1: {"0x1": {}, "0x2": {}}, - 2: {"0x2": {}}, - }, addrs) -}