Merge pull request #483 from iotaledger/feat/safemath-everywhere

Use safemath in accounts

components/debugapi/node.go

@@ -17,7 +17,12 @@ func validatorsSummary() (*ValidatorsSummaryResponse, error) {
 	}
 
 	var validatorSeats []*Validator
-	latestCommittee.Accounts().ForEach(func(id iotago.AccountID, pool *account.Pool) bool {
+	accounts, err := latestCommittee.Accounts()
+	if err != nil {
+		return nil, ierrors.Wrapf(err, "failed to get accounts from committee for slot %d", latestSlotIndex)
+	}
+
+	accounts.ForEach(func(id iotago.AccountID, pool *account.Pool) bool {
 		validatorSeats = append(validatorSeats, &Validator{
 			AccountID:      id,
 			SeatIndex:      uint8(lo.Return1(latestCommittee.GetSeat(id))),

components/restapi/core/accounts.go

@@ -188,7 +188,7 @@ func rewardsByOutputID(c echo.Context) (*apimodels.ManaRewardsResponse, error) {
 	}, nil
 }
 
-func selectedCommittee(c echo.Context) *apimodels.CommitteeResponse {
+func selectedCommittee(c echo.Context) (*apimodels.CommitteeResponse, error) {
 	timeProvider := deps.Protocol.CommittedAPI().TimeProvider()
 
 	var slot iotago.SlotIndex
@@ -206,11 +206,16 @@ func selectedCommittee(c echo.Context) *apimodels.CommitteeResponse {
 	if !exists {
 		return &apimodels.CommitteeResponse{
 			Epoch: epoch,
-		}
+		}, nil
+	}
+
+	accounts, err := seatedAccounts.Accounts()
+	if err != nil {
+		return nil, ierrors.Wrapf(err, "failed to get accounts from committee for slot %d", slot)
 	}
 
-	committee := make([]*apimodels.CommitteeMemberResponse, 0, seatedAccounts.Accounts().Size())
-	seatedAccounts.Accounts().ForEach(func(accountID iotago.AccountID, seat *account.Pool) bool {
+	committee := make([]*apimodels.CommitteeMemberResponse, 0, accounts.Size())
+	accounts.ForEach(func(accountID iotago.AccountID, seat *account.Pool) bool {
 		committee = append(committee, &apimodels.CommitteeMemberResponse{
 			AccountID:      accountID,
 			PoolStake:      seat.PoolStake,
@@ -224,7 +229,7 @@ func selectedCommittee(c echo.Context) *apimodels.CommitteeResponse {
 	return &apimodels.CommitteeResponse{
 		Epoch:               epoch,
 		Committee:           committee,
-		TotalStake:          seatedAccounts.Accounts().TotalStake(),
-		TotalValidatorStake: seatedAccounts.Accounts().TotalValidatorStake(),
-	}
+		TotalStake:          accounts.TotalStake(),
+		TotalValidatorStake: accounts.TotalValidatorStake(),
+	}, nil
 }

components/restapi/core/component.go

@@ -357,7 +357,10 @@ func configure() error {
 	}, checkNodeSynced())
 
 	routeGroup.GET(RouteCommittee, func(c echo.Context) error {
-		resp := selectedCommittee(c)
+		resp, err := selectedCommittee(c)
+		if err != nil {
+			return err
+		}
 
 		return responseByHeader(c, resp)
 	}, checkNodeSynced())

pkg/core/account/accounts.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"sync/atomic"
 
+	"github.com/iotaledger/hive.go/core/safemath"
 	"github.com/iotaledger/hive.go/ds/shrinkingmap"
 	"github.com/iotaledger/hive.go/ierrors"
 	"github.com/iotaledger/hive.go/runtime/syncutils"
@@ -71,31 +72,44 @@ func (a *Accounts) Get(id iotago.AccountID) (pool *Pool, exists bool) {
 }
 
 // setWithoutLocking sets the weight of the given identity.
-func (a *Accounts) setWithoutLocking(id iotago.AccountID, pool *Pool) {
+func (a *Accounts) setWithoutLocking(id iotago.AccountID, pool *Pool) error {
 	value, created := a.accountPools.GetOrCreate(id, func() *Pool {
 		return pool
 	})
 
+	var safeMathErr error
+
 	if !created {
 		// if there was already an entry, we need to subtract the former
 		// stake first and set the new value
-		// TODO: use safemath
-		a.totalStake -= value.PoolStake
-		a.totalValidatorStake -= value.ValidatorStake
+		if a.totalStake, safeMathErr = safemath.SafeSub(a.totalStake, value.PoolStake); safeMathErr != nil {
+			return ierrors.Wrapf(safeMathErr, "failed to subtract pool stake from total stake for account %s", id.String())
+		}
+
+		if a.totalValidatorStake, safeMathErr = safemath.SafeSub(a.totalValidatorStake, value.ValidatorStake); safeMathErr != nil {
+			return ierrors.Wrapf(safeMathErr, "failed to subtract validator stake from total validator stake for account %s", id.String())
+		}
 
 		a.accountPools.Set(id, pool)
 	}
 
-	a.totalStake += pool.PoolStake
-	a.totalValidatorStake += pool.ValidatorStake
+	if a.totalStake, safeMathErr = safemath.SafeAdd(a.totalStake, pool.PoolStake); safeMathErr != nil {
+		return ierrors.Wrapf(safeMathErr, "failed to add pool stake to total stake for account %s", id.String())
+	}
+
+	if a.totalValidatorStake, safeMathErr = safemath.SafeAdd(a.totalValidatorStake, pool.ValidatorStake); safeMathErr != nil {
+		return ierrors.Wrapf(safeMathErr, "failed to add validator stake to total validator stake for account %s", id.String())
+	}
+
+	return nil
 }
 
 // Set sets the weight of the given identity.
-func (a *Accounts) Set(id iotago.AccountID, pool *Pool) {
+func (a *Accounts) Set(id iotago.AccountID, pool *Pool) error {
 	a.mutex.Lock()
 	defer a.mutex.Unlock()
 
-	a.setWithoutLocking(id, pool)
+	return a.setWithoutLocking(id, pool)
 }
 
 func (a *Accounts) TotalStake() iotago.BaseToken {
@@ -168,7 +182,9 @@ func (a *Accounts) readFromReadSeeker(reader io.ReadSeeker) (n int, err error) {
 			return 0, ierrors.Wrap(err, "invalid pool bytes length")
 		}
 
-		a.setWithoutLocking(accountID, pool)
+		if err := a.setWithoutLocking(accountID, pool); err != nil {
+			return 0, ierrors.Wrapf(err, "failed to set pool for account %s", accountID.String())
+		}
 	}
 
 	var reused bool

pkg/core/account/accounts_test.go

@@ -20,11 +20,13 @@ func TestAccounts(t *testing.T) {
 
 	// check "Set"
 	for id, stake := range issuers {
-		accounts.Set(id, &account.Pool{
+		if err := accounts.Set(id, &account.Pool{
 			PoolStake:      iotago.BaseToken(stake),
 			ValidatorStake: iotago.BaseToken(stake) * 2,
 			FixedCost:      iotago.Mana(stake) * 3,
-		})
+		}); err != nil {
+			t.Fatal(err)
+		}
 	}
 
 	// check "Size"

pkg/core/account/seated_accounts.go

@@ -78,19 +78,22 @@ func (s *SeatedAccounts) SeatCount() int {
 	return s.seatsByAccount.Size()
 }
 
-func (s *SeatedAccounts) Accounts() *Accounts {
+func (s *SeatedAccounts) Accounts() (*Accounts, error) {
 	accounts := NewAccounts()
+	var err error
 	s.seatsByAccount.ForEachKey(func(id iotago.AccountID) bool {
 		pool, exists := s.accounts.Get(id)
 		if !exists {
 			panic("account not found")
 		}
-		accounts.Set(id, pool)
+		if err = accounts.Set(id, pool); err != nil {
+			return false
+		}
 
 		return true
 	})
 
-	return accounts
+	return accounts, err
 }
 
 func (s *SeatedAccounts) String() string {

pkg/core/account/seated_accounts_test.go

@@ -19,9 +19,9 @@ func TestSelectedAccounts(t *testing.T) {
 	account3 := iotago.AccountID([32]byte{3})
 	account4 := iotago.AccountID([32]byte{4})
 
-	accounts.Set(account1, &account.Pool{})
-	accounts.Set(account2, &account.Pool{})
-	accounts.Set(account3, &account.Pool{})
+	require.NoError(t, accounts.Set(account1, &account.Pool{}))
+	require.NoError(t, accounts.Set(account2, &account.Pool{}))
+	require.NoError(t, accounts.Set(account3, &account.Pool{}))
 
 	// Create a new set of selected accounts
 	seatedAccounts := account.NewSeatedAccounts(accounts, account1, account3)
@@ -65,7 +65,8 @@ func TestSelectedAccounts(t *testing.T) {
 	require.True(t, has)
 
 	// Test the "Members" method
-	members := seatedAccounts.Accounts()
+	members, err := seatedAccounts.Accounts()
+	require.NoError(t, err)
 	require.Equal(t, 2, members.Size())
 	require.True(t, members.Has(account2))
 	require.True(t, members.Has(account3))

pkg/protocol/commitment_verifier.go

@@ -25,10 +25,15 @@ func NewCommitmentVerifier(mainEngine *engine.Engine, lastCommonCommitmentBefore
 		return nil, ierrors.Errorf("committee in slot %d does not exist", lastCommonCommitmentBeforeFork.Slot())
 	}
 
+	accountsAtForkingPoint, err := committeeAtForkingPoint.Accounts()
+	if err != nil {
+		return nil, ierrors.Wrapf(err, "failed to get accounts from committee for slot %d", lastCommonCommitmentBeforeFork.Slot())
+	}
+
 	return &CommitmentVerifier{
 		engine:                  mainEngine,
 		cumulativeWeight:        lastCommonCommitmentBeforeFork.CumulativeWeight(),
-		validatorAccountsAtFork: lo.PanicOnErr(mainEngine.Ledger.PastAccounts(committeeAtForkingPoint.Accounts().IDs(), lastCommonCommitmentBeforeFork.Slot())),
+		validatorAccountsAtFork: lo.PanicOnErr(mainEngine.Ledger.PastAccounts(accountsAtForkingPoint.IDs(), lastCommonCommitmentBeforeFork.Slot())),
 		// TODO: what happens if the committee rotated after the fork?
 	}, nil
 }

pkg/protocol/engine/attestation/slotattestation/testframework_test.go

@@ -61,7 +61,9 @@ func NewTestFramework(test *testing.T) *TestFramework {
 		accounts := account.NewAccounts()
 		var members []iotago.AccountID
 		t.issuerByAlias.ForEach(func(alias string, issuer *issuer) bool {
-			accounts.Set(issuer.accountID, &account.Pool{}) // we don't care about pools with PoA
+			if err := accounts.Set(issuer.accountID, &account.Pool{}); err != nil { // we don't care about pools with PoA
+				test.Fatal(err)
+			}
 			members = append(members, issuer.accountID)
 			return true
 		})

pkg/protocol/engine/filter/blockfilter/filter_test.go

@@ -98,7 +98,9 @@ func (t *TestFramework) IssueBlockAtSlotWithVersion(alias string, slot iotago.Sl
 
 func mockedCommitteeFunc(validatorAccountID iotago.AccountID) func(iotago.SlotIndex) (*account.SeatedAccounts, bool) {
 	mockedAccounts := account.NewAccounts()
-	mockedAccounts.Set(validatorAccountID, new(account.Pool))
+	if err := mockedAccounts.Set(validatorAccountID, new(account.Pool)); err != nil {
+		panic(err)
+	}
 	seatedAccounts := account.NewSeatedAccounts(mockedAccounts)
 	seatedAccounts.Set(account.SeatIndex(0), validatorAccountID)
 

pkg/protocol/engine/mempool/conflictdag/tests/accounts_framework.go

@@ -74,7 +74,9 @@ func (f *AccountsTestFramework) CreateID(alias string) iotago.AccountID {
 	validatorID := iotago.AccountIDFromData(hashedAlias[:])
 	validatorID.RegisterAlias(alias)
 
-	f.Instance.Set(validatorID, &account.Pool{}) // we don't care about pools when doing PoA
+	if err := f.Instance.Set(validatorID, &account.Pool{}); err != nil { // we don't care about pools when doing PoA
+		f.test.Fatal(err)
+	}
 	f.Committee.Set(account.SeatIndex(f.Committee.SeatCount()), validatorID)
 
 	f.identitiesByAlias[alias] = validatorID

pkg/protocol/sybilprotection/seatmanager/mock/mockseatmanager.go

@@ -61,11 +61,13 @@ func NewManualPOAProvider() module.Provider[*engine.Engine, seatmanager.SeatMana
 func (m *ManualPOA) AddRandomAccount(alias string) iotago.AccountID {
 	id := iotago.AccountID(tpkg.Rand32ByteArray())
 	id.RegisterAlias(alias)
-	m.accounts.Set(id, &account.Pool{
+	if err := m.accounts.Set(id, &account.Pool{ // We don't care about pools with PoA, but need to set something to avoid division by zero errors.
 		PoolStake:      1,
 		ValidatorStake: 1,
 		FixedCost:      1,
-	}) // We don't care about pools with PoA, but need to set something to avoid division by zero errors.
+	}); err != nil {
+		panic(err)
+	}
 
 	m.aliases.Set(alias, id)
 
@@ -79,11 +81,13 @@ func (m *ManualPOA) AddRandomAccount(alias string) iotago.AccountID {
 }
 
 func (m *ManualPOA) AddAccount(id iotago.AccountID, alias string) iotago.AccountID {
-	m.accounts.Set(id, &account.Pool{
+	if err := m.accounts.Set(id, &account.Pool{ // We don't care about pools with PoA, but need to set something to avoid division by zero errors.
 		PoolStake:      1,
 		ValidatorStake: 1,
 		FixedCost:      1,
-	}) // We don't care about pools with PoA, but need to set something to avoid division by zero errors.
+	}); err != nil {
+		panic(err)
+	}
 	m.aliases.Set(alias, id)
 
 	m.committee = m.accounts.SelectCommittee(m.accounts.IDs()...)
@@ -164,11 +168,13 @@ func (m *ManualPOA) RotateCommittee(epoch iotago.EpochIndex, validators accounts
 		m.accounts = account.NewAccounts()
 
 		for _, validatorData := range validators {
-			m.accounts.Set(validatorData.ID, &account.Pool{
+			if err := m.accounts.Set(validatorData.ID, &account.Pool{
 				PoolStake:      validatorData.ValidatorStake + validatorData.DelegationStake,
 				ValidatorStake: validatorData.ValidatorStake,
 				FixedCost:      validatorData.FixedCost,
-			})
+			}); err != nil {
+				return nil, ierrors.Wrapf(err, "error while setting pool for epoch %d for validator %s", epoch, validatorData.ID.String())
+			}
 		}
 		m.committee = m.accounts.SelectCommittee(m.accounts.IDs()...)
 	}

pkg/protocol/sybilprotection/seatmanager/poa/poa.go

@@ -89,17 +89,23 @@ func (s *SeatManager) RotateCommittee(epoch iotago.EpochIndex, validators accoun
 		committeeAccounts := account.NewAccounts()
 
 		for _, validatorData := range validators {
-			committeeAccounts.Set(validatorData.ID, &account.Pool{
+			if err := committeeAccounts.Set(validatorData.ID, &account.Pool{
 				PoolStake:      validatorData.ValidatorStake + validatorData.DelegationStake,
 				ValidatorStake: validatorData.ValidatorStake,
 				FixedCost:      validatorData.FixedCost,
-			})
+			}); err != nil {
+				return nil, ierrors.Wrapf(err, "error while setting committee for epoch %d for validator %s", epoch, validatorData.ID.String())
+			}
 		}
 		s.committee = committeeAccounts.SelectCommittee(committeeAccounts.IDs()...)
 	}
 
-	err := s.committeeStore.Store(epoch, s.committee.Accounts())
+	accounts, err := s.committee.Accounts()
 	if err != nil {
+		return nil, ierrors.Wrapf(err, "error while getting accounts from committee for epoch %d", epoch)
+	}
+
+	if err := s.committeeStore.Store(epoch, accounts); err != nil {
 		return nil, ierrors.Wrapf(err, "error while storing committee for epoch %d", epoch)
 	}
 
@@ -186,8 +192,12 @@ func (s *SeatManager) SetCommittee(epoch iotago.EpochIndex, validators *account.
 
 	s.committee = validators.SelectCommittee(validators.IDs()...)
 
-	err := s.committeeStore.Store(epoch, s.committee.Accounts())
+	accounts, err := s.committee.Accounts()
 	if err != nil {
+		return ierrors.Wrapf(err, "failed to get accounts from committee for epoch %d", epoch)
+	}
+
+	if err := s.committeeStore.Store(epoch, accounts); err != nil {
 		return ierrors.Wrapf(err, "failed to set committee for epoch %d", epoch)
 	}