add RectifyChainOwner to ArbOwner precompile

arbos/addressSet/addressSet.go

@@ -4,6 +4,8 @@
 package addressSet
 
 import (
+	"errors"
+
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/offchainlabs/nitro/arbos/storage"
 	"github.com/offchainlabs/nitro/arbos/util"
@@ -84,6 +86,76 @@ func (aset *AddressSet) AllMembers(maxNumToReturn uint64) ([]common.Address, err
 	return ret, nil
 }
 
+func (aset *AddressSet) getMapping(addrHash common.Hash) (common.Hash, uint64, bool, error) {
+	index, err := aset.byAddress.GetUint64(addrHash)
+	if err != nil || index == 0 {
+		return common.Hash{}, index, false, errors.New("RectifyMapping: Address is not an owner")
+	}
+	atIndex, err := aset.backingStorage.GetByUint64(index)
+	if (err != nil || atIndex == common.Hash{}) {
+		return atIndex, index, true, errors.New("RectifyMapping: Invalid mapping")
+	}
+	return atIndex, index, true, nil
+}
+
+func (aset *AddressSet) syncListAndMap(addr common.Address) error {
+	// Iterate through the list and replace the value at first index with incorrect mapping occurance.
+	addrHash := common.BytesToHash(addr.Bytes())
+	size, err := aset.size.Get()
+	if err != nil || size == 0 {
+		return err
+	}
+	for i := uint64(1); i <= size; i++ {
+		tmpAddrHash, err := aset.backingStorage.GetByUint64(i)
+		if err != nil {
+			return err
+		}
+		tmpAddrHashInList, tmpIndex, _, err := aset.getMapping(tmpAddrHash)
+		if err != nil || tmpAddrHash != tmpAddrHashInList || tmpIndex != i {
+			err = aset.backingStorage.SetByUint64(i, addrHash)
+			if err != nil {
+				return err
+			}
+			err = aset.byAddress.Set(addrHash, util.UintToHash(i))
+			return err
+		}
+	}
+	// List is correctly aligned with map, only way to sync owner is to add a new entry to list.
+	err = aset.byAddress.Clear(addrHash)
+	if err != nil {
+		return err
+	}
+	err = aset.Add(addr)
+	return err
+}
+
+func (aset *AddressSet) RectifyMapping(addr common.Address) error {
+	addrHash := common.BytesToHash(addr.Bytes())
+	addrHashInList, index, inMap, err := aset.getMapping(addrHash)
+	// Key is not found in the Map.
+	if !inMap {
+		return err
+	}
+	// Map has incorrect list index for the address.
+	if err != nil {
+		err = aset.syncListAndMap(addr)
+		return err
+	}
+	// Map and list are correctly synced for this address.
+	if addrHash == addrHashInList {
+		return nil
+	}
+	// Not a correct mapping Or is a correct mapping and no collision, list value at 'index' can be edited.
+	tmpAddrHashInList, tmpIndex, _, err := aset.getMapping(addrHashInList)
+	if err != nil || addrHashInList != tmpAddrHashInList || index != tmpIndex {
+		err = aset.backingStorage.SetByUint64(index, addrHash)
+		return err
+	}
+	// Both keys addrHash and addrHashInList point to the same index and addrHashInList is correctly synced.
+	err = aset.syncListAndMap(addr)
+	return err
+}
+
 func (aset *AddressSet) Add(addr common.Address) error {
 	present, err := aset.IsMember(addr)
 	if present || err != nil {

arbos/addressSet/addressSet_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/offchainlabs/nitro/arbos/burn"
 	"github.com/offchainlabs/nitro/arbos/storage"
+	"github.com/offchainlabs/nitro/arbos/util"
 	"github.com/offchainlabs/nitro/util/colors"
 	"github.com/offchainlabs/nitro/util/testhelpers"
 )
@@ -173,6 +174,65 @@ func TestAddressSetAllMembers(t *testing.T) {
 	}
 }
 
+func TestRectifyMapping(t *testing.T) {
+	db := storage.NewMemoryBackedStateDB()
+	sto := storage.NewGeth(db, burn.NewSystemBurner(nil, false))
+	Require(t, Initialize(sto))
+	aset := OpenAddressSet(sto)
+
+	addr1 := testhelpers.RandomAddress()
+	addr2 := testhelpers.RandomAddress()
+	addr3 := testhelpers.RandomAddress()
+	possibleAddresses := []common.Address{addr1, addr2, addr3}
+
+	Require(t, aset.Add(addr1))
+	Require(t, aset.Add(addr2))
+	Require(t, aset.Add(addr3))
+
+	// RectifyMapping should not do anything if the mapping is correct
+	addr := possibleAddresses[rand.Intn(len(possibleAddresses))]
+	Require(t, aset.RectifyMapping(addr))
+
+	// Non owner's should not be able to call RectifyMapping
+	err := aset.RectifyMapping(testhelpers.RandomAddress())
+	if err == nil {
+		Fail(t, "RectifyMapping was succesfully called by non owner")
+	}
+
+	// Corrupt the list and verify if RectifyMapping fixes it
+	addrHash := common.BytesToHash(addr2.Bytes())
+	Require(t, aset.backingStorage.SetByUint64(uint64(1), addrHash))
+
+	Require(t, aset.RectifyMapping(addr1))
+	addrHash = common.BytesToHash(addr1.Bytes())
+	addrHashInList, index, _, _ := aset.getMapping(addrHash)
+	if addrHashInList != addrHash || index != uint64(1) {
+		Fail(t, "RectifyMapping did not rectify the corrupt list")
+	}
+
+	// Corrupt the map and verify if RectifyMapping fixes it
+	addrHash = common.BytesToHash(addr2.Bytes())
+	Require(t, aset.byAddress.Set(addrHash, util.UintToHash(uint64(6))))
+
+	Require(t, aset.RectifyMapping(addr2))
+	addrHashInList, index, _, _ = aset.getMapping(addrHash)
+	if addrHashInList != addrHash || index != uint64(2) {
+		Fail(t, "RectifyMapping did not rectify the corrupt map")
+	}
+
+	// Add a new owner to the map and verify if RectifyMapping syncs list with the map
+	// to check for the case where list has fewer owners than expected
+	addr4 := testhelpers.RandomAddress()
+	addrHash = common.BytesToHash(addr4.Bytes())
+	Require(t, aset.byAddress.Set(addrHash, util.UintToHash(uint64(1))))
+
+	Require(t, aset.RectifyMapping(addr4))
+	addrHashInList, _, _, _ = aset.getMapping(addrHash)
+	if addrHashInList != addrHash || size(t, aset) != 4 {
+		Fail(t, "RectifyMapping did not add the missing owner to the list")
+	}
+}
+
 func checkAllMembers(t *testing.T, aset *AddressSet, possibleAddresses []common.Address) {
 	allMembers, err := aset.AllMembers(1024)
 	Require(t, err)

precompiles/ArbOwner.go

@@ -44,6 +44,11 @@ func (con ArbOwner) RemoveChainOwner(c ctx, evm mech, addr addr) error {
 	return c.State.ChainOwners().Remove(addr, c.State.ArbOSVersion())
 }
 
+// RectifyChainOwner checks if the account is a chain owner
+func (con ArbOwner) RectifyChainOwner(c ctx, evm mech, addr addr) error {
+	return c.State.ChainOwners().RectifyMapping(addr)
+}
+
 // IsChainOwner checks if the account is a chain owner
 func (con ArbOwner) IsChainOwner(c ctx, evm mech, addr addr) (bool, error) {
 	return c.State.ChainOwners().IsMember(addr)