Handle push notification registration e2e

protocol/messenger.go

@@ -7,6 +7,7 @@ import (
 	"io/ioutil"
 	"math/rand"
 	"os"
+	"reflect"
 	"sync"
 	"time"
 
@@ -252,7 +253,16 @@ func NewMessenger(
 	}
 
 	pushNotificationClientPersistence := push_notification_client.NewPersistence(database)
-	pushNotificationClient := push_notification_client.New(pushNotificationClientPersistence, processor)
+	// Hardcoding for now
+	pushNotificationClientConfig := &push_notification_client.Config{
+		Identity:                   identity,
+		TokenType:                  protobuf.PushNotificationRegistration_APN_TOKEN,
+		SendEnabled:                true,
+		Logger:                     logger,
+		RemoteNotificationsEnabled: true,
+		InstallationID:             installationID,
+	}
+	pushNotificationClient := push_notification_client.New(pushNotificationClientPersistence, pushNotificationClientConfig, processor)
 
 	handler := newMessageHandler(identity, logger, &sqlitePersistence{db: database})
 
@@ -1689,6 +1699,7 @@ func (m *Messenger) syncContact(ctx context.Context, contact *Contact) error {
 // RetrieveAll retrieves messages from all filters, processes them and returns a
 // MessengerResponse to the client
 func (m *Messenger) RetrieveAll() (*MessengerResponse, error) {
+	m.logger.Info("RETRIEVING ALL", zap.String("installation-id", m.installationID))
 	chatWithMessages, err := m.transport.RetrieveRawAll()
 	if err != nil {
 		return nil, err
@@ -1920,18 +1931,6 @@ func (m *Messenger) handleRetrievedMessages(chatWithMessages map[transport.Filte
 							logger.Warn("failed to handle ContactUpdate", zap.Error(err))
 							continue
 						}
-					case protobuf.PushNotificationRegistration:
-						logger.Debug("Received PushNotificationRegistration")
-						if m.pushNotificationServer == nil {
-							continue
-						}
-						logger.Debug("Handling PushNotificationRegistration")
-						// TODO: Compare DST with Identity
-						if err := m.pushNotificationServer.HandlePushNotificationRegistration2(publicKey, msg.ParsedMessage.([]byte)); err != nil {
-							logger.Warn("failed to handle PushNotificationRegistration", zap.Error(err))
-						}
-						// We continue in any case, no changes to messenger
-						continue
 					case protobuf.PushNotificationQuery:
 						logger.Debug("Received PushNotificationQuery")
 						if m.pushNotificationServer == nil {
@@ -1958,7 +1957,22 @@ func (m *Messenger) handleRetrievedMessages(chatWithMessages map[transport.Filte
 						continue
 
 					default:
-						logger.Debug("message not handled")
+						// Check if is an encrypted PushNotificationRegistration
+						if msg.Type == protobuf.ApplicationMetadataMessage_PUSH_NOTIFICATION_REGISTRATION {
+							logger.Debug("Received PushNotificationRegistration")
+							if m.pushNotificationServer == nil {
+								continue
+							}
+							logger.Debug("Handling PushNotificationRegistration")
+							// TODO: Compare DST with Identity
+							if err := m.pushNotificationServer.HandlePushNotificationRegistration2(publicKey, msg.ParsedMessage.([]byte)); err != nil {
+								logger.Warn("failed to handle PushNotificationRegistration", zap.Error(err))
+							}
+							// We continue in any case, no changes to messenger
+							continue
+						}
+
+						logger.Debug("message not handled", zap.Any("messageType", reflect.TypeOf(msg.ParsedMessage)))
 
 					}
 				}

protocol/push_notification_client/push_notification.go

@@ -1,22 +1,25 @@
 package push_notification_client
 
 import (
+	"context"
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/ecdsa"
 	"crypto/rand"
 	"errors"
 	"io"
+	"time"
 
-	"golang.org/x/crypto/sha3"
+	"github.com/golang/protobuf/proto"
-
 	"github.com/google/uuid"
 
 	"github.com/status-im/status-go/eth-node/crypto/ecies"
 	"github.com/status-im/status-go/protocol/common"
 	"github.com/status-im/status-go/protocol/protobuf"
+	"go.uber.org/zap"
 )
 
+const encryptedPayloadKeyLength = 16
 const accessTokenKeyLength = 16
 
 type PushNotificationServer struct {
@@ -46,6 +49,11 @@ type Config struct {
 	PushNotificationServers []*PushNotificationServer
 	// InstallationID is the installation-id for this device
 	InstallationID string
+
+	Logger *zap.Logger
+
+	// TokenType is the type of token
+	TokenType protobuf.PushNotificationRegistration_TokenType
 }
 
 type Client struct {
@@ -67,14 +75,19 @@ type Client struct {
 	//messageProcessor is a message processor used to send and being notified of messages
 
 	messageProcessor *common.MessageProcessor
+	//pushNotificationQueryResponses is a channel that listens to pushNotificationResponse
+
+	pushNotificationQueryResponses chan *protobuf.PushNotificationQueryResponse
 }
 
-func New(persistence *Persistence, processor *common.MessageProcessor) *Client {
+func New(persistence *Persistence, config *Config, processor *common.MessageProcessor) *Client {
 	return &Client{
-		quit:             make(chan struct{}),
+		quit:                           make(chan struct{}),
-		messageProcessor: processor,
+		config:                         config,
-		persistence:      persistence,
+		pushNotificationQueryResponses: make(chan *protobuf.PushNotificationQueryResponse),
-		reader:           rand.Reader}
+		messageProcessor:               processor,
+		persistence:                    persistence,
+		reader:                         rand.Reader}
 }
 
 func (c *Client) Start() error {
@@ -129,7 +142,7 @@ func (p *Client) mutedChatIDsHashes(chatIDs []string) [][]byte {
 	var mutedChatListHashes [][]byte
 
 	for _, chatID := range chatIDs {
-		mutedChatListHashes = append(mutedChatListHashes, shake256(chatID))
+		mutedChatListHashes = append(mutedChatListHashes, common.Shake256([]byte(chatID)))
 	}
 
 	return mutedChatListHashes
@@ -174,6 +187,7 @@ func (p *Client) buildPushNotificationRegistrationMessage(contactIDs []*ecdsa.Pu
 
 	options := &protobuf.PushNotificationRegistration{
 		AccessToken:     token,
+		TokenType:       p.config.TokenType,
 		Version:         p.lastPushNotificationVersion + 1,
 		InstallationId:  p.config.InstallationID,
 		Token:           p.DeviceToken,
@@ -184,15 +198,53 @@ func (p *Client) buildPushNotificationRegistrationMessage(contactIDs []*ecdsa.Pu
 	return options, nil
 }
 
-func (p *Client) Register(deviceToken string, contactIDs []*ecdsa.PublicKey, mutedChatIDs []string) error {
+func (c *Client) Register(deviceToken string, contactIDs []*ecdsa.PublicKey, mutedChatIDs []string) error {
-	servers, err := p.persistence.GetServers()
+	c.DeviceToken = deviceToken
+	servers, err := c.persistence.GetServers()
 	if err != nil {
 		return err
 	}
+
 	if len(servers) == 0 {
 		return errors.New("no servers to register with")
 	}
-	return nil
+
+	registration, err := c.buildPushNotificationRegistrationMessage(contactIDs, mutedChatIDs)
+	if err != nil {
+		return err
+	}
+
+	marshaledRegistration, err := proto.Marshal(registration)
+	if err != nil {
+		return err
+	}
+
+	for _, server := range servers {
+
+		encryptedRegistration, err := c.encryptRegistration(server.publicKey, marshaledRegistration)
+		if err != nil {
+			return err
+		}
+		rawMessage := &common.RawMessage{
+			Payload:     encryptedRegistration,
+			MessageType: protobuf.ApplicationMetadataMessage_PUSH_NOTIFICATION_REGISTRATION,
+		}
+
+		_, err = c.messageProcessor.SendPrivate(context.Background(), server.publicKey, rawMessage)
+
+		// Send message and wait for reply
+
+	}
+	for {
+		select {
+		case <-c.quit:
+			return nil
+		case <-time.After(5 * time.Second):
+			return errors.New("no query response received")
+		case <-c.pushNotificationQueryResponses:
+			return nil
+		}
+	}
 }
 
 // HandlePushNotificationRegistrationResponse should check whether the response was successful or not, retry if necessary otherwise store the result in the database
@@ -206,7 +258,14 @@ func (p *Client) HandlePushNotificationAdvertisement(info *protobuf.PushNotifica
 }
 
 // HandlePushNotificationQueryResponse should update the data in the database for a given user
-func (p *Client) HandlePushNotificationQueryResponse(response *protobuf.PushNotificationQueryResponse) error {
+func (c *Client) HandlePushNotificationQueryResponse(response *protobuf.PushNotificationQueryResponse) error {
+
+	c.config.Logger.Debug("received push notification query response", zap.Any("response", response))
+	select {
+	case c.pushNotificationQueryResponses <- response:
+	default:
+		return errors.New("could not process push notification query response")
+	}
 	return nil
 }
 
@@ -216,6 +275,7 @@ func (p *Client) HandlePushNotificationResponse(ack *protobuf.PushNotificationRe
 }
 
 func (c *Client) AddPushNotificationServer(publicKey *ecdsa.PublicKey) error {
+	c.config.Logger.Debug("adding push notification server", zap.Any("public-key", publicKey))
 	currentServers, err := c.persistence.GetServers()
 	if err != nil {
 		return err
@@ -270,9 +330,19 @@ func encryptAccessToken(plaintext []byte, key []byte, reader io.Reader) ([]byte,
 	return gcm.Seal(nonce, nonce, plaintext, nil), nil
 }
 
-func shake256(input string) []byte {
+func (c *Client) encryptRegistration(publicKey *ecdsa.PublicKey, payload []byte) ([]byte, error) {
-	buf := []byte(input)
+	sharedKey, err := c.generateSharedKey(publicKey)
-	h := make([]byte, 64)
+	if err != nil {
-	sha3.ShakeSum256(h, buf)
+		return nil, err
-	return h
+	}
+
+	return common.Encrypt(payload, sharedKey, c.reader)
+}
+
+func (c *Client) generateSharedKey(publicKey *ecdsa.PublicKey) ([]byte, error) {
+	return ecies.ImportECDSA(c.config.Identity).GenerateShared(
+		ecies.ImportECDSAPublic(publicKey),
+		encryptedPayloadKeyLength,
+		encryptedPayloadKeyLength,
+	)
 }

protocol/push_notification_server/push_notification_server.go

@@ -203,14 +203,14 @@ func (p *Server) HandlePushNotificationRequest(request *protobuf.PushNotificatio
 	return response
 }
 
-func (p *Server) HandlePushNotificationRegistration(publicKey *ecdsa.PublicKey, payload []byte) *protobuf.PushNotificationRegistrationResponse {
+func (s *Server) HandlePushNotificationRegistration(publicKey *ecdsa.PublicKey, payload []byte) *protobuf.PushNotificationRegistrationResponse {
+
+	s.config.Logger.Debug("handling push notification registration")
 	response := &protobuf.PushNotificationRegistrationResponse{
 		RequestId: common.Shake256(payload),
 	}
 
-	registration, err := p.ValidateRegistration(publicKey, payload)
+	registration, err := s.ValidateRegistration(publicKey, payload)
-	if registration != nil {
-	}
 
 	if err != nil {
 		if err == ErrInvalidPushNotificationRegistrationVersion {
@@ -218,6 +218,7 @@ func (p *Server) HandlePushNotificationRegistration(publicKey *ecdsa.PublicKey,
 		} else {
 			response.Error = protobuf.PushNotificationRegistrationResponse_MALFORMED_MESSAGE
 		}
+		s.config.Logger.Warn("registration did not validate", zap.Error(err))
 		return response
 	}
 
@@ -227,18 +228,22 @@ func (p *Server) HandlePushNotificationRegistration(publicKey *ecdsa.PublicKey,
 			Version:        registration.Version,
 			InstallationId: registration.InstallationId,
 		}
-		if err := p.persistence.SavePushNotificationRegistration(common.HashPublicKey(publicKey), emptyRegistration); err != nil {
+		if err := s.persistence.SavePushNotificationRegistration(common.HashPublicKey(publicKey), emptyRegistration); err != nil {
 			response.Error = protobuf.PushNotificationRegistrationResponse_INTERNAL_ERROR
+			s.config.Logger.Error("failed to unregister ", zap.Error(err))
 			return response
 		}
 
-	} else if err := p.persistence.SavePushNotificationRegistration(common.HashPublicKey(publicKey), registration); err != nil {
+	} else if err := s.persistence.SavePushNotificationRegistration(common.HashPublicKey(publicKey), registration); err != nil {
 		response.Error = protobuf.PushNotificationRegistrationResponse_INTERNAL_ERROR
+		s.config.Logger.Error("failed to save registration", zap.Error(err))
 		return response
 	}
 
 	response.Success = true
 
+	s.config.Logger.Debug("handled push notification registration successfully")
+
 	return response
 }
 

protocol/push_notification_test.go

@@ -6,6 +6,7 @@ import (
 	"io/ioutil"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/suite"
@@ -14,6 +15,7 @@ import (
 	gethbridge "github.com/status-im/status-go/eth-node/bridge/geth"
 	"github.com/status-im/status-go/eth-node/crypto"
 	"github.com/status-im/status-go/eth-node/types"
+	"github.com/status-im/status-go/protocol/push_notification_server"
 	"github.com/status-im/status-go/protocol/tt"
 	"github.com/status-im/status-go/whisper/v6"
 )
@@ -46,16 +48,10 @@ func (s *MessengerPushNotificationSuite) SetupTest() {
 	s.privateKey = s.m.identity
 }
 
-func (s *MessengerPushNotificationSuite) newMessengerWithKey(shh types.Whisper, privateKey *ecdsa.PrivateKey) *Messenger {
+func (s *MessengerPushNotificationSuite) newMessengerWithOptions(shh types.Whisper, privateKey *ecdsa.PrivateKey, options []Option) *Messenger {
 	tmpFile, err := ioutil.TempFile("", "")
 	s.Require().NoError(err)
 
-	options := []Option{
-		WithCustomLogger(s.logger),
-		WithMessagesPersistenceEnabled(),
-		WithDatabaseConfig(tmpFile.Name(), "some-key"),
-		WithDatasync(),
-	}
 	m, err := NewMessenger(
 		privateKey,
 		&testNode{shh: shh},
@@ -72,6 +68,19 @@ func (s *MessengerPushNotificationSuite) newMessengerWithKey(shh types.Whisper,
 	return m
 }
 
+func (s *MessengerPushNotificationSuite) newMessengerWithKey(shh types.Whisper, privateKey *ecdsa.PrivateKey) *Messenger {
+	tmpFile, err := ioutil.TempFile("", "")
+	s.Require().NoError(err)
+
+	options := []Option{
+		WithCustomLogger(s.logger),
+		WithMessagesPersistenceEnabled(),
+		WithDatabaseConfig(tmpFile.Name(), "some-key"),
+		WithDatasync(),
+	}
+	return s.newMessengerWithOptions(shh, privateKey, options)
+}
+
 func (s *MessengerPushNotificationSuite) newMessenger(shh types.Whisper) *Messenger {
 	privateKey, err := crypto.GenerateKey()
 	s.Require().NoError(err)
@@ -79,16 +88,57 @@ func (s *MessengerPushNotificationSuite) newMessenger(shh types.Whisper) *Messen
 	return s.newMessengerWithKey(s.shh, privateKey)
 }
 
+func (s *MessengerPushNotificationSuite) newPushNotificationServer(shh types.Whisper) *Messenger {
+	privateKey, err := crypto.GenerateKey()
+	s.Require().NoError(err)
+
+	tmpFile, err := ioutil.TempFile("", "")
+	s.Require().NoError(err)
+
+	serverConfig := &push_notification_server.Config{
+		Logger:   s.logger,
+		Identity: privateKey,
+	}
+
+	options := []Option{
+		WithCustomLogger(s.logger),
+		WithMessagesPersistenceEnabled(),
+		WithDatabaseConfig(tmpFile.Name(), "some-key"),
+		WithPushNotificationServerConfig(serverConfig),
+		WithDatasync(),
+	}
+	return s.newMessengerWithOptions(shh, privateKey, options)
+}
+
 func (s *MessengerPushNotificationSuite) TestReceivePushNotification() {
+	errChan := make(chan error)
+
 	deviceToken := "token"
 
-	server := s.newMessenger(s.shh)
+	server := s.newPushNotificationServer(s.shh)
 	client2 := s.newMessenger(s.shh)
 
 	err := s.m.AddPushNotificationServer(context.Background(), &server.identity.PublicKey)
 	s.Require().NoError(err)
 
-	err = s.m.RegisterForPushNotifications(context.Background(), deviceToken)
+	go func() {
+		err := s.m.RegisterForPushNotifications(context.Background(), deviceToken)
+		errChan <- err
+	}()
+
+	time.Sleep(500 * time.Millisecond)
+	_, err = server.RetrieveAll()
+	s.Require().NoError(err)
+
+	time.Sleep(500 * time.Millisecond)
+	_, err = server.RetrieveAll()
+	s.Require().NoError(err)
+
+	time.Sleep(500 * time.Millisecond)
+	_, err = server.RetrieveAll()
+	s.Require().NoError(err)
+
+	err = <-errChan
 	s.Require().NoError(err)
 
 	info, err := client2.pushNotificationClient.RetrievePushNotificationInfo(&s.m.identity.PublicKey)