github.com/status-im/status-go@v1.1.0/protocol/common/message_sender_test.go (about)

     1  package common
     2  
     3  import (
     4  	"math"
     5  	"testing"
     6  
     7  	transport2 "github.com/status-im/status-go/protocol/transport"
     8  	"github.com/status-im/status-go/t/helpers"
     9  
    10  	"github.com/status-im/status-go/waku"
    11  
    12  	"github.com/golang/protobuf/proto"
    13  
    14  	"github.com/stretchr/testify/suite"
    15  	"go.uber.org/zap"
    16  
    17  	datasyncproto "github.com/status-im/mvds/protobuf"
    18  
    19  	gethbridge "github.com/status-im/status-go/eth-node/bridge/geth"
    20  	"github.com/status-im/status-go/eth-node/crypto"
    21  	"github.com/status-im/status-go/eth-node/types"
    22  	"github.com/status-im/status-go/protocol/datasync"
    23  	"github.com/status-im/status-go/protocol/encryption"
    24  	"github.com/status-im/status-go/protocol/protobuf"
    25  	"github.com/status-im/status-go/protocol/sqlite"
    26  	v1protocol "github.com/status-im/status-go/protocol/v1"
    27  
    28  	"github.com/status-im/status-go/appdatabase"
    29  )
    30  
    31  func TestMessageSenderSuite(t *testing.T) {
    32  	suite.Run(t, new(MessageSenderSuite))
    33  }
    34  
    35  type MessageSenderSuite struct {
    36  	suite.Suite
    37  
    38  	sender      *MessageSender
    39  	testMessage protobuf.ChatMessage
    40  	logger      *zap.Logger
    41  }
    42  
    43  func (s *MessageSenderSuite) SetupTest() {
    44  	s.testMessage = protobuf.ChatMessage{
    45  		Text:        "abc123",
    46  		ChatId:      "testing-adamb",
    47  		ContentType: protobuf.ChatMessage_TEXT_PLAIN,
    48  		MessageType: protobuf.MessageType_PUBLIC_GROUP,
    49  		Clock:       154593077368201,
    50  		Timestamp:   1545930773682,
    51  	}
    52  
    53  	var err error
    54  
    55  	s.logger, err = zap.NewDevelopment()
    56  	s.Require().NoError(err)
    57  
    58  	identity, err := crypto.GenerateKey()
    59  	s.Require().NoError(err)
    60  
    61  	database, err := helpers.SetupTestMemorySQLDB(appdatabase.DbInitializer{})
    62  	s.Require().NoError(err)
    63  	err = sqlite.Migrate(database)
    64  	s.Require().NoError(err)
    65  
    66  	encryptionProtocol := encryption.New(
    67  		database,
    68  		"installation-1",
    69  		s.logger,
    70  	)
    71  
    72  	wakuConfig := waku.DefaultConfig
    73  	wakuConfig.MinimumAcceptedPoW = 0
    74  	shh := waku.New(&wakuConfig, s.logger)
    75  	s.Require().NoError(shh.Start())
    76  
    77  	whisperTransport, err := transport2.NewTransport(
    78  		gethbridge.NewGethWakuWrapper(shh),
    79  		identity,
    80  		database,
    81  		"waku_keys",
    82  		nil,
    83  		nil,
    84  		s.logger,
    85  	)
    86  	s.Require().NoError(err)
    87  
    88  	s.sender, err = NewMessageSender(
    89  		identity,
    90  		database,
    91  		encryptionProtocol,
    92  		whisperTransport,
    93  		s.logger,
    94  		FeatureFlags{
    95  			Datasync: true,
    96  		},
    97  	)
    98  	s.Require().NoError(err)
    99  }
   100  
   101  func (s *MessageSenderSuite) TearDownTest() {
   102  	_ = s.logger.Sync()
   103  }
   104  
   105  func (s *MessageSenderSuite) TestHandleDecodedMessagesWrapped() {
   106  	relayerKey, err := crypto.GenerateKey()
   107  	s.Require().NoError(err)
   108  
   109  	authorKey, err := crypto.GenerateKey()
   110  	s.Require().NoError(err)
   111  
   112  	encodedPayload, err := proto.Marshal(&s.testMessage)
   113  	s.Require().NoError(err)
   114  
   115  	wrappedPayload, err := v1protocol.WrapMessageV1(encodedPayload, protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, authorKey)
   116  	s.Require().NoError(err)
   117  
   118  	message := &types.Message{}
   119  	message.Sig = crypto.FromECDSAPub(&relayerKey.PublicKey)
   120  	message.Payload = wrappedPayload
   121  
   122  	response, err := s.sender.HandleMessages(message)
   123  	s.Require().NoError(err)
   124  	decodedMessages := response.StatusMessages
   125  
   126  	s.Require().Equal(1, len(decodedMessages))
   127  	s.Require().Equal(&authorKey.PublicKey, decodedMessages[0].SigPubKey())
   128  	s.Require().Equal(v1protocol.MessageID(&authorKey.PublicKey, wrappedPayload), decodedMessages[0].ApplicationLayer.ID)
   129  	s.Require().Equal(encodedPayload, decodedMessages[0].ApplicationLayer.Payload)
   130  	s.Require().Equal(protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, decodedMessages[0].ApplicationLayer.Type)
   131  }
   132  
   133  func (s *MessageSenderSuite) TestHandleDecodedMessagesDatasync() {
   134  	relayerKey, err := crypto.GenerateKey()
   135  	s.Require().NoError(err)
   136  
   137  	authorKey, err := crypto.GenerateKey()
   138  	s.Require().NoError(err)
   139  
   140  	encodedPayload, err := proto.Marshal(&s.testMessage)
   141  	s.Require().NoError(err)
   142  
   143  	wrappedPayload, err := v1protocol.WrapMessageV1(encodedPayload, protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, authorKey)
   144  	s.Require().NoError(err)
   145  
   146  	ds := datasync.New(nil, nil, false, s.sender.logger)
   147  	s.sender.datasync = ds
   148  
   149  	dataSyncMessage := datasyncproto.Payload{
   150  		Messages: []*datasyncproto.Message{
   151  			{Body: wrappedPayload},
   152  		},
   153  	}
   154  	marshalledDataSyncMessage, err := proto.Marshal(&dataSyncMessage)
   155  	s.Require().NoError(err)
   156  	message := &types.Message{}
   157  	message.Sig = crypto.FromECDSAPub(&relayerKey.PublicKey)
   158  	message.Payload = marshalledDataSyncMessage
   159  
   160  	response, err := s.sender.HandleMessages(message)
   161  	s.Require().NoError(err)
   162  	decodedMessages := response.StatusMessages
   163  
   164  	// We send two messages, the unwrapped one will be attributed to the relayer, while the wrapped one will be attributed to the author
   165  	s.Require().Equal(1, len(decodedMessages))
   166  	s.Require().Equal(&authorKey.PublicKey, decodedMessages[0].SigPubKey())
   167  	s.Require().Equal(v1protocol.MessageID(&authorKey.PublicKey, wrappedPayload), decodedMessages[0].ApplicationLayer.ID)
   168  	s.Require().Equal(encodedPayload, decodedMessages[0].ApplicationLayer.Payload)
   169  	s.Require().Equal(protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, decodedMessages[0].ApplicationLayer.Type)
   170  }
   171  
   172  func (s *MessageSenderSuite) CalculatePoWTest() {
   173  	largeSizePayload := make([]byte, largeSizeInBytes)
   174  	s.Require().Equal(whisperLargeSizePoW, calculatePoW(largeSizePayload))
   175  	normalSizePayload := make([]byte, largeSizeInBytes-1)
   176  	s.Require().Equal(whisperDefaultPoW, calculatePoW(normalSizePayload))
   177  
   178  }
   179  func (s *MessageSenderSuite) TestHandleDecodedMessagesDatasyncEncrypted() {
   180  	relayerKey, err := crypto.GenerateKey()
   181  	s.Require().NoError(err)
   182  
   183  	authorKey, err := crypto.GenerateKey()
   184  	s.Require().NoError(err)
   185  
   186  	encodedPayload, err := proto.Marshal(&s.testMessage)
   187  	s.Require().NoError(err)
   188  
   189  	wrappedPayload, err := v1protocol.WrapMessageV1(encodedPayload, protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, authorKey)
   190  	s.Require().NoError(err)
   191  
   192  	dataSyncMessage := datasyncproto.Payload{
   193  		Messages: []*datasyncproto.Message{
   194  			{Body: wrappedPayload},
   195  		},
   196  	}
   197  	marshalledDataSyncMessage, err := proto.Marshal(&dataSyncMessage)
   198  	s.Require().NoError(err)
   199  
   200  	// Create sender encryption protocol.
   201  	senderDatabase, err := helpers.SetupTestMemorySQLDB(appdatabase.DbInitializer{})
   202  	s.Require().NoError(err)
   203  	err = sqlite.Migrate(senderDatabase)
   204  	s.Require().NoError(err)
   205  
   206  	senderEncryptionProtocol := encryption.New(
   207  		senderDatabase,
   208  		"installation-2",
   209  		s.logger,
   210  	)
   211  
   212  	messageSpec, err := senderEncryptionProtocol.BuildEncryptedMessage(
   213  		relayerKey,
   214  		&s.sender.identity.PublicKey,
   215  		marshalledDataSyncMessage,
   216  	)
   217  	s.Require().NoError(err)
   218  
   219  	encryptedPayload, err := proto.Marshal(messageSpec.Message)
   220  	s.Require().NoError(err)
   221  
   222  	message := &types.Message{}
   223  	message.Sig = crypto.FromECDSAPub(&relayerKey.PublicKey)
   224  	message.Payload = encryptedPayload
   225  
   226  	ds := datasync.New(nil, nil, false, s.sender.logger)
   227  	s.sender.datasync = ds
   228  
   229  	response, err := s.sender.HandleMessages(message)
   230  	s.Require().NoError(err)
   231  	decodedMessages := response.StatusMessages
   232  
   233  	// We send two messages, the unwrapped one will be attributed to the relayer,
   234  	// while the wrapped one will be attributed to the author.
   235  	s.Require().Equal(1, len(decodedMessages))
   236  	s.Require().Equal(&authorKey.PublicKey, decodedMessages[0].SigPubKey())
   237  	s.Require().Equal(v1protocol.MessageID(&authorKey.PublicKey, wrappedPayload), decodedMessages[0].ApplicationLayer.ID)
   238  	s.Require().Equal(encodedPayload, decodedMessages[0].ApplicationLayer.Payload)
   239  	s.Require().Equal(protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, decodedMessages[0].ApplicationLayer.Type)
   240  }
   241  
   242  func (s *MessageSenderSuite) TestHandleOutOfOrderHashRatchet() {
   243  	groupID := []byte("group-id")
   244  	senderKey, err := crypto.GenerateKey()
   245  	s.Require().NoError(err)
   246  
   247  	encodedPayload, err := proto.Marshal(&s.testMessage)
   248  	s.Require().NoError(err)
   249  
   250  	// Create sender encryption protocol.
   251  	senderDatabase, err := helpers.SetupTestMemorySQLDB(appdatabase.DbInitializer{})
   252  	s.Require().NoError(err)
   253  	err = sqlite.Migrate(senderDatabase)
   254  	s.Require().NoError(err)
   255  
   256  	senderEncryptionProtocol := encryption.New(
   257  		senderDatabase,
   258  		"installation-2",
   259  		s.logger,
   260  	)
   261  
   262  	ratchet, err := senderEncryptionProtocol.GenerateHashRatchetKey(groupID)
   263  	s.Require().NoError(err)
   264  
   265  	ratchets := []*encryption.HashRatchetKeyCompatibility{ratchet}
   266  
   267  	hashRatchetKeyExchangeMessage, err := senderEncryptionProtocol.BuildHashRatchetKeyExchangeMessage(senderKey, &s.sender.identity.PublicKey, groupID, ratchets)
   268  	s.Require().NoError(err)
   269  
   270  	encryptedPayload1, err := proto.Marshal(hashRatchetKeyExchangeMessage.Message)
   271  	s.Require().NoError(err)
   272  
   273  	wrappedPayload2, err := v1protocol.WrapMessageV1(encodedPayload, protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, senderKey)
   274  	s.Require().NoError(err)
   275  
   276  	messageSpec2, err := senderEncryptionProtocol.BuildHashRatchetMessage(
   277  		groupID,
   278  		wrappedPayload2,
   279  	)
   280  	s.Require().NoError(err)
   281  
   282  	encryptedPayload2, err := proto.Marshal(messageSpec2.Message)
   283  	s.Require().NoError(err)
   284  
   285  	message := &types.Message{}
   286  	message.Sig = crypto.FromECDSAPub(&senderKey.PublicKey)
   287  	message.Hash = []byte{0x1}
   288  	message.Payload = encryptedPayload2
   289  
   290  	_, err = s.sender.HandleMessages(message)
   291  	s.Require().NoError(err)
   292  
   293  	keyID, err := ratchet.GetKeyID()
   294  	s.Require().NoError(err)
   295  
   296  	msgs, err := s.sender.persistence.GetHashRatchetMessages(keyID)
   297  	s.Require().NoError(err)
   298  
   299  	s.Require().Len(msgs, 1)
   300  
   301  	message = &types.Message{}
   302  	message.Sig = crypto.FromECDSAPub(&senderKey.PublicKey)
   303  	message.Hash = []byte{0x2}
   304  	message.Payload = encryptedPayload1
   305  
   306  	response, err := s.sender.HandleMessages(message)
   307  	s.Require().NoError(err)
   308  	decodedMessages2 := response.StatusMessages
   309  	s.Require().NotNil(decodedMessages2)
   310  
   311  	// It should have 2 messages, the key exchange and the one from the database
   312  	s.Require().Len(decodedMessages2, 2)
   313  
   314  	// it deletes the messages after being processed
   315  	msgs, err = s.sender.persistence.GetHashRatchetMessages(keyID)
   316  	s.Require().NoError(err)
   317  
   318  	s.Require().Len(msgs, 0)
   319  
   320  }
   321  
   322  func (s *MessageSenderSuite) TestHandleSegmentMessages() {
   323  	relayerKey, err := crypto.GenerateKey()
   324  	s.Require().NoError(err)
   325  
   326  	authorKey, err := crypto.GenerateKey()
   327  	s.Require().NoError(err)
   328  
   329  	encodedPayload, err := proto.Marshal(&s.testMessage)
   330  	s.Require().NoError(err)
   331  
   332  	wrappedPayload, err := v1protocol.WrapMessageV1(encodedPayload, protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, authorKey)
   333  	s.Require().NoError(err)
   334  
   335  	segmentedMessages, err := segmentMessage(&types.NewMessage{Payload: wrappedPayload}, int(math.Ceil(float64(len(wrappedPayload))/2)))
   336  	s.Require().NoError(err)
   337  	s.Require().Len(segmentedMessages, 2)
   338  
   339  	message := &types.Message{}
   340  	message.Sig = crypto.FromECDSAPub(&relayerKey.PublicKey)
   341  	message.Payload = segmentedMessages[0].Payload
   342  
   343  	// First segment is received, no messages are decoded
   344  	response, err := s.sender.HandleMessages(message)
   345  	s.Require().NoError(err)
   346  	s.Require().Nil(response)
   347  
   348  	// Second (and final) segment is received, reassembled message is decoded
   349  	message.Payload = segmentedMessages[1].Payload
   350  	response, err = s.sender.HandleMessages(message)
   351  	s.Require().NoError(err)
   352  
   353  	decodedMessages := response.StatusMessages
   354  	s.Require().Len(decodedMessages, 1)
   355  	s.Require().Equal(&authorKey.PublicKey, decodedMessages[0].SigPubKey())
   356  	s.Require().Equal(v1protocol.MessageID(&authorKey.PublicKey, wrappedPayload), decodedMessages[0].ApplicationLayer.ID)
   357  	s.Require().Equal(encodedPayload, decodedMessages[0].ApplicationLayer.Payload)
   358  	s.Require().Equal(protobuf.ApplicationMetadataMessage_CHAT_MESSAGE, decodedMessages[0].ApplicationLayer.Type)
   359  
   360  	// Receiving another segment after the message has been reassembled is considered an error
   361  	_, err = s.sender.HandleMessages(message)
   362  	s.Require().ErrorIs(err, ErrMessageSegmentsAlreadyCompleted)
   363  }
   364  
   365  func (s *MessageSenderSuite) TestGetEphemeralKey() {
   366  	keyMap := make(map[string]bool)
   367  	for i := 0; i < maxMessageSenderEphemeralKeys; i++ {
   368  		key, err := s.sender.GetEphemeralKey()
   369  		s.Require().NoError(err)
   370  		s.Require().NotNil(key)
   371  		keyMap[PubkeyToHex(&key.PublicKey)] = true
   372  	}
   373  	s.Require().Len(keyMap, maxMessageSenderEphemeralKeys)
   374  	// Add one more
   375  	key, err := s.sender.GetEphemeralKey()
   376  	s.Require().NoError(err)
   377  	s.Require().NotNil(key)
   378  
   379  	s.Require().True(keyMap[PubkeyToHex(&key.PublicKey)])
   380  }