status-go/protocol/common/message_sender_test.go

package common

import (
	"math"
	"testing"

	transport2 "github.com/status-im/status-go/protocol/transport"
	"github.com/status-im/status-go/t/helpers"

	"github.com/status-im/status-go/waku"

	"github.com/golang/protobuf/proto"

	"github.com/stretchr/testify/suite"
	"go.uber.org/zap"

	datasyncproto "github.com/status-im/mvds/protobuf"

	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/datasync"
	"github.com/status-im/status-go/protocol/encryption"
	"github.com/status-im/status-go/protocol/protobuf"
	"github.com/status-im/status-go/protocol/sqlite"
	v1protocol "github.com/status-im/status-go/protocol/v1"

	"github.com/status-im/status-go/appdatabase"
)

func TestMessageSenderSuite(t *testing.T) {
	suite.Run(t, new(MessageSenderSuite))
}

type MessageSenderSuite struct {
	suite.Suite

	sender      *MessageSender
	testMessage protobuf.ChatMessage
	logger      *zap.Logger
}

func (s *MessageSenderSuite) SetupTest() {
	s.testMessage = protobuf.ChatMessage{
		Text:        "abc123",
		ChatId:      "testing-adamb",
		ContentType: protobuf.ChatMessage_TEXT_PLAIN,
		MessageType: protobuf.MessageType_PUBLIC_GROUP,
		Clock:       154593077368201,
		Timestamp:   1545930773682,
	}

	var err error

	s.logger, err = zap.NewDevelopment()
	s.Require().NoError(err)

	identity, err := crypto.GenerateKey()
	s.Require().NoError(err)

	database, err := helpers.SetupTestMemorySQLDB(appdatabase.DbInitializer{})
	s.Require().NoError(err)
	err = sqlite.Migrate(database)
	s.Require().NoError(err)

	encryptionProtocol := encryption.New(
		database,
		"installation-1",
		s.logger,
	)

	wakuConfig := waku.DefaultConfig
	wakuConfig.MinimumAcceptedPoW = 0
	shh := waku.New(&wakuConfig, s.logger)
	s.Require().NoError(shh.Start())

	whisperTransport, err := transport2.NewTransport(
		gethbridge.NewGethWakuWrapper(shh),
		identity,
		database,
		"waku_keys",
		nil,
		nil,
		s.logger,
	)
	s.Require().NoError(err)

	s.sender, err = NewMessageSender(
		identity,
		database,
		encryptionProtocol,
		whisperTransport,
		s.logger,
		FeatureFlags{
			Datasync: true,
		},
	)
	s.Require().NoError(err)
}

func (s *MessageSenderSuite) TearDownTest() {
	_ = s.logger.Sync()
}

func (s *MessageSenderSuite) TestHandleDecodedMessagesWrapped() {
	relayerKey, err := crypto.GenerateKey()
	s.Require().NoError(err)

	authorKey, err := crypto.GenerateKey()
	s.Require().NoError(err)

	encodedPayload, err := proto.Marshal(&s.testMessage)
	s.Require().NoError(err)

	wrappedPayload, err := v1protocol.WrapMessageV1(encodedPayload, protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, authorKey)
	s.Require().NoError(err)

	message := &types.Message{}
	message.Sig = crypto.FromECDSAPub(&relayerKey.PublicKey)
	message.Payload = wrappedPayload

	response, err := s.sender.HandleMessages(message)
	s.Require().NoError(err)
	decodedMessages := response.StatusMessages

	s.Require().Equal(1, len(decodedMessages))
	s.Require().Equal(&authorKey.PublicKey, decodedMessages[0].SigPubKey())
	s.Require().Equal(v1protocol.MessageID(&authorKey.PublicKey, wrappedPayload), decodedMessages[0].ApplicationLayer.ID)
	s.Require().Equal(encodedPayload, decodedMessages[0].ApplicationLayer.Payload)
	s.Require().Equal(protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, decodedMessages[0].ApplicationLayer.Type)
}

func (s *MessageSenderSuite) TestHandleDecodedMessagesDatasync() {
	relayerKey, err := crypto.GenerateKey()
	s.Require().NoError(err)

	authorKey, err := crypto.GenerateKey()
	s.Require().NoError(err)

	encodedPayload, err := proto.Marshal(&s.testMessage)
	s.Require().NoError(err)

	wrappedPayload, err := v1protocol.WrapMessageV1(encodedPayload, protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, authorKey)
	s.Require().NoError(err)

	ds := datasync.New(nil, nil, false, s.sender.logger)
	s.sender.datasync = ds

	dataSyncMessage := datasyncproto.Payload{
		Messages: []*datasyncproto.Message{
			{Body: wrappedPayload},
		},
	}
	marshalledDataSyncMessage, err := proto.Marshal(&dataSyncMessage)
	s.Require().NoError(err)
	message := &types.Message{}
	message.Sig = crypto.FromECDSAPub(&relayerKey.PublicKey)
	message.Payload = marshalledDataSyncMessage

	response, err := s.sender.HandleMessages(message)
	s.Require().NoError(err)
	decodedMessages := response.StatusMessages

	// We send two messages, the unwrapped one will be attributed to the relayer, while the wrapped one will be attributed to the author
	s.Require().Equal(1, len(decodedMessages))
	s.Require().Equal(&authorKey.PublicKey, decodedMessages[0].SigPubKey())
	s.Require().Equal(v1protocol.MessageID(&authorKey.PublicKey, wrappedPayload), decodedMessages[0].ApplicationLayer.ID)
	s.Require().Equal(encodedPayload, decodedMessages[0].ApplicationLayer.Payload)
	s.Require().Equal(protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, decodedMessages[0].ApplicationLayer.Type)
}

func (s *MessageSenderSuite) CalculatePoWTest() {
	largeSizePayload := make([]byte, largeSizeInBytes)
	s.Require().Equal(whisperLargeSizePoW, calculatePoW(largeSizePayload))
	normalSizePayload := make([]byte, largeSizeInBytes-1)
	s.Require().Equal(whisperDefaultPoW, calculatePoW(normalSizePayload))

}
func (s *MessageSenderSuite) TestHandleDecodedMessagesDatasyncEncrypted() {
	relayerKey, err := crypto.GenerateKey()
	s.Require().NoError(err)

	authorKey, err := crypto.GenerateKey()
	s.Require().NoError(err)

	encodedPayload, err := proto.Marshal(&s.testMessage)
	s.Require().NoError(err)

	wrappedPayload, err := v1protocol.WrapMessageV1(encodedPayload, protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, authorKey)
	s.Require().NoError(err)

	dataSyncMessage := datasyncproto.Payload{
		Messages: []*datasyncproto.Message{
			{Body: wrappedPayload},
		},
	}
	marshalledDataSyncMessage, err := proto.Marshal(&dataSyncMessage)
	s.Require().NoError(err)

	// Create sender encryption protocol.
	senderDatabase, err := helpers.SetupTestMemorySQLDB(appdatabase.DbInitializer{})
	s.Require().NoError(err)
	err = sqlite.Migrate(senderDatabase)
	s.Require().NoError(err)

	senderEncryptionProtocol := encryption.New(
		senderDatabase,
		"installation-2",
		s.logger,
	)

	messageSpec, err := senderEncryptionProtocol.BuildEncryptedMessage(
		relayerKey,
		&s.sender.identity.PublicKey,
		marshalledDataSyncMessage,
	)
	s.Require().NoError(err)

	encryptedPayload, err := proto.Marshal(messageSpec.Message)
	s.Require().NoError(err)

	message := &types.Message{}
	message.Sig = crypto.FromECDSAPub(&relayerKey.PublicKey)
	message.Payload = encryptedPayload

	ds := datasync.New(nil, nil, false, s.sender.logger)
	s.sender.datasync = ds

	response, err := s.sender.HandleMessages(message)
	s.Require().NoError(err)
	decodedMessages := response.StatusMessages

	// We send two messages, the unwrapped one will be attributed to the relayer,
	// while the wrapped one will be attributed to the author.
	s.Require().Equal(1, len(decodedMessages))
	s.Require().Equal(&authorKey.PublicKey, decodedMessages[0].SigPubKey())
	s.Require().Equal(v1protocol.MessageID(&authorKey.PublicKey, wrappedPayload), decodedMessages[0].ApplicationLayer.ID)
	s.Require().Equal(encodedPayload, decodedMessages[0].ApplicationLayer.Payload)
	s.Require().Equal(protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, decodedMessages[0].ApplicationLayer.Type)
}

func (s *MessageSenderSuite) TestHandleOutOfOrderHashRatchet() {
	groupID := []byte("group-id")
	senderKey, err := crypto.GenerateKey()
	s.Require().NoError(err)

	encodedPayload, err := proto.Marshal(&s.testMessage)
	s.Require().NoError(err)

	// Create sender encryption protocol.
	senderDatabase, err := helpers.SetupTestMemorySQLDB(appdatabase.DbInitializer{})
	s.Require().NoError(err)
	err = sqlite.Migrate(senderDatabase)
	s.Require().NoError(err)

	senderEncryptionProtocol := encryption.New(
		senderDatabase,
		"installation-2",
		s.logger,
	)

	ratchet, err := senderEncryptionProtocol.GenerateHashRatchetKey(groupID)
	s.Require().NoError(err)

	ratchets := []*encryption.HashRatchetKeyCompatibility{ratchet}

	hashRatchetKeyExchangeMessage, err := senderEncryptionProtocol.BuildHashRatchetKeyExchangeMessage(senderKey, &s.sender.identity.PublicKey, groupID, ratchets)
	s.Require().NoError(err)

	encryptedPayload1, err := proto.Marshal(hashRatchetKeyExchangeMessage.Message)
	s.Require().NoError(err)

	wrappedPayload2, err := v1protocol.WrapMessageV1(encodedPayload, protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, senderKey)
	s.Require().NoError(err)

	messageSpec2, err := senderEncryptionProtocol.BuildHashRatchetMessage(
		groupID,
		wrappedPayload2,
	)
	s.Require().NoError(err)

	encryptedPayload2, err := proto.Marshal(messageSpec2.Message)
	s.Require().NoError(err)

	message := &types.Message{}
	message.Sig = crypto.FromECDSAPub(&senderKey.PublicKey)
	message.Hash = []byte{0x1}
	message.Payload = encryptedPayload2

	_, err = s.sender.HandleMessages(message)
	s.Require().NoError(err)

	keyID, err := ratchet.GetKeyID()
	s.Require().NoError(err)

	msgs, err := s.sender.persistence.GetHashRatchetMessages(keyID)
	s.Require().NoError(err)

	s.Require().Len(msgs, 1)

	message = &types.Message{}
	message.Sig = crypto.FromECDSAPub(&senderKey.PublicKey)
	message.Hash = []byte{0x2}
	message.Payload = encryptedPayload1

	response, err := s.sender.HandleMessages(message)
	s.Require().NoError(err)
	decodedMessages2 := response.StatusMessages
	s.Require().NotNil(decodedMessages2)

	// It should have 2 messages, the key exchange and the one from the database
	s.Require().Len(decodedMessages2, 2)

	// it deletes the messages after being processed
	msgs, err = s.sender.persistence.GetHashRatchetMessages(keyID)
	s.Require().NoError(err)

	s.Require().Len(msgs, 0)

}

func (s *MessageSenderSuite) TestHandleSegmentMessages() {
	relayerKey, err := crypto.GenerateKey()
	s.Require().NoError(err)

	authorKey, err := crypto.GenerateKey()
	s.Require().NoError(err)

	encodedPayload, err := proto.Marshal(&s.testMessage)
	s.Require().NoError(err)

	wrappedPayload, err := v1protocol.WrapMessageV1(encodedPayload, protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, authorKey)
	s.Require().NoError(err)

	segmentedMessages, err := segmentMessage(&types.NewMessage{Payload: wrappedPayload}, int(math.Ceil(float64(len(wrappedPayload))/2)))
	s.Require().NoError(err)
	s.Require().Len(segmentedMessages, 2)

	message := &types.Message{}
	message.Sig = crypto.FromECDSAPub(&relayerKey.PublicKey)
	message.Payload = segmentedMessages[0].Payload

	// First segment is received, no messages are decoded
	response, err := s.sender.HandleMessages(message)
	s.Require().NoError(err)
	s.Require().Nil(response)

	// Second (and final) segment is received, reassembled message is decoded
	message.Payload = segmentedMessages[1].Payload
	response, err = s.sender.HandleMessages(message)
	s.Require().NoError(err)

	decodedMessages := response.StatusMessages
	s.Require().Len(decodedMessages, 1)
	s.Require().Equal(&authorKey.PublicKey, decodedMessages[0].SigPubKey())
	s.Require().Equal(v1protocol.MessageID(&authorKey.PublicKey, wrappedPayload), decodedMessages[0].ApplicationLayer.ID)
	s.Require().Equal(encodedPayload, decodedMessages[0].ApplicationLayer.Payload)
	s.Require().Equal(protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, decodedMessages[0].ApplicationLayer.Type)

	// Receiving another segment after the message has been reassembled is considered an error
	_, 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)])
}