status-go/protocol/encryption/persistence_test.go

package encryption

import (
	"reflect"
	"testing"

	"github.com/stretchr/testify/suite"

	"github.com/status-im/status-go/appdatabase"
	"github.com/status-im/status-go/eth-node/crypto"
	"github.com/status-im/status-go/t/helpers"

	"github.com/status-im/status-go/protocol/encryption/multidevice"
	"github.com/status-im/status-go/protocol/sqlite"
)

func TestSQLLitePersistenceTestSuite(t *testing.T) {
	suite.Run(t, new(SQLLitePersistenceTestSuite))
}

type SQLLitePersistenceTestSuite struct {
	suite.Suite
	service *sqlitePersistence
}

func (s *SQLLitePersistenceTestSuite) SetupTest() {
	db, err := helpers.SetupTestMemorySQLDB(appdatabase.DbInitializer{})
	s.Require().NoError(err)
	err = sqlite.Migrate(db)
	s.Require().NoError(err)

	s.service = newSQLitePersistence(db)
}

func (s *SQLLitePersistenceTestSuite) TestPrivateBundle() {
	installationID := "1"

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

	actualKey, err := s.service.GetPrivateKeyBundle([]byte("non-existing"))
	s.Require().NoError(err, "Error was not returned even though bundle is not there")
	s.Nil(actualKey)

	anyPrivateBundle, err := s.service.GetAnyPrivateBundle([]byte("non-existing-id"), []*multidevice.Installation{{ID: installationID, Version: 1}})
	s.Require().NoError(err)
	s.Nil(anyPrivateBundle)

	bundle, err := NewBundleContainer(key, installationID)
	s.Require().NoError(err)

	err = s.service.AddPrivateBundle(bundle)
	s.Require().NoError(err)

	bundleID := bundle.GetBundle().GetSignedPreKeys()[installationID].GetSignedPreKey()

	actualKey, err = s.service.GetPrivateKeyBundle(bundleID)
	s.Require().NoError(err)
	s.Equal(bundle.GetPrivateSignedPreKey(), actualKey, "It returns the same key")

	identity := crypto.CompressPubkey(&key.PublicKey)
	anyPrivateBundle, err = s.service.GetAnyPrivateBundle(identity, []*multidevice.Installation{{ID: installationID, Version: 1}})
	s.Require().NoError(err)
	s.NotNil(anyPrivateBundle)
	s.Equal(bundle.GetBundle().GetSignedPreKeys()[installationID].SignedPreKey, anyPrivateBundle.GetBundle().GetSignedPreKeys()[installationID].SignedPreKey, "It returns the same bundle")
}

func (s *SQLLitePersistenceTestSuite) TestPublicBundle() {
	key, err := crypto.GenerateKey()
	s.Require().NoError(err)

	actualBundle, err := s.service.GetPublicBundle(&key.PublicKey, []*multidevice.Installation{{ID: "1", Version: 1}})
	s.Require().NoError(err, "Error was not returned even though bundle is not there")
	s.Nil(actualBundle)

	bundleContainer, err := NewBundleContainer(key, "1")
	s.Require().NoError(err)

	bundle := bundleContainer.GetBundle()
	err = s.service.AddPublicBundle(bundle)
	s.Require().NoError(err)

	actualBundle, err = s.service.GetPublicBundle(&key.PublicKey, []*multidevice.Installation{{ID: "1", Version: 1}})
	s.Require().NoError(err)
	s.Equal(bundle.GetIdentity(), actualBundle.GetIdentity(), "It sets the right identity")
	s.Equal(bundle.GetSignedPreKeys(), actualBundle.GetSignedPreKeys(), "It sets the right prekeys")
}

func (s *SQLLitePersistenceTestSuite) TestUpdatedBundle() {
	key, err := crypto.GenerateKey()
	s.Require().NoError(err)

	actualBundle, err := s.service.GetPublicBundle(&key.PublicKey, []*multidevice.Installation{{ID: "1", Version: 1}})
	s.Require().NoError(err, "Error was not returned even though bundle is not there")
	s.Nil(actualBundle)

	// Create & add initial bundle
	bundleContainer, err := NewBundleContainer(key, "1")
	s.Require().NoError(err)

	bundle := bundleContainer.GetBundle()
	err = s.service.AddPublicBundle(bundle)
	s.Require().NoError(err)

	// Create & add a new bundle
	bundleContainer, err = NewBundleContainer(key, "1")
	s.Require().NoError(err)
	bundle = bundleContainer.GetBundle()
	// We set the version
	bundle.GetSignedPreKeys()["1"].Version = 1

	err = s.service.AddPublicBundle(bundle)
	s.Require().NoError(err)

	actualBundle, err = s.service.GetPublicBundle(&key.PublicKey, []*multidevice.Installation{{ID: "1", Version: 1}})
	s.Require().NoError(err)
	s.Equal(bundle.GetIdentity(), actualBundle.GetIdentity(), "It sets the right identity")
	s.Equal(bundle.GetSignedPreKeys(), actualBundle.GetSignedPreKeys(), "It sets the right prekeys")
}

func (s *SQLLitePersistenceTestSuite) TestOutOfOrderBundles() {
	key, err := crypto.GenerateKey()
	s.Require().NoError(err)

	actualBundle, err := s.service.GetPublicBundle(&key.PublicKey, []*multidevice.Installation{{ID: "1", Version: 1}})
	s.Require().NoError(err, "Error was not returned even though bundle is not there")
	s.Nil(actualBundle)

	// Create & add initial bundle
	bundleContainer, err := NewBundleContainer(key, "1")
	s.Require().NoError(err)

	bundle1 := bundleContainer.GetBundle()
	err = s.service.AddPublicBundle(bundle1)
	s.Require().NoError(err)

	// Create & add a new bundle
	bundleContainer, err = NewBundleContainer(key, "1")
	s.Require().NoError(err)

	bundle2 := bundleContainer.GetBundle()
	// We set the version
	bundle2.GetSignedPreKeys()["1"].Version = 1

	err = s.service.AddPublicBundle(bundle2)
	s.Require().NoError(err)

	// Add again the initial bundle
	err = s.service.AddPublicBundle(bundle1)
	s.Require().NoError(err)

	actualBundle, err = s.service.GetPublicBundle(&key.PublicKey, []*multidevice.Installation{{ID: "1", Version: 1}})
	s.Require().NoError(err)
	s.Equal(bundle2.GetIdentity(), actualBundle.GetIdentity(), "It sets the right identity")
	s.Equal(bundle2.GetSignedPreKeys()["1"].GetVersion(), uint32(1))
	s.Equal(bundle2.GetSignedPreKeys()["1"].GetSignedPreKey(), actualBundle.GetSignedPreKeys()["1"].GetSignedPreKey(), "It sets the right prekeys")
}

func (s *SQLLitePersistenceTestSuite) TestMultiplePublicBundle() {
	key, err := crypto.GenerateKey()
	s.Require().NoError(err)

	actualBundle, err := s.service.GetPublicBundle(&key.PublicKey, []*multidevice.Installation{{ID: "1", Version: 1}})
	s.Require().NoError(err, "Error was not returned even though bundle is not there")
	s.Nil(actualBundle)

	bundleContainer, err := NewBundleContainer(key, "1")
	s.Require().NoError(err)

	bundle := bundleContainer.GetBundle()
	err = s.service.AddPublicBundle(bundle)
	s.Require().NoError(err)

	// Adding it again does not throw an error
	err = s.service.AddPublicBundle(bundle)
	s.Require().NoError(err)

	// Adding a different bundle
	bundleContainer, err = NewBundleContainer(key, "1")
	s.Require().NoError(err)
	// We set the version
	bundle = bundleContainer.GetBundle()
	bundle.GetSignedPreKeys()["1"].Version = 1

	err = s.service.AddPublicBundle(bundle)
	s.Require().NoError(err)

	// Returns the most recent bundle
	actualBundle, err = s.service.GetPublicBundle(&key.PublicKey, []*multidevice.Installation{{ID: "1", Version: 1}})
	s.Require().NoError(err)

	s.Equal(bundle.GetIdentity(), actualBundle.GetIdentity(), "It sets the identity")
	s.Equal(bundle.GetSignedPreKeys(), actualBundle.GetSignedPreKeys(), "It sets the signed pre keys")

}

func (s *SQLLitePersistenceTestSuite) TestMultiDevicePublicBundle() {
	key, err := crypto.GenerateKey()
	s.Require().NoError(err)

	actualBundle, err := s.service.GetPublicBundle(&key.PublicKey, []*multidevice.Installation{{ID: "1", Version: 1}})
	s.Require().NoError(err, "Error was not returned even though bundle is not there")
	s.Nil(actualBundle)

	bundleContainer, err := NewBundleContainer(key, "1")
	s.Require().NoError(err)

	bundle := bundleContainer.GetBundle()
	err = s.service.AddPublicBundle(bundle)
	s.Require().NoError(err)

	// Adding it again does not throw an error
	err = s.service.AddPublicBundle(bundle)
	s.Require().NoError(err)

	// Adding a different bundle from a different instlation id
	bundleContainer, err = NewBundleContainer(key, "2")
	s.Require().NoError(err)

	bundle = bundleContainer.GetBundle()
	err = s.service.AddPublicBundle(bundle)
	s.Require().NoError(err)

	// Returns the most recent bundle
	actualBundle, err = s.service.GetPublicBundle(&key.PublicKey,
		[]*multidevice.Installation{
			{ID: "1", Version: 1},
			{ID: "2", Version: 1},
		})
	s.Require().NoError(err)

	s.Equal(bundle.GetIdentity(), actualBundle.GetIdentity(), "It sets the identity")
	s.NotNil(actualBundle.GetSignedPreKeys()["1"])
	s.NotNil(actualBundle.GetSignedPreKeys()["2"])
}

func (s *SQLLitePersistenceTestSuite) TestRatchetInfoPrivateBundle() {
	key, err := crypto.GenerateKey()
	s.Require().NoError(err)

	// Add a private bundle
	bundle, err := NewBundleContainer(key, "2")
	s.Require().NoError(err)

	err = s.service.AddPrivateBundle(bundle)
	s.Require().NoError(err)

	err = s.service.AddRatchetInfo(
		[]byte("symmetric-key"),
		[]byte("their-public-key"),
		bundle.GetBundle().GetSignedPreKeys()["2"].GetSignedPreKey(),
		[]byte("ephemeral-public-key"),
		"1",
	)
	s.Require().NoError(err)

	ratchetInfo, err := s.service.GetRatchetInfo(bundle.GetBundle().GetSignedPreKeys()["2"].GetSignedPreKey(), []byte("their-public-key"), "1")

	s.Require().NoError(err)
	s.Require().NotNil(ratchetInfo)
	s.NotNil(ratchetInfo.ID, "It adds an id")
	s.Equal(ratchetInfo.PrivateKey, bundle.GetPrivateSignedPreKey(), "It returns the private key")
	s.Equal(ratchetInfo.Sk, []byte("symmetric-key"), "It returns the symmetric key")
	s.Equal(ratchetInfo.Identity, []byte("their-public-key"), "It returns the identity of the contact")
	s.Equal(ratchetInfo.PublicKey, bundle.GetBundle().GetSignedPreKeys()["2"].GetSignedPreKey(), "It  returns the public key of the bundle")
	s.Equal(bundle.GetBundle().GetSignedPreKeys()["2"].GetSignedPreKey(), ratchetInfo.BundleID, "It returns the bundle id")
	s.Equal([]byte("ephemeral-public-key"), ratchetInfo.EphemeralKey, "It returns the ratchet ephemeral key")
	s.Equal("1", ratchetInfo.InstallationID, "It returns the right installation id")
}

func (s *SQLLitePersistenceTestSuite) TestRatchetInfoPublicBundle() {
	installationID := "1"
	theirPublicKey := []byte("their-public-key")
	key, err := crypto.GenerateKey()
	s.Require().NoError(err)

	// Add a private bundle
	bundle, err := NewBundleContainer(key, installationID)
	s.Require().NoError(err)

	err = s.service.AddPublicBundle(bundle.GetBundle())
	s.Require().NoError(err)

	signedPreKey := bundle.GetBundle().GetSignedPreKeys()[installationID].GetSignedPreKey()

	err = s.service.AddRatchetInfo(
		[]byte("symmetric-key"),
		theirPublicKey,
		signedPreKey,
		[]byte("public-ephemeral-key"),
		installationID,
	)
	s.Require().NoError(err)

	ratchetInfo, err := s.service.GetRatchetInfo(signedPreKey, theirPublicKey, installationID)

	s.Require().NoError(err)
	s.Require().NotNil(ratchetInfo, "It returns the ratchet info")

	s.NotNil(ratchetInfo.ID, "It adds an id")
	s.Nil(ratchetInfo.PrivateKey, "It does not return the private key")
	s.Equal(ratchetInfo.Sk, []byte("symmetric-key"), "It returns the symmetric key")
	s.Equal(ratchetInfo.Identity, theirPublicKey, "It returns the identity of the contact")
	s.Equal(ratchetInfo.PublicKey, signedPreKey, "It  returns the public key of the bundle")
	s.Equal(installationID, ratchetInfo.InstallationID, "It returns the right installationID")
	s.Nilf(ratchetInfo.PrivateKey, "It does not return the private key")

	ratchetInfo, err = s.service.GetAnyRatchetInfo(theirPublicKey, installationID)
	s.Require().NoError(err)
	s.Require().NotNil(ratchetInfo, "It returns the ratchet info")
	s.NotNil(ratchetInfo.ID, "It adds an id")
	s.Nil(ratchetInfo.PrivateKey, "It does not return the private key")
	s.Equal(ratchetInfo.Sk, []byte("symmetric-key"), "It returns the symmetric key")
	s.Equal(ratchetInfo.Identity, theirPublicKey, "It returns the identity of the contact")
	s.Equal(ratchetInfo.PublicKey, signedPreKey, "It  returns the public key of the bundle")
	s.Equal(signedPreKey, ratchetInfo.BundleID, "It returns the bundle id")
	s.Equal(installationID, ratchetInfo.InstallationID, "It saves the right installation ID")
}

func (s *SQLLitePersistenceTestSuite) TestRatchetInfoNoBundle() {
	err := s.service.AddRatchetInfo(
		[]byte("symmetric-key"),
		[]byte("their-public-key"),
		[]byte("non-existing-bundle"),
		[]byte("non-existing-ephemeral-key"),
		"none",
	)

	s.Error(err, "It returns an error")

	_, err = s.service.GetRatchetInfo([]byte("non-existing-bundle"), []byte("their-public-key"), "none")
	s.Require().NoError(err)

	ratchetInfo, err := s.service.GetAnyRatchetInfo([]byte("their-public-key"), "4")
	s.Require().NoError(err)
	s.Nil(ratchetInfo, "It returns nil when no bundle is there")
}

// TODO: Add test for MarkBundleExpired

func (s *SQLLitePersistenceTestSuite) TestGetHashRatchetKeyByID() {
	key := &HashRatchetKeyCompatibility{
		GroupID:   []byte{1, 2, 3},
		keyID:     []byte{4, 5, 6},
		Timestamp: 1,
		Key:       []byte{7, 8, 9},
	}
	err := s.service.SaveHashRatchetKey(key)
	s.Require().NoError(err)

	retrievedKey, err := s.service.GetHashRatchetKeyByID(key.keyID)
	s.Require().NoError(err)
	s.Require().True(reflect.DeepEqual(key.GroupID, retrievedKey.GroupID))
	s.Require().True(reflect.DeepEqual(key.keyID, retrievedKey.keyID))
	s.Require().True(reflect.DeepEqual(key.Key, retrievedKey.Key))
	s.Require().Equal(key.Timestamp, retrievedKey.Timestamp)

	cachedKey, err := s.service.GetHashRatchetCache(retrievedKey, 0)
	s.Require().NoError(err)
	s.Require().True(reflect.DeepEqual(key.keyID, cachedKey.KeyID))
	s.Require().True(reflect.DeepEqual(key.Key, cachedKey.Key))
	s.Require().EqualValues(0, cachedKey.SeqNo)

	var newSeqNo uint32 = 1
	newHash := []byte{10, 11, 12}
	err = s.service.SaveHashRatchetKeyHash(key, newHash, newSeqNo)
	s.Require().NoError(err)

	cachedKey, err = s.service.GetHashRatchetCache(retrievedKey, 0)
	s.Require().NoError(err)
	s.Require().True(reflect.DeepEqual(key.keyID, cachedKey.KeyID))
	s.Require().True(reflect.DeepEqual(key.Key, cachedKey.Key))
	s.Require().EqualValues(1, cachedKey.SeqNo)

	newSeqNo = 4
	newHash = []byte{10, 11, 13}
	err = s.service.SaveHashRatchetKeyHash(key, newHash, newSeqNo)
	s.Require().NoError(err)

	cachedKey, err = s.service.GetHashRatchetCache(retrievedKey, 0)
	s.Require().NoError(err)
	s.Require().True(reflect.DeepEqual(key.keyID, cachedKey.KeyID))
	s.Require().True(reflect.DeepEqual(key.Key, cachedKey.Key))
	s.Require().EqualValues(4, cachedKey.SeqNo)

	cachedKey, err = s.service.GetHashRatchetCache(retrievedKey, 1)
	s.Require().NoError(err)
	s.Require().True(reflect.DeepEqual(key.keyID, cachedKey.KeyID))
	s.Require().True(reflect.DeepEqual(key.Key, cachedKey.Key))
	s.Require().EqualValues(1, cachedKey.SeqNo)
}