chore_: limit max number of ephemeral keys to 3

This commit limits the number of ephemeral keys used by the push
notification client to 3, as it was noticed that it would greatly
increase the number of filters installed and was unbound.

protocol/common/message_sender.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/ecdsa"
 	"database/sql"
+	"math/rand"
 	"sync"
 	"time"
 
@@ -33,8 +34,9 @@ const (
 	whisperLargeSizePoW = 0.000002
 	// largeSizeInBytes is when should we be using a lower POW.
 	// Roughly this is 50KB
-	largeSizeInBytes = 50000
+	largeSizeInBytes              = 50000
-	whisperPoWTime   = 5
+	whisperPoWTime                = 5
+	maxMessageSenderEphemeralKeys = 3
 )
 
 // RekeyCompatibility indicates whether we should be sending
@@ -1245,15 +1247,37 @@ func (s *MessageSender) JoinPublic(id string) (*transport.Filter, error) {
 	return s.transport.JoinPublic(id)
 }
 
-// AddEphemeralKey adds an ephemeral key that we will be listening to
+func (s *MessageSender) getRandomEphemeralKey() *ecdsa.PrivateKey {
-// note that we never removed them from now, as waku/whisper does not
+	k := rand.Intn(len(s.ephemeralKeys)) //nolint: gosec
-// recalculate topics on removal, so effectively there's no benefit.
+	for _, key := range s.ephemeralKeys {
-// On restart they will be gone.
+		if k == 0 {
-func (s *MessageSender) AddEphemeralKey(privateKey *ecdsa.PrivateKey) (*transport.Filter, error) {
+			return key
+		}
+		k--
+	}
+	return nil
+}
+
+func (s *MessageSender) GetEphemeralKey() (*ecdsa.PrivateKey, error) {
 	s.ephemeralKeysMutex.Lock()
+	if len(s.ephemeralKeys) >= maxMessageSenderEphemeralKeys {
+		s.ephemeralKeysMutex.Unlock()
+		return s.getRandomEphemeralKey(), nil
+	}
+	privateKey, err := crypto.GenerateKey()
+	if err != nil {
+		s.ephemeralKeysMutex.Unlock()
+		return nil, err
+	}
+
 	s.ephemeralKeys[types.EncodeHex(crypto.FromECDSAPub(&privateKey.PublicKey))] = privateKey
 	s.ephemeralKeysMutex.Unlock()
-	return s.transport.LoadKeyFilters(privateKey)
+	_, err = s.transport.LoadKeyFilters(privateKey)
+	if err != nil {
+		return nil, err
+	}
+
+	return privateKey, nil
 }
 
 func MessageSpecToWhisper(spec *encryption.ProtocolMessageSpec) (*types.NewMessage, error) {

protocol/common/message_sender_test.go

@@ -361,3 +361,20 @@ func (s *MessageSenderSuite) TestHandleSegmentMessages() {
 	_, err = s.sender.HandleMessages(message)
 	s.Require().ErrorIs(err, ErrMessageSegmentsAlreadyCompleted)
 }
+
+func (s *MessageSenderSuite) TestGetEphemeralKey() {
+	keyMap := make(map[string]bool)
+	for i := 0; i < maxMessageSenderEphemeralKeys; i++ {
+		key, err := s.sender.GetEphemeralKey()
+		s.Require().NoError(err)
+		s.Require().NotNil(key)
+		keyMap[PubkeyToHex(&key.PublicKey)] = true
+	}
+	s.Require().Len(keyMap, maxMessageSenderEphemeralKeys)
+	// Add one more
+	key, err := s.sender.GetEphemeralKey()
+	s.Require().NoError(err)
+	s.Require().NotNil(key)
+
+	s.Require().True(keyMap[PubkeyToHex(&key.PublicKey)])
+}

protocol/pushnotificationclient/client.go

@@ -1375,12 +1375,7 @@ func (c *Client) SendNotification(publicKey *ecdsa.PublicKey, installationIDs []
 
 	c.config.Logger.Debug("actionable info", zap.Int("count", len(actionableInfos)))
 
-	// add ephemeral key and listen to it
+	ephemeralKey, err := c.messageSender.GetEphemeralKey()
-	ephemeralKey, err := crypto.GenerateKey()
-	if err != nil {
-		return nil, err
-	}
-	_, err = c.messageSender.AddEphemeralKey(ephemeralKey)
 	if err != nil {
 		return nil, err
 	}
@@ -1688,7 +1683,7 @@ func (c *Client) queryPushNotificationInfo(publicKey *ecdsa.PublicKey) error {
 		return err
 	}
 
-	ephemeralKey, err := crypto.GenerateKey()
+	ephemeralKey, err := c.messageSender.GetEphemeralKey()
 	if err != nil {
 		return err
 	}
@@ -1701,11 +1696,6 @@ func (c *Client) queryPushNotificationInfo(publicKey *ecdsa.PublicKey) error {
 		MessageType:         protobuf.ApplicationMetadataMessage_PUSH_NOTIFICATION_QUERY,
 	}
 
-	_, err = c.messageSender.AddEphemeralKey(ephemeralKey)
-	if err != nil {
-		return err
-	}
-
 	// this is the topic of message
 	encodedPublicKey := hex.EncodeToString(hashedPublicKey)
 	messageID, err := c.messageSender.SendPublic(context.Background(), encodedPublicKey, rawMessage)