telegabber/e2ee/omemo/libsignal/memstore_test.go

package libsignal_test

import (
	"bytes"
	"fmt"
	"strconv"
	"strings"
	"sync"

	"dev.narayana.im/narayana/telegabber/e2ee/omemo/libsignal"
)

// memStore is a minimal in-memory libsignal.Store used only to give the
// round-trip tests in this package somewhere to persist state during a
// session. It is not the real Store implementation - that's badgerstore,
// which has its own dedicated tests.
type memStore struct {
	mu sync.Mutex

	identityPublic  []byte
	identityPrivate []byte
	registrationID  uint32

	preKeys          map[uint32][]byte
	signedPreKeys    map[uint32][]byte
	remoteIdentities map[string][]byte
	sessions         map[string][]byte
}

func newMemStore(public, private []byte, registrationID uint32) *memStore {
	return &memStore{
		identityPublic:   public,
		identityPrivate:  private,
		registrationID:   registrationID,
		preKeys:          map[uint32][]byte{},
		signedPreKeys:    map[uint32][]byte{},
		remoteIdentities: map[string][]byte{},
		sessions:         map[string][]byte{},
	}
}

func addrKey(addr libsignal.Address) string {
	return fmt.Sprintf("%s:%d", addr.Name, addr.DeviceID)
}

func (m *memStore) GetIdentityKeyPair() ([]byte, []byte, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.identityPublic, m.identityPrivate, nil
}

func (m *memStore) GetLocalRegistrationID() (uint32, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.registrationID, nil
}

func (m *memStore) SaveIdentity(addr libsignal.Address, key []byte) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if key == nil {
		delete(m.remoteIdentities, addrKey(addr))
		return nil
	}
	m.remoteIdentities[addrKey(addr)] = append([]byte{}, key...)
	return nil
}

func (m *memStore) IsTrustedIdentity(addr libsignal.Address, key []byte) (bool, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	stored, ok := m.remoteIdentities[addrKey(addr)]
	if !ok {
		return true, nil
	}
	return bytes.Equal(stored, key), nil
}

func (m *memStore) LoadPreKey(id uint32) ([]byte, bool, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	v, ok := m.preKeys[id]
	return v, ok, nil
}

func (m *memStore) StorePreKey(id uint32, record []byte) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.preKeys[id] = append([]byte{}, record...)
	return nil
}

func (m *memStore) ContainsPreKey(id uint32) bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	_, ok := m.preKeys[id]
	return ok
}

func (m *memStore) RemovePreKey(id uint32) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.preKeys, id)
	return nil
}

func (m *memStore) LoadSignedPreKey(id uint32) ([]byte, bool, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	v, ok := m.signedPreKeys[id]
	return v, ok, nil
}

func (m *memStore) StoreSignedPreKey(id uint32, record []byte) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.signedPreKeys[id] = append([]byte{}, record...)
	return nil
}

func (m *memStore) ContainsSignedPreKey(id uint32) bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	_, ok := m.signedPreKeys[id]
	return ok
}

func (m *memStore) RemoveSignedPreKey(id uint32) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.signedPreKeys, id)
	return nil
}

func (m *memStore) LoadSession(addr libsignal.Address) ([]byte, bool, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	v, ok := m.sessions[addrKey(addr)]
	return v, ok, nil
}

func (m *memStore) GetSubDeviceSessions(name string) ([]uint32, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	var ids []uint32
	prefix := name + ":"
	for k := range m.sessions {
		if strings.HasPrefix(k, prefix) {
			if id, err := strconv.ParseUint(strings.TrimPrefix(k, prefix), 10, 32); err == nil {
				ids = append(ids, uint32(id))
			}
		}
	}
	return ids, nil
}

func (m *memStore) StoreSession(addr libsignal.Address, record []byte) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.sessions[addrKey(addr)] = append([]byte{}, record...)
	return nil
}

func (m *memStore) ContainsSession(addr libsignal.Address) bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	_, ok := m.sessions[addrKey(addr)]
	return ok
}

func (m *memStore) DeleteSession(addr libsignal.Address) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.sessions, addrKey(addr))
	return nil
}

func (m *memStore) DeleteAllSessions(name string) (int, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	prefix := name + ":"
	count := 0
	for k := range m.sessions {
		if strings.HasPrefix(k, prefix) {
			delete(m.sessions, k)
			count++
		}
	}
	return count, nil
}

var _ libsignal.Store = (*memStore)(nil)