refactor: use a map to store credentials instead of an array

.codeclimate.yml

@@ -9,3 +9,6 @@ plugins:
     #enabled: true
 exclude_patterns:
   - "."
+  - "**/*.pb.go"
+  - "**/rln/contracts/*.go"
+  - "**/bindata.go"

cmd/waku/flags_rln.go

@@ -6,6 +6,7 @@ package main
 import (
 	cli "github.com/urfave/cli/v2"
 	wcli "github.com/waku-org/go-waku/waku/cliutils"
+	"github.com/waku-org/go-waku/waku/v2/protocol/rln/keystore"
 )
 
 func rlnFlags() []cli.Flag {
@@ -17,10 +18,10 @@ func rlnFlags() []cli.Flag {
 			Destination: &options.RLNRelay.Enable,
 		},
 		&cli.UintFlag{
-			Name:        "rln-relay-membership-group-index",
+			Name:        "rln-relay-membership-index",
 			Value:       0,
-			Usage:       "the index of credentials to use, within a specific rln membership set",
+			Usage:       "the index of credentials to use",
-			Destination: &options.RLNRelay.MembershipGroupIndex,
+			Destination: &options.RLNRelay.MembershipIndex,
 		},
 		&cli.BoolFlag{
 			Name:        "rln-relay-dynamic",
@@ -30,12 +31,12 @@ func rlnFlags() []cli.Flag {
 		&cli.PathFlag{
 			Name:        "rln-relay-cred-path",
 			Usage:       "RLN relay membership credentials file",
-			Value:       "",
+			Value:       keystore.DefaultCredentialsFilename,
 			Destination: &options.RLNRelay.CredentialsPath,
 		},
 		&cli.StringFlag{
 			Name:        "rln-relay-cred-password",
-			Value:       "",
+			Value:       keystore.DefaultCredentialsPassword,
 			Usage:       "Password for encrypting RLN credentials",
 			Destination: &options.RLNRelay.CredentialsPassword,
 		},
@@ -45,12 +46,6 @@ func rlnFlags() []cli.Flag {
 			Usage:       "Path to the RLN merkle tree sled db (https://github.com/spacejam/sled)",
 			Destination: &options.RLNRelay.TreePath,
 		},
-		&cli.UintFlag{
-			Name:        "rln-relay-membership-index",
-			Value:       0,
-			Usage:       "the index of credentials to use",
-			Destination: &options.RLNRelay.CredentialsIndex,
-		},
 		&cli.StringFlag{
 			Name:        "rln-relay-eth-client-address",
 			Usage:       "Ethereum testnet client address",

cmd/waku/node_rln.go

@@ -17,17 +17,16 @@ func checkForRLN(logger *zap.Logger, options NodeOptions, nodeOpts *[]node.WakuN
 			failOnErr(errors.New("relay not available"), "Could not enable RLN Relay")
 		}
 		if !options.RLNRelay.Dynamic {
-			*nodeOpts = append(*nodeOpts, node.WithStaticRLNRelay(rln.MembershipIndex(options.RLNRelay.MembershipGroupIndex), nil))
+			*nodeOpts = append(*nodeOpts, node.WithStaticRLNRelay(rln.MembershipIndex(options.RLNRelay.MembershipIndex), nil))
 		} else {
 			// TODO: too many parameters in this function
 			// consider passing a config struct instead
 			*nodeOpts = append(*nodeOpts, node.WithDynamicRLNRelay(
 				options.RLNRelay.CredentialsPath,
 				options.RLNRelay.CredentialsPassword,
-				options.RLNRelay.CredentialsIndex,
 				options.RLNRelay.TreePath,
 				options.RLNRelay.MembershipContractAddress,
-				rln.MembershipIndex(options.RLNRelay.MembershipGroupIndex),
+				rln.MembershipIndex(options.RLNRelay.MembershipIndex),
 				nil,
 				options.RLNRelay.ETHClientAddress,
 			))

cmd/waku/options.go

@@ -37,9 +37,8 @@ type RLNRelayOptions struct {
 	Enable                    bool
 	CredentialsPath           string
 	CredentialsPassword       string
-	CredentialsIndex          uint
 	TreePath                  string
-	MembershipGroupIndex      uint
+	MembershipIndex           uint
 	Dynamic                   bool
 	ETHClientAddress          string
 	MembershipContractAddress common.Address

cmd/waku/rlngenerate/command_rln.go

@@ -103,27 +103,24 @@ func execute(ctx context.Context) error {
 	return nil
 }
 
-func persistCredentials(identityCredential *rln.IdentityCredential, membershipIndex rln.MembershipIndex, chainID *big.Int) error {
+func persistCredentials(identityCredential *rln.IdentityCredential, treeIndex rln.MembershipIndex, chainID *big.Int) error {
 	appKeystore, err := keystore.New(options.CredentialsPath, dynamic.RLNAppInfo, logger)
 	if err != nil {
 		return err
 	}
 
-	membershipGroup := keystore.MembershipGroup{
+	membershipCredential := keystore.MembershipCredentials{
-		TreeIndex: membershipIndex,
+		IdentityCredential:     identityCredential,
-		MembershipContract: keystore.MembershipContract{
+		TreeIndex:              treeIndex,
-			ChainID: fmt.Sprintf("0x%X", chainID.Int64()),
+		MembershipContractInfo: keystore.NewMembershipContractInfo(chainID, options.MembershipContractAddress),
-			Address: options.MembershipContractAddress.String(),
-		},
 	}
 
-	membershipGroupIndex, err := appKeystore.AddMembershipCredentials(identityCredential, membershipGroup, options.CredentialsPassword)
+	err = appKeystore.AddMembershipCredentials(membershipCredential, options.CredentialsPassword)
 	if err != nil {
 		return fmt.Errorf("failed to persist credentials: %w", err)
 	}
 
-	// TODO: obtain keystore index?
+	logger.Info("persisted credentials succesfully")
-	logger.Info("persisted credentials succesfully", zap.Uint("membershipGroupIndex", membershipGroupIndex))
 
 	return nil
 }

examples/chat2/chat.go

@@ -468,11 +468,7 @@ func (c *Chat) welcomeMessage() {
 		fmt.Println(err.Error())
 	}
 
-	idx, err := c.node.RLNRelay().MembershipIndex()
+	idx := c.node.RLNRelay().MembershipIndex()
-	if err != nil {
-		c.ui.Quit()
-		fmt.Println(err.Error())
-	}
 
 	idTrapdoor := credential.IDTrapdoor
 	idNullifier := credential.IDSecretHash

examples/chat2/exec.go

@@ -52,7 +52,6 @@ func execute(options Options) {
 			opts = append(opts, node.WithDynamicRLNRelay(
 				options.RLNRelay.CredentialsPath,
 				options.RLNRelay.CredentialsPassword,
-				options.RLNRelay.CredentialsIndex,
 				"", // Will use default tree path
 				options.RLNRelay.MembershipContractAddress,
 				uint(options.RLNRelay.MembershipIndex),

examples/chat2/flags.go

@@ -185,17 +185,11 @@ func getFlags() []cli.Flag {
 			Usage:       "Enable spam protection through rln-relay",
 			Destination: &options.RLNRelay.Enable,
 		},
-		&cli.UintFlag{
-			Name:        "rln-relay-membership-group-index",
-			Value:       0,
-			Usage:       "the index of credentials to use, within a specific rln membership set",
-			Destination: &options.RLNRelay.MembershipGroupIndex,
-		},
 		&cli.UintFlag{
 			Name:        "rln-relay-membership-index",
 			Value:       0,
 			Usage:       "the index of credentials to use",
-			Destination: &options.RLNRelay.CredentialsIndex,
+			Destination: &options.RLNRelay.MembershipIndex,
 		},
 		&cli.BoolFlag{
 			Name:        "rln-relay-dynamic",

examples/chat2/options.go

@@ -31,8 +31,6 @@ type RLNRelayOptions struct {
 	Enable                    bool
 	CredentialsPath           string
 	CredentialsPassword       string
-	CredentialsIndex          uint
-	MembershipGroupIndex      uint
 	MembershipIndex           uint
 	Dynamic                   bool
 	ETHClientAddress          string

examples/rln/main.go

@@ -29,8 +29,7 @@ const ethClientAddress = "wss://sepolia.infura.io/ws/v3/API_KEY_GOES_HERE"
 const contractAddress = "0x9C09146844C1326c2dBC41c451766C7138F88155"
 const keystorePath = ""     // Empty to store in current folder
 const keystorePassword = "" // Empty to use default
-const keystoreIndex = 0
+const membershipIndex = 0
-const membershipGroupIndex = 0
 
 var contentTopic = protocol.NewContentTopic("rln", 1, "test", "proto").String()
 var pubsubTopic = protocol.DefaultPubsubTopic()
@@ -63,14 +62,11 @@ func main() {
 		node.WithNTP(),
 		node.WithWakuRelay(),
 		node.WithDynamicRLNRelay(
-			pubsubTopic.String(),
-			contentTopic,
 			keystorePath,
 			keystorePassword,
-			keystoreIndex,
 			"", // Will use default tree path
 			common.HexToAddress(contractAddress),
-			membershipGroupIndex,
+			membershipIndex,
 			spamHandler,
 			ethClientAddress,
 		),

waku/v2/node/wakunode2.go

@@ -70,7 +70,7 @@ type SpamHandler = func(message *pb.WakuMessage) error
 
 type RLNRelay interface {
 	IdentityCredential() (IdentityCredential, error)
-	MembershipIndex() (uint, error)
+	MembershipIndex() uint
 	AppendRLNProof(msg *pb.WakuMessage, senderEpochTime time.Time) error
 	Validator(spamHandler SpamHandler) func(ctx context.Context, peerID peer.ID, message *pubsub.Message) bool
 	Start(ctx context.Context) error

waku/v2/node/wakunode2_rln.go

@@ -56,7 +56,6 @@ func (w *WakuNode) setupRLNRelay() error {
 			w.opts.rlnRelayMemIndex,
 			appKeystore,
 			w.opts.keystorePassword,
-			w.opts.keystoreIndex,
 			w.opts.prometheusReg,
 			w.log,
 		)

waku/v2/node/wakuoptions.go

@@ -100,7 +100,6 @@ type WakuNodeParameters struct {
 	rlnETHClientAddress          string
 	keystorePath                 string
 	keystorePassword             string
-	keystoreIndex                uint
 	rlnTreePath                  string
 	rlnMembershipContractAddress common.Address
 

waku/v2/node/wakuoptions_rln.go

@@ -23,17 +23,16 @@ func WithStaticRLNRelay(memberIndex r.MembershipIndex, spamHandler rln.SpamHandl
 
 // WithDynamicRLNRelay enables the Waku V2 RLN protocol in onchain mode.
 // Requires the `gowaku_rln` build constrain (or the env variable RLN=true if building go-waku)
-func WithDynamicRLNRelay(keystorePath string, keystorePassword string, keystoreIndex uint, treePath string, membershipContract common.Address, membershipGroupIndex uint, spamHandler rln.SpamHandler, ethClientAddress string) WakuNodeOption {
+func WithDynamicRLNRelay(keystorePath string, keystorePassword string, treePath string, membershipContract common.Address, membershipIndex uint, spamHandler rln.SpamHandler, ethClientAddress string) WakuNodeOption {
 	return func(params *WakuNodeParameters) error {
 		params.enableRLN = true
 		params.rlnRelayDynamic = true
 		params.keystorePassword = keystorePassword
 		params.keystorePath = keystorePath
-		params.keystoreIndex = keystoreIndex
 		params.rlnSpamHandler = spamHandler
 		params.rlnETHClientAddress = ethClientAddress
 		params.rlnMembershipContractAddress = membershipContract
-		params.rlnRelayMemIndex = membershipGroupIndex
+		params.rlnRelayMemIndex = membershipIndex
 		params.rlnTreePath = treePath
 		return nil
 	}

waku/v2/protocol/rln/group_manager/dynamic/dynamic.go

@@ -3,7 +3,6 @@ package dynamic
 import (
 	"context"
 	"errors"
-	"fmt"
 	"math/big"
 	"sync"
 	"time"
@@ -25,7 +24,7 @@ import (
 var RLNAppInfo = keystore.AppInfo{
 	Application:   "waku-rln-relay",
 	AppIdentifier: "01234567890abcdef",
-	Version:       "0.1",
+	Version:       "0.2",
 }
 
 type DynamicGroupManager struct {
@@ -37,10 +36,9 @@ type DynamicGroupManager struct {
 	wg     sync.WaitGroup
 
 	identityCredential *rln.IdentityCredential
-	membershipIndex    *rln.MembershipIndex
+	membershipIndex    rln.MembershipIndex
 
 	membershipContractAddress common.Address
-	membershipGroupIndex      uint
 	ethClientAddress          string
 	ethClient                 *ethclient.Client
 
@@ -53,7 +51,6 @@ type DynamicGroupManager struct {
 
 	appKeystore      *keystore.AppKeystore
 	keystorePassword string
-	keystoreIndex    uint
 
 	rootTracker *group_manager.MerkleRootTracker
 }
@@ -120,23 +117,21 @@ type RegistrationHandler = func(tx *types.Transaction)
 func NewDynamicGroupManager(
 	ethClientAddr string,
 	memContractAddr common.Address,
-	membershipGroupIndex uint,
+	membershipIndex uint,
 	appKeystore *keystore.AppKeystore,
 	keystorePassword string,
-	keystoreIndex uint,
 	reg prometheus.Registerer,
 	log *zap.Logger,
 ) (*DynamicGroupManager, error) {
 	log = log.Named("rln-dynamic")
 
 	return &DynamicGroupManager{
-		membershipGroupIndex:      membershipGroupIndex,
+		membershipIndex:           membershipIndex,
 		membershipContractAddress: memContractAddr,
 		ethClientAddress:          ethClientAddr,
 		eventHandler:              handler,
 		appKeystore:               appKeystore,
 		keystorePassword:          keystorePassword,
-		keystoreIndex:             keystoreIndex,
 		log:                       log,
 		metrics:                   newMetrics(reg),
 	}, nil
@@ -198,30 +193,18 @@ func (gm *DynamicGroupManager) loadCredential() error {
 
 	credentials, err := gm.appKeystore.GetMembershipCredentials(
 		gm.keystorePassword,
-		nil,
+		gm.membershipIndex,
-		[]keystore.MembershipContract{{
+		keystore.NewMembershipContractInfo(gm.chainId, gm.membershipContractAddress))
-			ChainID: fmt.Sprintf("0x%X", gm.chainId),
-			Address: gm.membershipContractAddress.Hex(),
-		}})
 	if err != nil {
 		return err
 	}
 	gm.metrics.RecordMembershipCredentialsImportDuration(time.Since(start))
 
-	if len(credentials) == 0 {
+	if credentials == nil {
 		return errors.New("no credentials available")
 	}
 
-	if int(gm.keystoreIndex) > len(credentials)-1 {
+	gm.identityCredential = credentials.IdentityCredential
-		return errors.New("invalid keystore index")
-	}
-
-	if int(gm.membershipGroupIndex) > len(credentials[gm.keystoreIndex].MembershipGroups)-1 {
-		return errors.New("invalid membership group index")
-	}
-
-	gm.identityCredential = credentials[gm.keystoreIndex].IdentityCredential
-	gm.membershipIndex = &credentials[gm.keystoreIndex].MembershipGroups[gm.membershipGroupIndex].TreeIndex
 
 	return nil
 }
@@ -282,12 +265,8 @@ func (gm *DynamicGroupManager) IdentityCredentials() (rln.IdentityCredential, er
 	return *gm.identityCredential, nil
 }
 
-func (gm *DynamicGroupManager) MembershipIndex() (rln.MembershipIndex, error) {
+func (gm *DynamicGroupManager) MembershipIndex() rln.MembershipIndex {
-	if gm.membershipIndex == nil {
+	return gm.membershipIndex
-		return 0, errors.New("membership index has not been setup")
-	}
-
-	return *gm.membershipIndex, nil
 }
 
 // Stop stops all go-routines, eth client and closes the rln database

waku/v2/protocol/rln/group_manager/static/static.go

@@ -14,7 +14,7 @@ type StaticGroupManager struct {
 	log *zap.Logger
 
 	identityCredential *rln.IdentityCredential
-	membershipIndex    *rln.MembershipIndex
+	membershipIndex    rln.MembershipIndex
 
 	group       []rln.IDCommitment
 	rootTracker *group_manager.MerkleRootTracker
@@ -36,7 +36,7 @@ func NewStaticGroupManager(
 		log:                log.Named("rln-static"),
 		group:              group,
 		identityCredential: &identityCredential,
-		membershipIndex:    &index,
+		membershipIndex:    index,
 	}, nil
 }
 
@@ -85,12 +85,8 @@ func (gm *StaticGroupManager) IdentityCredentials() (rln.IdentityCredential, err
 	return *gm.identityCredential, nil
 }
 
-func (gm *StaticGroupManager) MembershipIndex() (rln.MembershipIndex, error) {
+func (gm *StaticGroupManager) MembershipIndex() rln.MembershipIndex {
-	if gm.membershipIndex == nil {
+	return gm.membershipIndex
-		return 0, errors.New("membership index has not been setup")
-	}
-
-	return *gm.membershipIndex, nil
 }
 
 // Stop is a function created just to comply with the GroupManager interface (it does nothing)

waku/v2/protocol/rln/keystore/keystore.go

@@ -2,28 +2,29 @@ package keystore
 
 import (
 	"bytes"
+	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"os"
-	"sort"
+	"strings"
 
 	"github.com/ethereum/go-ethereum/accounts/keystore"
+	"github.com/waku-org/go-waku/waku/v2/hash"
 	"github.com/waku-org/go-zerokit-rln/rln"
 	"go.uber.org/zap"
 )
 
-// DefaultCredentialsFilename is the default filename for the rln credentials keystore
+// DefaultCredentialsFilename is the suggested default filename for the rln credentials keystore
-const DefaultCredentialsFilename = "rlnKeystore.json"
+const DefaultCredentialsFilename = "./rlnKeystore.json"
 
-// DefaultCredentialsPassword contains the default password used when no password is specified
+// DefaultCredentialsPassword is the suggested default password for the rln credentials store
 const DefaultCredentialsPassword = "password"
 
 // New creates a new instance of a rln credentials keystore
-func New(keystorePath string, appInfo AppInfo, logger *zap.Logger) (*AppKeystore, error) {
+func New(path string, appInfo AppInfo, logger *zap.Logger) (*AppKeystore, error) {
 	logger = logger.Named("rln-keystore")
 
-	path := keystorePath
 	if path == "" {
 		logger.Warn("keystore: no credentials path set, using default path", zap.String("path", DefaultCredentialsFilename))
 		path = DefaultCredentialsFilename
@@ -53,13 +54,18 @@ func New(keystorePath string, appInfo AppInfo, logger *zap.Logger) (*AppKeystore
 		}
 
 		keystore := new(AppKeystore)
-		keystore.logger = logger
-		keystore.path = path
 		err := json.Unmarshal(keystoreBytes, keystore)
 		if err != nil {
 			continue
 		}
 
+		keystore.logger = logger
+		keystore.path = path
+
+		if keystore.Credentials == nil {
+			keystore.Credentials = map[Key]appKeystoreCredential{}
+		}
+
 		if keystore.AppIdentifier == appInfo.AppIdentifier && keystore.Application == appInfo.Application && keystore.Version == appInfo.Version {
 			return keystore, nil
 		}
@@ -68,128 +74,67 @@ func New(keystorePath string, appInfo AppInfo, logger *zap.Logger) (*AppKeystore
 	return nil, errors.New("no keystore found")
 }
 
+func getKey(treeIndex rln.MembershipIndex, filterMembershipContract MembershipContractInfo) (Key, error) {
+	keyStr := fmt.Sprintf("%s%s%d", filterMembershipContract.ChainID, filterMembershipContract.Address, treeIndex)
+	hash := hash.SHA256([]byte(keyStr))
+	return Key(strings.ToUpper(hex.EncodeToString(hash))), nil
+}
+
 // GetMembershipCredentials decrypts and retrieves membership credentials from the keystore applying filters
-func (k *AppKeystore) GetMembershipCredentials(keystorePassword string, filterIdentityCredentials []MembershipCredentials, filterMembershipContracts []MembershipContract) ([]MembershipCredentials, error) {
+func (k *AppKeystore) GetMembershipCredentials(keystorePassword string, treeIndex rln.MembershipIndex, filterMembershipContract MembershipContractInfo) (*MembershipCredentials, error) {
-	password := keystorePassword
+	key, err := getKey(treeIndex, filterMembershipContract)
-	if password == "" {
+	if err != nil {
-		k.logger.Warn("keystore: no credentials password set, using default password", zap.String("password", DefaultCredentialsPassword))
+		return nil, err
-		password = DefaultCredentialsPassword
 	}
 
-	var result []MembershipCredentials
+	credential, ok := k.Credentials[key]
-
+	if !ok {
-	for _, credential := range k.Credentials {
+		return nil, nil
-		credentialsBytes, err := keystore.DecryptDataV3(credential.Crypto, password)
-		if err != nil {
-			return nil, err
-		}
-
-		var credentials MembershipCredentials
-		err = json.Unmarshal(credentialsBytes, &credentials)
-		if err != nil {
-			return nil, err
-		}
-
-		filteredCredential := filterCredential(credentials, filterIdentityCredentials, filterMembershipContracts)
-		if filteredCredential != nil {
-			result = append(result, *filteredCredential)
-		}
 	}
 
-	return result, nil
+	credentialsBytes, err := keystore.DecryptDataV3(credential.Crypto, keystorePassword)
+	if err != nil {
+		return nil, err
+	}
+
+	credentials := new(MembershipCredentials)
+	err = json.Unmarshal(credentialsBytes, credentials)
+	if err != nil {
+		return nil, err
+	}
+
+	return credentials, nil
 }
 
 // AddMembershipCredentials inserts a membership credential to the keystore matching the application, appIdentifier and version filters.
-func (k *AppKeystore) AddMembershipCredentials(newIdentityCredential *rln.IdentityCredential, newMembershipGroup MembershipGroup, password string) (membershipGroupIndex uint, err error) {
+func (k *AppKeystore) AddMembershipCredentials(newCredential MembershipCredentials, password string) error {
-	// A flag to tell us if the keystore contains a credential associated to the input identity credential, i.e. membershipCredential
+	credentials, err := k.GetMembershipCredentials(password, newCredential.TreeIndex, newCredential.MembershipContractInfo)
-	found := false
+	if err != nil {
-	for i, existingCredentials := range k.Credentials {
+		return err
-		credentialsBytes, err := keystore.DecryptDataV3(existingCredentials.Crypto, password)
-		if err != nil {
-			continue
-		}
-
-		var credentials MembershipCredentials
-		err = json.Unmarshal(credentialsBytes, &credentials)
-		if err != nil {
-			continue
-		}
-
-		if rln.IdentityCredentialEquals(*credentials.IdentityCredential, *newIdentityCredential) {
-			// idCredential is present in keystore. We add the input credential membership group to the one contained in the decrypted keystore credential (we deduplicate groups using sets)
-			allMembershipsMap := make(map[MembershipGroup]struct{})
-			for _, m := range credentials.MembershipGroups {
-				allMembershipsMap[m] = struct{}{}
-			}
-			allMembershipsMap[newMembershipGroup] = struct{}{}
-
-			// We sort membership groups, otherwise we will not have deterministic results in tests
-			var allMemberships []MembershipGroup
-			for k := range allMembershipsMap {
-				allMemberships = append(allMemberships, k)
-			}
-			sort.Slice(allMemberships, func(i, j int) bool {
-				return allMemberships[i].MembershipContract.Address < allMemberships[j].MembershipContract.Address
-			})
-
-			// we define the updated credential with the updated membership sets
-			updatedCredential := MembershipCredentials{
-				IdentityCredential: newIdentityCredential,
-				MembershipGroups:   allMemberships,
-			}
-
-			// we re-encrypt creating a new keyfile
-			b, err := json.Marshal(updatedCredential)
-			if err != nil {
-				return 0, err
-			}
-
-			encryptedCredentials, err := keystore.EncryptDataV3(b, []byte(password), keystore.StandardScryptN, keystore.StandardScryptP)
-			if err != nil {
-				return 0, err
-			}
-
-			// we update the original credential field in keystoreCredentials
-			k.Credentials[i] = appKeystoreCredential{Crypto: encryptedCredentials}
-
-			found = true
-
-			// We setup the return values
-			membershipGroupIndex = uint(len(allMemberships))
-			for mIdx, mg := range updatedCredential.MembershipGroups {
-				if mg.MembershipContract.Equals(newMembershipGroup.MembershipContract) {
-					membershipGroupIndex = uint(mIdx)
-					break
-				}
-			}
-
-			// We stop decrypting other credentials in the keystore
-			break
-		}
 	}
 
-	if !found { // Not found
+	key, err := getKey(newCredential.TreeIndex, newCredential.MembershipContractInfo)
-		newCredential := MembershipCredentials{
+	if err != nil {
-			IdentityCredential: newIdentityCredential,
+		return err
-			MembershipGroups:   []MembershipGroup{newMembershipGroup},
-		}
-
-		b, err := json.Marshal(newCredential)
-		if err != nil {
-			return 0, err
-		}
-
-		encryptedCredentials, err := keystore.EncryptDataV3(b, []byte(password), keystore.StandardScryptN, keystore.StandardScryptP)
-		if err != nil {
-			return 0, err
-		}
-
-		k.Credentials = append(k.Credentials, appKeystoreCredential{Crypto: encryptedCredentials})
-
-		membershipGroupIndex = uint(len(newCredential.MembershipGroups) - 1)
 	}
 
-	return membershipGroupIndex, save(k, k.path)
+	if credentials != nil {
+		return errors.New("credential already present")
+	}
+
+	b, err := json.Marshal(newCredential)
+	if err != nil {
+		return err
+	}
+
+	encryptedCredentials, err := keystore.EncryptDataV3(b, []byte(password), keystore.StandardScryptN, keystore.StandardScryptP)
+	if err != nil {
+		return err
+	}
+
+	k.Credentials[key] = appKeystoreCredential{Crypto: encryptedCredentials}
+
+	return save(k, k.path)
 }
 
 func createAppKeystore(path string, appInfo AppInfo, separator string) error {
@@ -201,6 +146,7 @@ func createAppKeystore(path string, appInfo AppInfo, separator string) error {
 		Application:   appInfo.Application,
 		AppIdentifier: appInfo.AppIdentifier,
 		Version:       appInfo.Version,
+		Credentials:   make(map[Key]appKeystoreCredential),
 	}
 
 	b, err := json.Marshal(keystore)
@@ -220,46 +166,6 @@ func createAppKeystore(path string, appInfo AppInfo, separator string) error {
 	return os.WriteFile(path, buffer.Bytes(), 0600)
 }
 
-func filterCredential(credential MembershipCredentials, filterIdentityCredentials []MembershipCredentials, filterMembershipContracts []MembershipContract) *MembershipCredentials {
-	if len(filterIdentityCredentials) != 0 {
-		found := false
-		for _, filterCreds := range filterIdentityCredentials {
-			if filterCreds.Equals(credential) {
-				found = true
-			}
-		}
-		if !found {
-			return nil
-		}
-	}
-
-	if len(filterMembershipContracts) != 0 {
-		var membershipGroupsIntersection []MembershipGroup
-		for _, filterContract := range filterMembershipContracts {
-			for _, credentialGroups := range credential.MembershipGroups {
-				if filterContract.Equals(credentialGroups.MembershipContract) {
-					membershipGroupsIntersection = append(membershipGroupsIntersection, credentialGroups)
-				}
-			}
-		}
-		if len(membershipGroupsIntersection) != 0 {
-			// If we have a match on some groups, we return the credential with filtered groups
-			return &MembershipCredentials{
-				IdentityCredential: credential.IdentityCredential,
-				MembershipGroups:   membershipGroupsIntersection,
-			}
-		} else {
-			return nil
-		}
-	}
-
-	// We hit this return only if
-	// - filterIdentityCredentials.len() == 0 and filterMembershipContracts.len() == 0 (no filter)
-	// - filterIdentityCredentials.len() != 0 and filterMembershipContracts.len() == 0 (filter only on identity credential)
-	// Indeed, filterMembershipContracts.len() != 0 will have its exclusive return based on all values of membershipGroupsIntersection.len()
-	return &credential
-}
-
 // Safely saves a Keystore's JsonNode to disk.
 // If exists, the destination file is renamed with extension .bkp; the file is written at its destination and the .bkp file is removed if write is successful, otherwise is restored
 func save(keystore *AppKeystore, path string) error {

waku/v2/protocol/rln/keystore/types.go

@@ -1,59 +1,93 @@
 package keystore
 
 import (
+	"fmt"
+	"math/big"
+	"strings"
+
 	"github.com/ethereum/go-ethereum/accounts/keystore"
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/common/hexutil"
 	"github.com/waku-org/go-zerokit-rln/rln"
 	"go.uber.org/zap"
 )
 
-// MembershipContract contains information about a membership smart contract address and the chain in which it is deployed
+// MembershipContractInfo contains information about a membership smart contract address and the chain in which it is deployed
-type MembershipContract struct {
+type MembershipContractInfo struct {
-	ChainID string `json:"chainId"`
+	ChainID ChainID         `json:"chainId"`
-	Address string `json:"address"`
+	Address ContractAddress `json:"address"`
+}
+
+// NewMembershipContractInfo generates a new MembershipContract instance
+func NewMembershipContractInfo(chainID *big.Int, address common.Address) MembershipContractInfo {
+	return MembershipContractInfo{
+		ChainID: ChainID{
+			chainID,
+		},
+		Address: ContractAddress(address),
+	}
+}
+
+// ContractAddress is a common.Address created to comply with the expected marshalling for the credentials
+type ContractAddress common.Address
+
+// MarshalText is used to convert a ContractAddress into a valid value expected by the json encoder
+func (c ContractAddress) MarshalText() ([]byte, error) {
+	return []byte(common.Address(c).Hex()), nil
+}
+
+// UnmarshalText converts a byte slice into a ContractAddress
+func (c *ContractAddress) UnmarshalText(text []byte) error {
+	b, err := hexutil.Decode(string(text))
+	if err != nil {
+		return err
+	}
+	copy(c[:], b[:])
+	return nil
+}
+
+// ChainID is a helper struct created to comply with the expected marshalling for the credentials
+type ChainID struct {
+	*big.Int
+}
+
+// String returns a string with the expected chainId format for the credentials
+func (c ChainID) String() string {
+	return fmt.Sprintf(`"%s"`, hexutil.EncodeBig(c.Int))
+}
+
+// MarshalJSON is used to convert a ChainID into a valid value expected by the json encoder
+func (c ChainID) MarshalJSON() (text []byte, err error) {
+	return []byte(c.String()), nil
+}
+
+// UnmarshalJSON converts a byte slice into a ChainID
+func (c *ChainID) UnmarshalJSON(text []byte) error {
+	hexVal := strings.ReplaceAll(string(text), `"`, "")
+	b, err := hexutil.DecodeBig(hexVal)
+	if err != nil {
+		return err
+	}
+
+	c.Int = b
+	return nil
 }
 
 // Equals is used to compare MembershipContract
-func (m MembershipContract) Equals(other MembershipContract) bool {
+func (m MembershipContractInfo) Equals(other MembershipContractInfo) bool {
-	return m.Address == other.Address && m.ChainID == other.ChainID
+	return m.Address == other.Address && m.ChainID.Int64() == other.ChainID.Int64()
-}
-
-// MembershipGroup contains information about the index in which a credential is stored in the merkle tree and the contract associated to this credential
-type MembershipGroup struct {
-	MembershipContract MembershipContract  `json:"membershipContract"`
-	TreeIndex          rln.MembershipIndex `json:"treeIndex"`
-}
-
-// Equals is used to compare MembershipGroup
-func (m MembershipGroup) Equals(other MembershipGroup) bool {
-	return m.MembershipContract.Equals(other.MembershipContract) && m.TreeIndex == other.TreeIndex
 }
 
 // MembershipCredentials contains all the information about an RLN Identity Credential and membership group it belongs to
 type MembershipCredentials struct {
-	IdentityCredential *rln.IdentityCredential `json:"identityCredential"`
+	IdentityCredential     *rln.IdentityCredential `json:"identityCredential"`
-	MembershipGroups   []MembershipGroup       `json:"membershipGroups"`
+	MembershipContractInfo MembershipContractInfo  `json:"membershipContract"`
+	TreeIndex              rln.MembershipIndex     `json:"treeIndex"`
 }
 
 // Equals is used to compare MembershipCredentials
 func (m MembershipCredentials) Equals(other MembershipCredentials) bool {
-	if !rln.IdentityCredentialEquals(*m.IdentityCredential, *other.IdentityCredential) {
+	return rln.IdentityCredentialEquals(*m.IdentityCredential, *other.IdentityCredential) && m.MembershipContractInfo.Equals(other.MembershipContractInfo) && m.TreeIndex == other.TreeIndex
-		return false
-	}
-
-	for _, x := range m.MembershipGroups {
-		found := false
-		for _, y := range other.MembershipGroups {
-			if x.Equals(y) {
-				found = true
-				break
-			}
-		}
-		if !found {
-			return false
-		}
-	}
-
-	return true
 }
 
 // AppInfo is a helper structure that contains information about the application that uses these credentials
@@ -63,12 +97,15 @@ type AppInfo struct {
 	Version       string `json:"version"`
 }
 
+// Key is a helper type created to represent the key in a map of credentials
+type Key string
+
 // AppKeystore represents the membership credentials to be used in RLN
 type AppKeystore struct {
-	Application   string                  `json:"application"`
+	Application   string                        `json:"application"`
-	AppIdentifier string                  `json:"appIdentifier"`
+	AppIdentifier string                        `json:"appIdentifier"`
-	Credentials   []appKeystoreCredential `json:"credentials"`
+	Credentials   map[Key]appKeystoreCredential `json:"credentials"`
-	Version       string                  `json:"version"`
+	Version       string                        `json:"version"`
 
 	path   string
 	logger *zap.Logger

waku/v2/protocol/rln/onchain_test.go

@@ -6,7 +6,6 @@ package rln
 import (
 	"context"
 	"crypto/ecdsa"
-	"fmt"
 	"math/big"
 	"os"
 	"testing"
@@ -103,7 +102,7 @@ func (s *WakuRLNRelayDynamicSuite) generateCredentials(rlnInstance *rln.RLN) *rl
 	return identityCredential
 }
 
-func (s *WakuRLNRelayDynamicSuite) register(identityCredential *rln.IdentityCredential, privKey *ecdsa.PrivateKey, keystorePath string) (rln.MembershipIndex, uint) {
+func (s *WakuRLNRelayDynamicSuite) register(appKeystore *keystore.AppKeystore, identityCredential *rln.IdentityCredential, privKey *ecdsa.PrivateKey) rln.MembershipIndex {
 
 	auth, err := bind.NewKeyedTransactorWithChainID(privKey, s.chainID)
 	s.Require().NoError(err)
@@ -123,18 +122,16 @@ func (s *WakuRLNRelayDynamicSuite) register(identityCredential *rln.IdentityCred
 
 	membershipIndex := rln.MembershipIndex(uint(evt.Index.Int64()))
 
-	membershipGroup := keystore.MembershipGroup{
+	membershipCredential := keystore.MembershipCredentials{
-		TreeIndex: membershipIndex,
+		IdentityCredential:     identityCredential,
-		MembershipContract: keystore.MembershipContract{
+		TreeIndex:              membershipIndex,
-			ChainId: fmt.Sprintf("0x%X", s.chainID.Int64()),
+		MembershipContractInfo: keystore.NewMembershipContractInfo(s.chainID, s.rlnAddr),
-			Address: s.rlnAddr.String(),
-		},
 	}
 
-	membershipGroupIndex, err := keystore.AddMembershipCredentials(keystorePath, identityCredential, membershipGroup, keystorePassword, dynamic.RLNAppInfo, keystore.DefaultSeparator)
+	err = appKeystore.AddMembershipCredentials(membershipCredential, keystorePassword)
 	s.Require().NoError(err)
 
-	return membershipIndex, membershipGroupIndex
+	return membershipIndex
 }
 
 func (s *WakuRLNRelayDynamicSuite) TestDynamicGroupManagement() {
@@ -147,10 +144,13 @@ func (s *WakuRLNRelayDynamicSuite) TestDynamicGroupManagement() {
 
 	u1Credentials := s.generateCredentials(rlnInstance)
 	keystorePath1 := "./test_onchain.json"
-	_, membershipGroupIndex := s.register(u1Credentials, s.u1PrivKey, keystorePath1)
+	appKeystore, err := keystore.New(keystorePath1, dynamic.RLNAppInfo, utils.Logger())
+	s.Require().NoError(err)
+
+	membershipIndex := s.register(appKeystore, u1Credentials, s.u1PrivKey)
 	defer s.removeCredentials(keystorePath1)
 
-	gm, err := dynamic.NewDynamicGroupManager(s.clientAddr, s.rlnAddr, membershipGroupIndex, keystorePath1, keystorePassword, 0, false, prometheus.DefaultRegisterer, utils.Logger())
+	gm, err := dynamic.NewDynamicGroupManager(s.clientAddr, s.rlnAddr, membershipIndex, appKeystore, keystorePassword, prometheus.DefaultRegisterer, utils.Logger())
 	s.Require().NoError(err)
 
 	// initialize the WakuRLNRelay
@@ -167,7 +167,10 @@ func (s *WakuRLNRelayDynamicSuite) TestDynamicGroupManagement() {
 
 	u2Credentials := s.generateCredentials(rlnInstance)
 	keystorePath2 := "./test_onchain2.json"
-	membershipIndex, _ := s.register(u2Credentials, s.u2PrivKey, keystorePath2)
+	appKeystore2, err := keystore.New(keystorePath2, dynamic.RLNAppInfo, utils.Logger())
+	s.Require().NoError(err)
+
+	membershipIndex = s.register(appKeystore2, u2Credentials, s.u2PrivKey)
 	defer s.removeCredentials(keystorePath2)
 
 	time.Sleep(1 * time.Second)
@@ -187,7 +190,10 @@ func (s *WakuRLNRelayDynamicSuite) TestInsertKeyMembershipContract() {
 	credentials3 := s.generateCredentials(rlnInstance)
 
 	keystorePath1 := "./test_onchain.json"
-	s.register(credentials1, s.u1PrivKey, keystorePath1)
+	appKeystore, err := keystore.New(keystorePath1, dynamic.RLNAppInfo, utils.Logger())
+	s.Require().NoError(err)
+
+	s.register(appKeystore, credentials1, s.u1PrivKey)
 	defer s.removeCredentials(keystorePath1)
 
 	// Batch Register
@@ -224,14 +230,18 @@ func (s *WakuRLNRelayDynamicSuite) TestMerkleTreeConstruction() {
 	s.Require().NoError(err)
 
 	// register the members to the contract
-	_, membershipGroupIndex := s.register(credentials1, s.u1PrivKey, "./test_onchain.json")
+	keystorePath1 := "./test_onchain.json"
-	_, membershipGroupIndex = s.register(credentials2, s.u1PrivKey, "./test_onchain.json")
+	appKeystore, err := keystore.New(keystorePath1, dynamic.RLNAppInfo, utils.Logger())
+	s.Require().NoError(err)
+
+	membershipIndex := s.register(appKeystore, credentials1, s.u1PrivKey)
+	membershipIndex = s.register(appKeystore, credentials2, s.u1PrivKey)
+
+	defer s.removeCredentials(keystorePath1)
 
 	// mount the rln relay protocol in the on-chain/dynamic mode
-	// TODO: This assumes the keystoreIndex is 0, but there are two possible credentials in this keystore due to using the same contract address
+	gm, err := dynamic.NewDynamicGroupManager(s.clientAddr, s.rlnAddr, membershipIndex, appKeystore, keystorePassword, prometheus.DefaultRegisterer, utils.Logger())
-	// when credentials1 and credentials2 were registered. We should remove this hardcoded value and obtain the correct value when the credentials are persisted
+
-	keystoreIndex := uint(0)
-	gm, err := dynamic.NewDynamicGroupManager(s.clientAddr, s.rlnAddr, membershipGroupIndex, "./test_onchain.json", keystorePassword, keystoreIndex, false, prometheus.DefaultRegisterer, utils.Logger())
 	s.Require().NoError(err)
 
 	rlnRelay, err := New(gm, "test-merkle-tree.db", timesource.NewDefaultClock(), prometheus.DefaultRegisterer, utils.Logger())
@@ -261,11 +271,13 @@ func (s *WakuRLNRelayDynamicSuite) TestCorrectRegistrationOfPeers() {
 	// Register credentials1 in contract and keystore1
 	credentials1 := s.generateCredentials(rlnInstance)
 	keystorePath1 := "./test_onchain.json"
-	_, membershipGroupIndex := s.register(credentials1, s.u1PrivKey, keystorePath1)
+	appKeystore, err := keystore.New(keystorePath1, dynamic.RLNAppInfo, utils.Logger())
+	s.Require().NoError(err)
+	membershipGroupIndex := s.register(appKeystore, credentials1, s.u1PrivKey)
 	defer s.removeCredentials(keystorePath1)
 
 	// mount the rln relay protocol in the on-chain/dynamic mode
-	gm1, err := dynamic.NewDynamicGroupManager(s.clientAddr, s.rlnAddr, membershipGroupIndex, keystorePath1, keystorePassword, 0, false, prometheus.DefaultRegisterer, utils.Logger())
+	gm1, err := dynamic.NewDynamicGroupManager(s.clientAddr, s.rlnAddr, membershipGroupIndex, appKeystore, keystorePassword, prometheus.DefaultRegisterer, utils.Logger())
 	s.Require().NoError(err)
 
 	rlnRelay1, err := New(gm1, "test-correct-registration-1.db", timesource.NewDefaultClock(), prometheus.DefaultRegisterer, utils.Logger())
@@ -278,11 +290,13 @@ func (s *WakuRLNRelayDynamicSuite) TestCorrectRegistrationOfPeers() {
 	// Register credentials2 in contract and keystore2
 	credentials2 := s.generateCredentials(rlnInstance)
 	keystorePath2 := "./test_onchain2.json"
-	_, membershipGroupIndex = s.register(credentials2, s.u2PrivKey, keystorePath2)
+	appKeystore2, err := keystore.New(keystorePath2, dynamic.RLNAppInfo, utils.Logger())
+	s.Require().NoError(err)
+	membershipGroupIndex = s.register(appKeystore2, credentials2, s.u2PrivKey)
 	defer s.removeCredentials(keystorePath2)
 
 	// mount the rln relay protocol in the on-chain/dynamic mode
-	gm2, err := dynamic.NewDynamicGroupManager(s.clientAddr, s.rlnAddr, membershipGroupIndex, keystorePath2, keystorePassword, 0, false, prometheus.DefaultRegisterer, utils.Logger())
+	gm2, err := dynamic.NewDynamicGroupManager(s.clientAddr, s.rlnAddr, membershipGroupIndex, appKeystore2, keystorePassword, prometheus.DefaultRegisterer, utils.Logger())
 	s.Require().NoError(err)
 
 	rlnRelay2, err := New(gm2, "test-correct-registration-2.db", timesource.NewDefaultClock(), prometheus.DefaultRegisterer, utils.Logger())
@@ -294,9 +308,8 @@ func (s *WakuRLNRelayDynamicSuite) TestCorrectRegistrationOfPeers() {
 	// the two nodes should be registered into the contract
 	// since nodes are spun up sequentially
 	// the first node has index 0 whereas the second node gets index 1
-	idx1, err := rlnRelay1.groupManager.MembershipIndex()
+	idx1 := rlnRelay1.groupManager.MembershipIndex()
-	s.Require().NoError(err)
+	idx2 := rlnRelay2.groupManager.MembershipIndex()
-	idx2, err := rlnRelay2.groupManager.MembershipIndex()
 	s.Require().NoError(err)
 
 	s.Require().Equal(rln.MembershipIndex(0), idx1)

waku/v2/protocol/rln/waku_rln_relay.go

@@ -25,7 +25,7 @@ import (
 type GroupManager interface {
 	Start(ctx context.Context, rln *rln.RLN, rootTracker *group_manager.MerkleRootTracker) error
 	IdentityCredentials() (rln.IdentityCredential, error)
-	MembershipIndex() (rln.MembershipIndex, error)
+	MembershipIndex() rln.MembershipIndex
 	Stop() error
 }
 
@@ -356,10 +356,7 @@ func (rlnRelay *WakuRLNRelay) generateProof(input []byte, epoch rln.Epoch) (*pb.
 		return nil, err
 	}
 
-	membershipIndex, err := rlnRelay.groupManager.MembershipIndex()
+	membershipIndex := rlnRelay.groupManager.MembershipIndex()
-	if err != nil {
-		return nil, err
-	}
 
 	proof, err := rlnRelay.RLN.GenerateProof(input, identityCredentials, membershipIndex, epoch)
 	if err != nil {
@@ -381,6 +378,6 @@ func (rlnRelay *WakuRLNRelay) IdentityCredential() (rln.IdentityCredential, erro
 	return rlnRelay.groupManager.IdentityCredentials()
 }
 
-func (rlnRelay *WakuRLNRelay) MembershipIndex() (uint, error) {
+func (rlnRelay *WakuRLNRelay) MembershipIndex() uint {
 	return rlnRelay.groupManager.MembershipIndex()
 }