Encrypted Badger storage backend for OMEMO keys

e2ee/omemo/libsignal/store.go

@@ -117,8 +117,8 @@ type SignedPreKeyStore interface {
 type SessionStore interface {
 	LoadSession(addr Address) (record []byte, found bool, err error)
 	// GetSubDeviceSessions returns the device IDs of all known sessions
-	// for name, excluding the sentinel device ID 1 (per library
+	// for name (signal_protocol.h's get_sub_device_sessions_func: "all
-	// convention - see get_sub_device_sessions_func in signal_protocol.h).
+	// known devices with active sessions for a recipient").
 	GetSubDeviceSessions(name string) ([]uint32, error)
 	StoreSession(addr Address, record []byte) error
 	ContainsSession(addr Address) bool

e2ee/store/badgerstore/db.go

@@ -0,0 +1,56 @@
+// Package badgerstore is the Badger-backed implementation of
+// libsignal.Store: one bridged chat's OMEMO identity, prekeys, and its
+// remote peers' sessions/trust state.
+//
+// This is a dedicated database, separate from telegabber's existing
+// id-mapping Badger DB (badger.IdsDB) - it holds private key material, so
+// it gets its own lifecycle (optionally encrypted at rest) rather than
+// sharing a file with the purely-operational id-mapping store.
+package badgerstore
+
+import (
+	badger "github.com/dgraph-io/badger/v4"
+
+	"dev.narayana.im/narayana/telegabber/e2ee/omemo/libsignal"
+)
+
+// DB is the OMEMO secrets database, shared across every bridged chat's
+// identity (each gets its own key namespace within it - see Store).
+type DB struct {
+	db *badger.DB
+}
+
+// Open opens (creating if necessary) the OMEMO secrets database at path.
+// If encryptionKey is non-empty, the database is encrypted at rest
+// (badger.Options.WithEncryptionKey); its default BlockCacheSize (256MB)
+// is left as-is, which badger requires to be non-zero whenever encryption
+// is enabled.
+func Open(path string, encryptionKey []byte) (*DB, error) {
+	opts := badger.DefaultOptions(path)
+	if len(encryptionKey) > 0 {
+		opts = opts.WithEncryptionKey(encryptionKey)
+	}
+	bdb, err := badger.Open(opts)
+	if err != nil {
+		return nil, err
+	}
+	return &DB{db: bdb}, nil
+}
+
+// Store returns the libsignal.Store view of one bridged chat's OMEMO
+// identity, scoped by (login, ownerPeerJID). ctx is used to split a stored
+// combined identity key pair record into the separate public/private
+// buffers libsignal.IdentityStore.GetIdentityKeyPair must return.
+func (db *DB) Store(ctx *libsignal.Context, login, ownerPeerJID string) *Store {
+	return &Store{db: db.db, ctx: ctx, login: login, owner: ownerPeerJID}
+}
+
+// Gc compacts the value log.
+func (db *DB) Gc() {
+	db.db.RunValueLogGC(0.7)
+}
+
+// Close closes the database.
+func (db *DB) Close() error {
+	return db.db.Close()
+}

e2ee/store/badgerstore/store.go

@@ -0,0 +1,395 @@
+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>/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>/trust/<deviceId>
+//	<login>/<owner>/remote/<remote>/session/<deviceId>
+//	<login>/<owner>/remote/<remote>/devicelist
+
+func (s *Store) enabledKey() []byte {
+	return []byte(fmt.Sprintf("%s/%s/enabled", 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) 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))
+}
+
+// --- 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
+}
+
+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) {
+	var start uint32
+	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})
+}
+
+// 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)
+}