telegabber/e2ee/store/badgerstore/store.go

package badgerstore

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"strconv"

	badger "github.com/dgraph-io/badger/v4"
	log "github.com/sirupsen/logrus"

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

// Store implements libsignal.Store for one bridged chat's OMEMO identity,
// scoped by (login, ownerPeerJID) - see db.go's key layout note. Obtain one
// via DB.Store.
type Store struct {
	db    *badger.DB
	ctx   *libsignal.Context
	login string
	owner string // the bridged chat's pseudo-JID; this identity's own
}

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

// Key layout (mirrors badger/ids.go's prefix-then-type-then-suffix
// convention):
//
//	<login>/<owner>/enabled
//	<login>/<owner>/negotiatedversion
//	<login>/<owner>/override
//	<login>/<owner>/identity/keypair
//	<login>/<owner>/identity/registrationid
//	<login>/<owner>/identity/signedprekey/<id>
//	<login>/<owner>/identity/prekey/<id>
//	<login>/<owner>/identity/prekey/counter
//	<login>/<owner>/remote/<remote>/identity/<deviceId>
//	<login>/<owner>/remote/<remote>/blocked/<deviceId>
//	<login>/<owner>/remote/<remote>/session/<deviceId>
//	<login>/<owner>/remote/<remote>/devicelist
//
// (identity/<deviceId> doubles as the TOFU trust record - see
// IsTrustedIdentity; there's no separate "trust" key.)

func (s *Store) enabledKey() []byte {
	return []byte(fmt.Sprintf("%s/%s/enabled", s.login, s.owner))
}

func (s *Store) negotiatedVersionKey() []byte {
	return []byte(fmt.Sprintf("%s/%s/negotiatedversion", s.login, s.owner))
}

func (s *Store) overrideKey() []byte {
	return []byte(fmt.Sprintf("%s/%s/override", s.login, s.owner))
}

func (s *Store) identityKeyPairKey() []byte {
	return []byte(fmt.Sprintf("%s/%s/identity/keypair", s.login, s.owner))
}

func (s *Store) registrationIDKey() []byte {
	return []byte(fmt.Sprintf("%s/%s/identity/registrationid", s.login, s.owner))
}

func (s *Store) signedPreKeyKey(id uint32) []byte {
	return []byte(fmt.Sprintf("%s/%s/identity/signedprekey/%d", s.login, s.owner, id))
}

func (s *Store) preKeyKey(id uint32) []byte {
	return []byte(fmt.Sprintf("%s/%s/identity/prekey/%d", s.login, s.owner, id))
}

func (s *Store) preKeyPrefix() []byte {
	return []byte(fmt.Sprintf("%s/%s/identity/prekey/", s.login, s.owner))
}

func (s *Store) preKeyCounterKey() []byte {
	return []byte(fmt.Sprintf("%s/%s/identity/prekey/counter", s.login, s.owner))
}

func (s *Store) remoteIdentityKey(remote string, deviceID uint32) []byte {
	return []byte(fmt.Sprintf("%s/%s/remote/%s/identity/%d", s.login, s.owner, remote, deviceID))
}

func (s *Store) remoteSessionKey(remote string, deviceID uint32) []byte {
	return []byte(fmt.Sprintf("%s/%s/remote/%s/session/%d", s.login, s.owner, remote, deviceID))
}

func (s *Store) remoteSessionPrefix(remote string) []byte {
	return []byte(fmt.Sprintf("%s/%s/remote/%s/session/", s.login, s.owner, remote))
}

func (s *Store) remoteDeviceListKey(remote string) []byte {
	return []byte(fmt.Sprintf("%s/%s/remote/%s/devicelist", s.login, s.owner, remote))
}

func (s *Store) blockedKey(remote string, deviceID uint32) []byte {
	return []byte(fmt.Sprintf("%s/%s/remote/%s/blocked/%d", s.login, s.owner, remote, deviceID))
}

// --- generic byte get/set/delete/has, mirroring badger/ids.go's style ---

func (s *Store) get(key []byte) ([]byte, error) {
	var val []byte
	err := s.db.View(func(txn *badger.Txn) error {
		item, err := txn.Get(key)
		if err == badger.ErrKeyNotFound {
			return nil
		}
		if err != nil {
			return err
		}
		val, err = item.ValueCopy(nil)
		return err
	})
	return val, err
}

func (s *Store) set(key, value []byte) error {
	return s.db.Update(func(txn *badger.Txn) error {
		return txn.Set(key, value)
	})
}

func (s *Store) delete(key []byte) error {
	return s.db.Update(func(txn *badger.Txn) error {
		err := txn.Delete(key)
		if err == badger.ErrKeyNotFound {
			return nil
		}
		return err
	})
}

func (s *Store) has(key []byte) (bool, error) {
	found := false
	err := s.db.View(func(txn *badger.Txn) error {
		_, err := txn.Get(key)
		if err == badger.ErrKeyNotFound {
			return nil
		}
		if err != nil {
			return err
		}
		found = true
		return nil
	})
	return found, err
}

func encodeUint32(v uint32) []byte {
	b := make([]byte, 4)
	binary.BigEndian.PutUint32(b, v)
	return b
}

func decodeUint32(b []byte) uint32 {
	if len(b) < 4 {
		return 0
	}
	return binary.BigEndian.Uint32(b)
}

// --- libsignal.IdentityStore ---

func (s *Store) GetIdentityKeyPair() (public, private []byte, err error) {
	record, err := s.get(s.identityKeyPairKey())
	if err != nil {
		return nil, nil, err
	}
	if record == nil {
		return nil, nil, errors.New("badgerstore: no identity key pair stored for this chat - call SaveIdentityKeyPair first")
	}
	return libsignal.SplitIdentityKeyPair(s.ctx, record)
}

func (s *Store) GetLocalRegistrationID() (uint32, error) {
	val, err := s.get(s.registrationIDKey())
	if err != nil {
		return 0, err
	}
	if val == nil {
		return 0, errors.New("badgerstore: no registration id stored for this chat - call SaveRegistrationID first")
	}
	return decodeUint32(val), nil
}

func (s *Store) SaveIdentity(addr libsignal.Address, key []byte) error {
	k := s.remoteIdentityKey(addr.Name, addr.DeviceID)
	if key == nil {
		return s.delete(k)
	}
	return s.set(k, key)
}

// IsTrustedIdentity implements TOFU: a device never seen before is trusted
// implicitly (matching the library's own documented convention); a device
// whose stored key doesn't match the presented one is not. This method
// never writes - SaveIdentity (called by the library itself as a side
// effect of a successful session build) is what persists new identities.
func (s *Store) IsTrustedIdentity(addr libsignal.Address, key []byte) (bool, error) {
	stored, err := s.get(s.remoteIdentityKey(addr.Name, addr.DeviceID))
	if err != nil {
		return false, err
	}
	if stored == nil {
		return true, nil
	}
	return bytes.Equal(stored, key), nil
}

// --- libsignal.PreKeyStore ---

func (s *Store) LoadPreKey(id uint32) ([]byte, bool, error) {
	val, err := s.get(s.preKeyKey(id))
	if err != nil {
		return nil, false, err
	}
	return val, val != nil, nil
}

func (s *Store) StorePreKey(id uint32, record []byte) error {
	return s.set(s.preKeyKey(id), record)
}

func (s *Store) ContainsPreKey(id uint32) bool {
	found, err := s.has(s.preKeyKey(id))
	if err != nil {
		log.WithError(err).Warn("badgerstore: ContainsPreKey read failed")
	}
	return found
}

// ListPreKeyIDs returns the ids of every one-time prekey currently stored
// (i.e. not yet consumed - libomemo-c's session_cipher automatically
// removes a prekey via RemovePreKey the moment it's used to decrypt an
// inbound prekey message, so the surviving ids are neither contiguous nor
// derivable from the counter NextPreKeyIDs advances). Used to publish an
// accurate bundle - see e2ee/omemo's PublishedBundle.
func (s *Store) ListPreKeyIDs() ([]uint32, error) {
	prefix := s.preKeyPrefix()
	var ids []uint32
	err := s.db.View(func(txn *badger.Txn) error {
		opts := badger.DefaultIteratorOptions
		opts.PrefetchValues = false
		it := txn.NewIterator(opts)
		defer it.Close()
		for it.Seek(prefix); it.ValidForPrefix(prefix); it.Next() {
			suffix := it.Item().Key()[len(prefix):]
			// Skips the "counter" key, which shares this prefix - its
			// suffix doesn't parse as a uint32, so it's silently excluded.
			id, err := strconv.ParseUint(string(suffix), 10, 32)
			if err != nil {
				continue
			}
			ids = append(ids, uint32(id))
		}
		return nil
	})
	return ids, err
}

func (s *Store) RemovePreKey(id uint32) error {
	return s.delete(s.preKeyKey(id))
}

// --- libsignal.SignedPreKeyStore ---

func (s *Store) LoadSignedPreKey(id uint32) ([]byte, bool, error) {
	val, err := s.get(s.signedPreKeyKey(id))
	if err != nil {
		return nil, false, err
	}
	return val, val != nil, nil
}

func (s *Store) StoreSignedPreKey(id uint32, record []byte) error {
	return s.set(s.signedPreKeyKey(id), record)
}

func (s *Store) ContainsSignedPreKey(id uint32) bool {
	found, err := s.has(s.signedPreKeyKey(id))
	if err != nil {
		log.WithError(err).Warn("badgerstore: ContainsSignedPreKey read failed")
	}
	return found
}

func (s *Store) RemoveSignedPreKey(id uint32) error {
	return s.delete(s.signedPreKeyKey(id))
}

// --- libsignal.SessionStore ---

func (s *Store) LoadSession(addr libsignal.Address) ([]byte, bool, error) {
	val, err := s.get(s.remoteSessionKey(addr.Name, addr.DeviceID))
	if err != nil {
		return nil, false, err
	}
	return val, val != nil, nil
}

func (s *Store) GetSubDeviceSessions(name string) ([]uint32, error) {
	prefix := s.remoteSessionPrefix(name)
	var deviceIDs []uint32
	err := s.db.View(func(txn *badger.Txn) error {
		opts := badger.DefaultIteratorOptions
		opts.PrefetchValues = false
		it := txn.NewIterator(opts)
		defer it.Close()
		for it.Seek(prefix); it.ValidForPrefix(prefix); it.Next() {
			suffix := it.Item().Key()[len(prefix):]
			id, err := strconv.ParseUint(string(suffix), 10, 32)
			if err != nil {
				continue
			}
			deviceIDs = append(deviceIDs, uint32(id))
		}
		return nil
	})
	return deviceIDs, err
}

func (s *Store) StoreSession(addr libsignal.Address, record []byte) error {
	return s.set(s.remoteSessionKey(addr.Name, addr.DeviceID), record)
}

func (s *Store) ContainsSession(addr libsignal.Address) bool {
	found, err := s.has(s.remoteSessionKey(addr.Name, addr.DeviceID))
	if err != nil {
		log.WithError(err).Warn("badgerstore: ContainsSession read failed")
	}
	return found
}

func (s *Store) DeleteSession(addr libsignal.Address) error {
	return s.delete(s.remoteSessionKey(addr.Name, addr.DeviceID))
}

func (s *Store) DeleteAllSessions(name string) (int, error) {
	prefix := s.remoteSessionPrefix(name)
	count := 0
	err := s.db.Update(func(txn *badger.Txn) error {
		opts := badger.DefaultIteratorOptions
		opts.PrefetchValues = false
		it := txn.NewIterator(opts)
		var keys [][]byte
		for it.Seek(prefix); it.ValidForPrefix(prefix); it.Next() {
			keys = append(keys, append([]byte{}, it.Item().Key()...))
		}
		it.Close()

		for _, k := range keys {
			if err := txn.Delete(k); err != nil {
				return err
			}
			count++
		}
		return nil
	})
	return count, err
}

// --- bootstrap helpers (not part of libsignal.Store - the library only
// ever reads its own identity/registration id, it never writes them; the
// e2ee/omemo package writes these once, via libsignal.GenerateIdentityKeyPair
// / GenerateRegistrationID, when first creating a chat's OMEMO identity) ---

// HasIdentityKeyPair reports whether this chat's identity has already been
// generated, so EnsureIdentity-style callers can skip regenerating it.
func (s *Store) HasIdentityKeyPair() (bool, error) {
	return s.has(s.identityKeyPairKey())
}

// SaveIdentityKeyPair persists the combined identity key pair record (as
// produced by libsignal.GenerateIdentityKeyPair).
func (s *Store) SaveIdentityKeyPair(record []byte) error {
	return s.set(s.identityKeyPairKey(), record)
}

// SaveRegistrationID persists this chat's own registration id (as produced
// by libsignal.GenerateRegistrationID).
func (s *Store) SaveRegistrationID(id uint32) error {
	return s.set(s.registrationIDKey(), encodeUint32(id))
}

// NextPreKeyIDs atomically reserves a contiguous block of count prekey ids
// for libsignal.GeneratePreKeys, returning the first id in the block.
func (s *Store) NextPreKeyIDs(count uint32) (uint32, error) {
	// Starts at 1, not 0: XEP-0384 0.9.1 explicitly requires positive
	// (non-zero) prekey ids ("Fix using id=0 in examples. Spec requires
	// positive numbers.").
	start := uint32(1)
	err := s.db.Update(func(txn *badger.Txn) error {
		item, err := txn.Get(s.preKeyCounterKey())
		if err != nil && err != badger.ErrKeyNotFound {
			return err
		}
		if err == nil {
			val, err := item.ValueCopy(nil)
			if err != nil {
				return err
			}
			start = decodeUint32(val)
		}
		return txn.Set(s.preKeyCounterKey(), encodeUint32(start+count))
	})
	return start, err
}

// Enabled reports whether OMEMO is active for this chat (see the plan's
// two-trigger enablement model: an account-wide config default, or having
// auto-upgraded on first successfully-decrypted inbound message).
func (s *Store) Enabled() (bool, error) {
	val, err := s.get(s.enabledKey())
	if err != nil {
		return false, err
	}
	return len(val) == 1 && val[0] == 1, nil
}

// SetEnabled persists this chat's OMEMO-active flag.
func (s *Store) SetEnabled(enabled bool) error {
	v := byte(0)
	if enabled {
		v = 1
	}
	return s.set(s.enabledKey(), []byte{v})
}

// Override returns this chat's manual on/off override of the account-wide
// OMEMO mode (see e2ee.Mode/e2ee.ShouldEncrypt), and whether one has been
// set at all - isSet=false means "no override, follow the account-wide
// mode", not "forced off".
func (s *Store) Override() (on bool, isSet bool, err error) {
	val, err := s.get(s.overrideKey())
	if err != nil {
		return false, false, err
	}
	if len(val) != 1 {
		return false, false, nil
	}
	return val[0] == 1, true, nil
}

// SetOverride persists this chat's manual on/off override.
func (s *Store) SetOverride(on bool) error {
	v := byte(0)
	if on {
		v = 1
	}
	return s.set(s.overrideKey(), []byte{v})
}

// ClearOverride removes this chat's manual override, reverting it to
// following the account-wide OMEMO mode.
func (s *Store) ClearOverride() error {
	return s.delete(s.overrideKey())
}

// NegotiatedVersion returns the wire-format version last negotiated for
// this chat's outgoing content (see e2ee/omemo's Version enum - this
// package stays version-agnostic and only stores the raw byte), and
// whether one has been recorded yet (false before the first successful
// inbound decrypt - callers should fall back to a default version).
func (s *Store) NegotiatedVersion() (byte, bool, error) {
	val, err := s.get(s.negotiatedVersionKey())
	if err != nil {
		return 0, false, err
	}
	if len(val) != 1 {
		return 0, false, nil
	}
	return val[0], true, nil
}

// SetNegotiatedVersion persists the wire-format version to use for this
// chat's outgoing content.
func (s *Store) SetNegotiatedVersion(v byte) error {
	return s.set(s.negotiatedVersionKey(), []byte{v})
}

// RemoteDeviceListCache returns the last raw device-list document cached
// for remote (see e2ee.DeviceListDoc), and whether one was found - used to
// decide whether a fresh PEP fetch is needed before encrypting.
func (s *Store) RemoteDeviceListCache(remote string) ([]byte, bool, error) {
	val, err := s.get(s.remoteDeviceListKey(remote))
	if err != nil {
		return nil, false, err
	}
	return val, val != nil, nil
}

// SaveRemoteDeviceListCache caches the raw device-list document last fetched
// for remote.
func (s *Store) SaveRemoteDeviceListCache(remote string, doc []byte) error {
	return s.set(s.remoteDeviceListKey(remote), doc)
}

// IsDeviceBlocked reports whether remote's device deviceID has been
// manually blocked - a manual override layered on top of TOFU trust
// (IsTrustedIdentity), which always stays on regardless; this only ever
// narrows it further. Works independent of whether a session currently
// exists for that device.
func (s *Store) IsDeviceBlocked(remote string, deviceID uint32) (bool, error) {
	return s.has(s.blockedKey(remote, deviceID))
}

// SetDeviceBlocked sets or clears remote's device deviceID's manual block.
func (s *Store) SetDeviceBlocked(remote string, deviceID uint32, blocked bool) error {
	if !blocked {
		return s.delete(s.blockedKey(remote, deviceID))
	}
	return s.set(s.blockedKey(remote, deviceID), []byte{1})
}