fix de

Author: nkitlabs

cylinder/store/store.go

@@ -184,6 +184,11 @@ func (s *Store) GetDE(pubDE types.DE) (DE, error) {
 	return de, err
 }
 
+func (s *Store) HasDE(pubDE types.DE) bool {
+	bytes, err := s.DB.Get(DEStoreKey(pubDE))
+	return err == nil && bytes != nil
+}
+
 // DeleteDE deletes the private (d, E) by the given public (D, E)
 func (s *Store) DeleteDE(pubDE types.DE) error {
 	return s.DB.DeleteSync(DEStoreKey(pubDE))

cylinder/workers/de/de.go

@@ -121,7 +121,11 @@ func (de *DE) updateDE() {
 	de.logger.Info(":delivery_truck: Updating DE")
 
 	// Generate new DE pairs
-	privDEs, err := GenerateDEs(de.context.Config.MinDE, de.context.Config.RandomSecret)
+	privDEs, err := GenerateDEs(
+		de.context.Config.MinDE,
+		de.context.Config.RandomSecret,
+		de.context.Store,
+	)
 	if err != nil {
 		de.logger.Error(":cold_sweat: Failed to generate new DE pairs: %s", err)
 		return

cylinder/workers/de/helpers.go

@@ -1,14 +1,27 @@
 package de
 
 import (
+	"errors"
+
 	"github.com/bandprotocol/chain/v2/cylinder/store"
 	"github.com/bandprotocol/chain/v2/pkg/tss"
 	"github.com/bandprotocol/chain/v2/x/tss/types"
 )
 
+const (
+	MaxDuplicateDEAttempts = 5
+)
+
+type DEGetter interface {
+	HasDE(de types.DE) bool
+}
+
 // GenerateDEs generates n pairs of DE by using secret value as a random factor
-func GenerateDEs(n uint64, secret tss.Scalar) (privDEs []store.DE, err error) {
-	for i := uint64(1); i <= n; i++ {
+func GenerateDEs(n uint64, secret tss.Scalar, db DEGetter) (privDEs []store.DE, err error) {
+	privDEs = make([]store.DE, 0, n)
+	attempt := 0
+
+	for len(privDEs) < int(n) && attempt < MaxDuplicateDEAttempts {
 		privD, err := tss.GenerateSigningNonce(secret)
 		if err != nil {
 			return nil, err
@@ -19,15 +32,27 @@ func GenerateDEs(n uint64, secret tss.Scalar) (privDEs []store.DE, err error) {
 			return nil, err
 		}
 
+		pubDE := types.DE{
+			PubD: privD.Point(),
+			PubE: privE.Point(),
+		}
+
+		if db.HasDE(pubDE) {
+			attempt += 1
+			continue
+		}
+
 		privDEs = append(privDEs, store.DE{
-			PubDE: types.DE{
-				PubD: privD.Point(),
-				PubE: privE.Point(),
-			},
+			PubDE: pubDE,
 			PrivD: privD,
 			PrivE: privE,
 		})
 	}
 
+	// Unlikely to occur, but included for fail-safe measures.
+	if attempt >= MaxDuplicateDEAttempts {
+		return nil, errors.New("reach maximum generating duplicated DE")
+	}
+
 	return privDEs, nil
 }

cylinder/workers/de/helpers_test.go

@@ -6,10 +6,15 @@ import (
 	"github.com/stretchr/testify/assert"
 
 	"github.com/bandprotocol/chain/v2/cylinder/workers/de"
+	"github.com/bandprotocol/chain/v2/x/tss/types"
 )
 
+type MockDEGetter struct{}
+
+func (m *MockDEGetter) HasDE(de types.DE) bool { return false }
+
 func TestGenerateDEs(t *testing.T) {
-	privDEs, err := de.GenerateDEs(10, []byte("secret"))
+	privDEs, err := de.GenerateDEs(10, []byte("secret"), &MockDEGetter{})
 	assert.NoError(t, err)
 
 	for _, privDE := range privDEs {

proto/tss/v1beta1/tss.proto

@@ -59,16 +59,6 @@ message DE {
   bytes pub_e = 2 [(gogoproto.casttype) = "github.com/bandprotocol/chain/v2/pkg/tss.Point"];
 }
 
-// DEQueue is a simple queue data structure for holding DE objects.
-message DEQueue {
-  // address is the address of the de holder.
-  string address = 1 [(cosmos_proto.scalar) = "cosmos.AddressString"];
-  // head is the index of the first element in the queue.
-  uint64 head = 2;
-  // tail is the index of the last element in the queue.
-  uint64 tail = 3;
-}
-
 // Signing contains all necessary information for handling a signing request.
 message Signing {
   // id is the unique identifier of the signing.

x/bandtss/keeper/keeper_test.go

@@ -155,14 +155,19 @@ func (s *KeeperTestSuite) setupDE() {
 	for _, tc := range testutil.TestCases {
 		for _, m := range tc.Group.Members {
 			// Submit DEs for each member
+			des := make([]tsstypes.DE, 0, 5)
+			for i := 0; i < 5; i++ {
+				privD, err := tss.GenerateSigningNonce(m.PrivKey)
+				s.Require().NoError(err)
+
+				privE, err := tss.GenerateSigningNonce(m.PrivKey)
+				s.Require().NoError(err)
+
+				des = append(des, tsstypes.DE{PubD: privD.Point(), PubE: privE.Point()})
+			}
+
 			_, err := tssMsgSrvr.SubmitDEs(ctx, &tsstypes.MsgSubmitDEs{
-				DEs: []tsstypes.DE{
-					{PubD: PubD, PubE: PubE},
-					{PubD: PubD, PubE: PubE},
-					{PubD: PubD, PubE: PubE},
-					{PubD: PubD, PubE: PubE},
-					{PubD: PubD, PubE: PubE},
-				},
+				DEs:     des,
 				Address: sdk.AccAddress(m.PubKey()).String(),
 			})
 			s.Require().NoError(err)

x/tss/README.md

@@ -100,7 +100,7 @@ var (
 	ConfirmComplainCountStoreKeyPrefix = []byte{0x0a}
 	ConfirmStoreKeyPrefix = []byte{0x0b}
 	DEStoreKeyPrefix = []byte{0x0c}
-	DEQueueStoreKeyPrefix = []byte{0x0d}
+	DECountStoreKeyPrefix = []byte{0x0d}
 	SigningStoreKeyPrefix = []byte{0x0e}
 	PartialSignatureCountStoreKeyPrefix = []byte{0x0f}
 	PartialSignatureStoreKeyPrefix = []byte{0x10}

x/tss/genesis.go

@@ -37,15 +37,11 @@ func InitGenesis(ctx sdk.Context, k *keeper.Keeper, data *types.GenesisState) {
 			panic(fmt.Sprintf("DEsGenesis of %s size exceeds MaxDESize", addr))
 		}
 
-		k.SetDEQueue(ctx, types.DEQueue{
-			Address: addr,
-			Head:    0,
-			Tail:    uint64(len(des)),
-		})
-
 		acc := sdk.MustAccAddressFromBech32(addr)
-		for i, de := range des {
-			k.SetDE(ctx, acc, uint64(i), de)
+		k.SetDECount(ctx, acc, uint64(len(des)))
+
+		for _, de := range des {
+			k.SetDE(ctx, acc, de)
 		}
 	}
 }

x/tss/genesis_test.go

@@ -89,14 +89,6 @@ func TestExportGenesis(t *testing.T) {
 	require.Equal(t, data.DEsGenesis, exported.DEsGenesis)
 
 	require.Equal(t, uint64(2), k.GetDECount(ctx, addr1))
-	de, err := k.GetDE(ctx, addr1, 1)
-	require.NoError(t, err)
-	require.Equal(
-		t,
-		types.DE{
-			PubD: []byte("pubD2"),
-			PubE: []byte("pubE2"),
-		},
-		de,
-	)
+	hasDE := k.HasDE(ctx, addr1, types.DE{PubD: []byte("pubD2"), PubE: []byte("pubE2")})
+	require.True(t, hasDE)
 }

x/tss/keeper/de.go

@@ -1,87 +1,54 @@
 package keeper
 
 import (
-	"fmt"
-
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 
 	"github.com/bandprotocol/chain/v2/x/tss/types"
 )
 
-// SetDEQueue sets the DEQueue for a given address in the context's KVStore.
-func (k Keeper) SetDEQueue(ctx sdk.Context, deQueue types.DEQueue) {
-	address := sdk.MustAccAddressFromBech32(deQueue.Address)
-	ctx.KVStore(k.storeKey).Set(types.DEQueueKeyStoreKey(address), k.cdc.MustMarshal(&deQueue))
+// SetDECount sets the number of existing DE for a given address.
+func (k Keeper) SetDECount(ctx sdk.Context, address sdk.AccAddress, count uint64) {
+	ctx.KVStore(k.storeKey).Set(types.DECountStoreKey(address), sdk.Uint64ToBigEndian(count))
 }
 
-// GetDEQueue retrieves the DEQueue for a given address from the context's KVStore.
-func (k Keeper) GetDEQueue(ctx sdk.Context, address sdk.AccAddress) types.DEQueue {
-	bz := ctx.KVStore(k.storeKey).Get(types.DEQueueKeyStoreKey(address))
-	if bz == nil {
-		return types.DEQueue{
-			Address: address.String(),
-			Head:    0,
-			Tail:    0,
-		}
-	}
-	var deQueue types.DEQueue
-	k.cdc.MustUnmarshal(bz, &deQueue)
-	return deQueue
+// GetDECount retrieves the number of existing DE for a given address.
+func (k Keeper) GetDECount(ctx sdk.Context, address sdk.AccAddress) uint64 {
+	return sdk.BigEndianToUint64(ctx.KVStore(k.storeKey).Get(types.DECountStoreKey(address)))
 }
 
-// GetDEQueueIterator function gets an iterator over all de queue from the context's KVStore.
-func (k Keeper) GetDEQueueIterator(ctx sdk.Context) sdk.Iterator {
-	return sdk.KVStorePrefixIterator(ctx.KVStore(k.storeKey), types.DEQueueStoreKeyPrefix)
+// GetDEByAddressIterator function gets an iterator over the DEs of the given address.
+func (k Keeper) GetDEByAddressIterator(ctx sdk.Context, addr sdk.AccAddress) sdk.Iterator {
+	return sdk.KVStorePrefixIterator(ctx.KVStore(k.storeKey), types.DEStoreKeyPerAddressPrefix(addr))
 }
 
-// GetDEQueues retrieves all DEQueues from the context's KVStore.
-func (k Keeper) GetDEQueues(ctx sdk.Context) []types.DEQueue {
-	var deQueues []types.DEQueue
-	iterator := k.GetDEQueueIterator(ctx)
-	defer iterator.Close()
-	for ; iterator.Valid(); iterator.Next() {
-		var deQueue types.DEQueue
-		k.cdc.MustUnmarshal(iterator.Value(), &deQueue)
-		deQueues = append(deQueues, deQueue)
-	}
-	return deQueues
-}
-
-// GetDECount retrieves the current count of DE for a given address from the context's KVStore.
-func (k Keeper) GetDECount(ctx sdk.Context, address sdk.AccAddress) uint64 {
-	deQueue := k.GetDEQueue(ctx, address)
-	if deQueue.Head <= deQueue.Tail {
-		return deQueue.Tail - deQueue.Head
-	}
-
-	return k.GetParams(ctx).MaxDESize - (deQueue.Head - deQueue.Tail)
+// HasDE checks if a DE object exists in the context's KVStore.
+func (k Keeper) HasDE(ctx sdk.Context, address sdk.AccAddress, de types.DE) bool {
+	return ctx.KVStore(k.storeKey).Has(types.DEStoreKey(address, de))
 }
 
-// SetDE sets a DE object in the context's KVStore for a given address and index.
-func (k Keeper) SetDE(ctx sdk.Context, address sdk.AccAddress, index uint64, de types.DE) {
-	ctx.KVStore(k.storeKey).Set(types.DEIndexStoreKey(address, index), k.cdc.MustMarshal(&de))
+// SetDE sets a DE object in the context's KVStore for a given address.
+func (k Keeper) SetDE(ctx sdk.Context, address sdk.AccAddress, de types.DE) {
+	ctx.KVStore(k.storeKey).Set(types.DEStoreKey(address, de), []byte{1})
 }
 
-// GetDE retrieves a DE object from the context's KVStore for a given address and index.
+// GetFirstDE retrieves a DE object from the context's KVStore for a given address.
 // Returns an error if DE is not found.
-func (k Keeper) GetDE(ctx sdk.Context, address sdk.AccAddress, index uint64) (types.DE, error) {
-	bz := ctx.KVStore(k.storeKey).Get(types.DEIndexStoreKey(address, index))
-	if bz == nil {
-		return types.DE{}, types.ErrDENotFound.Wrapf(
-			"failed to get DE with address %s index %d",
-			address,
-			index,
-		)
+func (k Keeper) GetFirstDE(ctx sdk.Context, address sdk.AccAddress) (types.DE, error) {
+	iterator := k.GetDEByAddressIterator(ctx, address)
+	defer iterator.Close()
+
+	if !iterator.Valid() {
+		return types.DE{}, types.ErrDENotFound.Wrapf("failed to get DE with address %s", address)
 	}
-	var de types.DE
-	k.cdc.MustUnmarshal(bz, &de)
+
+	_, de := types.ExtractValueFromDEStoreKey(iterator.Key())
 	return de, nil
 }
 
 // DeleteDE removes a DE object from the context's KVStore for a given address and index.
-func (k Keeper) DeleteDE(ctx sdk.Context, address sdk.AccAddress, index uint64) {
-	ctx.KVStore(k.storeKey).Delete(types.DEIndexStoreKey(address, index))
+func (k Keeper) DeleteDE(ctx sdk.Context, address sdk.AccAddress, de types.DE) {
+	ctx.KVStore(k.storeKey).Delete(types.DEStoreKey(address, de))
 }
 
 // GetDEIterator function gets an iterator over all de from the context's KVStore
@@ -94,10 +61,9 @@ func (k Keeper) GetDEsGenesis(ctx sdk.Context) []types.DEGenesis {
 	var des []types.DEGenesis
 	iterator := k.GetDEIterator(ctx)
 	defer iterator.Close()
+
 	for ; iterator.Valid(); iterator.Next() {
-		var de types.DE
-		k.cdc.MustUnmarshal(iterator.Value(), &de)
-		address, _ := types.AddressAndIndexFromDEStoreKey(iterator.Key())
+		address, de := types.ExtractValueFromDEStoreKey(iterator.Key())
 		des = append(des, types.DEGenesis{
 			Address: address.String(),
 			DE:      de,
@@ -106,45 +72,38 @@ func (k Keeper) GetDEsGenesis(ctx sdk.Context) []types.DEGenesis {
 	return des
 }
 
-// NextQueueValue returns next value of head/tail for DE queue
-func (k Keeper) NextQueueValue(ctx sdk.Context, val uint64) uint64 {
-	nextVal := (val + 1) % k.GetParams(ctx).MaxDESize
-	return nextVal
-}
-
 // HandleSetDEs sets multiple DE objects for a given address in the context's KVStore,
-// if tail reaches to head, return err as DE is full
+// if the given DEs reach maximum limit, return err as DE will be over the limit.
 func (k Keeper) HandleSetDEs(ctx sdk.Context, address sdk.AccAddress, des []types.DE) error {
-	deQueue := k.GetDEQueue(ctx, address)
-
+	added := uint64(0)
 	for _, de := range des {
-		k.SetDE(ctx, address, deQueue.Tail, de)
-		deQueue.Tail = k.NextQueueValue(ctx, deQueue.Tail)
-
-		if deQueue.Tail == deQueue.Head {
-			return types.ErrDEQueueFull.Wrap(fmt.Sprintf("DE size exceeds %d", k.GetParams(ctx).MaxDESize))
+		if k.HasDE(ctx, address, de) {
+			continue
 		}
-	}
 
-	k.SetDEQueue(ctx, deQueue)
+		k.SetDE(ctx, address, de)
+		added++
+	}
 
+	cnt := k.GetDECount(ctx, address)
+	if cnt+added > k.GetParams(ctx).MaxDESize {
+		return types.ErrDEReachMaximumLimit.Wrapf("DE size exceeds %d", k.GetParams(ctx).MaxDESize)
+	}
+	k.SetDECount(ctx, address, cnt+uint64(len(des)))
 	return nil
 }
 
-// PollDE retrieves and removes the DE object at the head of the DEQueue for a given address,
-// then increments the head index in the DEQueue.
-// Returns an error if the DE object could not be retrieved.
+// PollDE retrieves and removes the first DE object being retrieved from the iterator of
+// the given address. Returns an error if the DE object could not be retrieved.
 func (k Keeper) PollDE(ctx sdk.Context, address sdk.AccAddress) (types.DE, error) {
-	deQueue := k.GetDEQueue(ctx, address)
-	de, err := k.GetDE(ctx, address, deQueue.Head)
+	de, err := k.GetFirstDE(ctx, address)
 	if err != nil {
 		return types.DE{}, err
 	}
 
-	k.DeleteDE(ctx, address, deQueue.Head)
-
-	deQueue.Head = k.NextQueueValue(ctx, deQueue.Head)
-	k.SetDEQueue(ctx, deQueue)
+	cnt := k.GetDECount(ctx, address)
+	k.SetDECount(ctx, address, cnt-1)
+	k.DeleteDE(ctx, address, de)
 
 	return de, nil
 }