status-go/services/shhext/chat/protocol_test.go

117 lines
2.9 KiB
Go

package chat
import (
"os"
"testing"
"github.com/ethereum/go-ethereum/crypto"
"github.com/golang/protobuf/proto"
"github.com/stretchr/testify/suite"
)
func TestProtocolServiceTestSuite(t *testing.T) {
suite.Run(t, new(ProtocolServiceTestSuite))
}
type ProtocolServiceTestSuite struct {
suite.Suite
alice *ProtocolService
bob *ProtocolService
}
func (s *ProtocolServiceTestSuite) SetupTest() {
aliceDBPath := "/tmp/alice.db"
aliceDBKey := "alice"
bobDBPath := "/tmp/bob.db"
bobDBKey := "bob"
os.Remove(aliceDBPath)
os.Remove(bobDBPath)
alicePersistence, err := NewSQLLitePersistence(aliceDBPath, aliceDBKey)
if err != nil {
panic(err)
}
bobPersistence, err := NewSQLLitePersistence(bobDBPath, bobDBKey)
if err != nil {
panic(err)
}
s.alice = NewProtocolService(NewEncryptionService(alicePersistence, "1"))
s.bob = NewProtocolService(NewEncryptionService(bobPersistence, "2"))
}
func (s *ProtocolServiceTestSuite) TestBuildDirectMessage() {
bobKey, err := crypto.GenerateKey()
s.NoError(err)
aliceKey, err := crypto.GenerateKey()
s.NoError(err)
payload, err := proto.Marshal(&ChatMessagePayload{
Content: "Test content",
ClockValue: 1,
ContentType: "a",
MessageType: "some type",
})
s.NoError(err)
marshaledMsg, err := s.alice.BuildDirectMessage(aliceKey, payload, &bobKey.PublicKey, &aliceKey.PublicKey)
s.NoError(err)
s.NotNil(marshaledMsg, "It creates a message")
s.NotNil(marshaledMsg[&aliceKey.PublicKey], "It creates a single message")
unmarshaledMsg := &ProtocolMessage{}
err = proto.Unmarshal(marshaledMsg[&bobKey.PublicKey], unmarshaledMsg)
s.NoError(err)
s.NotNilf(unmarshaledMsg.GetBundle(), "It adds a bundle to the message")
directMessage := unmarshaledMsg.GetDirectMessage()
s.NotNilf(directMessage, "It sets the direct message")
encryptedPayload := directMessage["none"].GetPayload()
s.NotNilf(encryptedPayload, "It sets the payload of the message")
s.NotEqualf(payload, encryptedPayload, "It encrypts the payload")
}
func (s *ProtocolServiceTestSuite) TestBuildAndReadDirectMessage() {
bobKey, err := crypto.GenerateKey()
s.NoError(err)
aliceKey, err := crypto.GenerateKey()
s.NoError(err)
payload := ChatMessagePayload{
Content: "Test content",
ClockValue: 1,
ContentType: "a",
MessageType: "some type",
}
marshaledPayload, err := proto.Marshal(&payload)
s.NoError(err)
// Message is sent with DH
marshaledMsg, err := s.alice.BuildDirectMessage(aliceKey, marshaledPayload, &bobKey.PublicKey)
s.NoError(err)
// Bob is able to decrypt the message
response, err := s.bob.HandleMessage(bobKey, &aliceKey.PublicKey, marshaledMsg[&bobKey.PublicKey])
s.NoError(err)
unmarshaledMsg := response.Message
s.NotNil(unmarshaledMsg)
recoveredPayload := ChatMessagePayload{}
err = proto.Unmarshal(unmarshaledMsg, &recoveredPayload)
s.NoError(err)
s.Equalf(proto.Equal(&payload, &recoveredPayload), true, "It successfully unmarshal the decrypted message")
}