+package badger

+

+import (

+ "bytes"

+ "errors"

+ "fmt"

+ "strconv"

+

+ badger "github.com/dgraph-io/badger/v4"

+ log "github.com/sirupsen/logrus"

+)

+

+// IdsDB represents a Badger database

+type IdsDB struct {

+ db *badger.DB

+}

+

+// IdsDBOpen returns a new DB object

+func IdsDBOpen(path string) IdsDB {

+ bdb, err := badger.Open(badger.DefaultOptions(path))

+ if err != nil {

+ log.Errorf("Failed to open ids database: %v, falling back to in-memory database", path)

+ bdb, err = badger.Open(badger.DefaultOptions("").WithInMemory(true))

+ if err != nil {

+ log.Fatalf("Couldn't initialize the ids database")

+ }

+ }

+

+ return IdsDB{

+ db: bdb,

+ }

+}

+

+// Set stores an id pair

+func (db *IdsDB) Set(tgAccount, xmppAccount string, tgChatId, tgMsgId int64, xmppId string) error {

+ bPrefix := toKeyPrefix(tgAccount, xmppAccount)

+ bTgId := toTgByteString(tgChatId, tgMsgId)

+ bXmppId := toXmppByteString(xmppId)

+ bTgKey := toByteKey(bPrefix, bTgId, "tg")

+ bXmppKey := toByteKey(bPrefix, bXmppId, "xmpp")

+

+ return db.db.Update(func(txn *badger.Txn) error {

+ if err := txn.Set(bTgKey, bXmppId); err != nil {

+ return err

+ }

+ return txn.Set(bXmppKey, bTgId)

+ })

+}

+

+func (db *IdsDB) getByteValue(key []byte) ([]byte, error) {

+ var valCopy []byte

+ err := db.db.View(func(txn *badger.Txn) error {

+ item, err := txn.Get(key)

+ if err != nil {

+ return err

+ }

+

+ valCopy, err = item.ValueCopy(nil)

+ return err

+ })

+ return valCopy, err

+}

+

+// GetByTgIds obtains an XMPP id by Telegram chat/message ids

+func (db *IdsDB) GetByTgIds(tgAccount, xmppAccount string, tgChatId, tgMsgId int64) (string, error) {

+ val, err := db.getByteValue(toByteKey(

+ toKeyPrefix(tgAccount, xmppAccount),

+ toTgByteString(tgChatId, tgMsgId),

+ "tg",

+ ))

+ if err != nil {

+ return "", err

+ }

+ return string(val), nil

+}

+

+// GetByXmppId obtains Telegram chat/message ids by an XMPP id

+func (db *IdsDB) GetByXmppId(tgAccount, xmppAccount, xmppId string) (int64, int64, error) {

+ val, err := db.getByteValue(toByteKey(

+ toKeyPrefix(tgAccount, xmppAccount),

+ toXmppByteString(xmppId),

+ "xmpp",

+ ))

+ if err != nil {

+ return 0, 0, err

+ }

+ return splitTgByteString(val)

+}

+

+func toKeyPrefix(tgAccount, xmppAccount string) []byte {

+ return []byte(fmt.Sprintf("%v/%v/", tgAccount, xmppAccount))

+}

+

+func toByteKey(prefix, suffix []byte, typ string) []byte {

+ key := make([]byte, 0, len(prefix)+len(suffix)+6)

+ key = append(key, prefix...)

+ key = append(key, []byte(typ)...)

+ key = append(key, []byte("/")...)

+ key = append(key, suffix...)

+ return key

+}

+

+func toTgByteString(tgChatId, tgMsgId int64) []byte {

+ return []byte(fmt.Sprintf("%v/%v", tgChatId, tgMsgId))

+}

+

+func toXmppByteString(xmppId string) []byte {

+ return []byte(xmppId)

+}

+

+func splitTgByteString(val []byte) (int64, int64, error) {

+ parts := bytes.Split(val, []byte("/"))

+ if len(parts) != 2 {

+ return 0, 0, errors.New("Couldn't parse tg id pair")

+ }

+ tgChatId, err := strconv.ParseInt(string(parts[0]), 10, 64)

+ if err != nil {

+ return 0, 0, err

+ }

+ tgMsgId, err := strconv.ParseInt(string(parts[1]), 10, 64)

+ return tgChatId, tgMsgId, err

+}

+

+// ReplaceIdPair replaces an old entry by XMPP ID with both new XMPP and Tg ID

+func (db *IdsDB) ReplaceIdPair(tgAccount, xmppAccount, oldXmppId, newXmppId string, newMsgId int64) error {

+ // read old pair

+ chatId, oldMsgId, err := db.GetByXmppId(tgAccount, xmppAccount, oldXmppId)

+ if err != nil {

+ return err

+ }

+

+ bPrefix := toKeyPrefix(tgAccount, xmppAccount)

+

+ bOldTgId := toTgByteString(chatId, oldMsgId)

+ bOldXmppId := toXmppByteString(oldXmppId)

+ bOldTgKey := toByteKey(bPrefix, bOldTgId, "tg")

+ bOldXmppKey := toByteKey(bPrefix, bOldXmppId, "xmpp")

+

+ bTgId := toTgByteString(chatId, newMsgId)

+ bXmppId := toXmppByteString(newXmppId)

+ bTgKey := toByteKey(bPrefix, bTgId, "tg")

+ bXmppKey := toByteKey(bPrefix, bXmppId, "xmpp")

+

+ return db.db.Update(func(txn *badger.Txn) error {

+ // save new pair

+ if err := txn.Set(bTgKey, bXmppId); err != nil {

+ return err

+ }

+ if err := txn.Set(bXmppKey, bTgId); err != nil {

+ return err

+ }

+ // delete old pair

+ if err := txn.Delete(bOldTgKey); err != nil {

+ return err

+ }

+ return txn.Delete(bOldXmppKey)

+ })

+}

+

+// ReplaceXmppId replaces an old XMPP ID with new XMPP ID and keeps Tg ID intact

+func (db *IdsDB) ReplaceXmppId(tgAccount, xmppAccount, oldXmppId, newXmppId string) error {

+ // read old Tg IDs

+ chatId, msgId, err := db.GetByXmppId(tgAccount, xmppAccount, oldXmppId)

+ if err != nil {

+ return err

+ }

+

+ bPrefix := toKeyPrefix(tgAccount, xmppAccount)

+

+ bOldXmppId := toXmppByteString(oldXmppId)

+ bOldXmppKey := toByteKey(bPrefix, bOldXmppId, "xmpp")

+

+ bTgId := toTgByteString(chatId, msgId)

+ bXmppId := toXmppByteString(newXmppId)

+ bTgKey := toByteKey(bPrefix, bTgId, "tg")

+ bXmppKey := toByteKey(bPrefix, bXmppId, "xmpp")

+

+ return db.db.Update(func(txn *badger.Txn) error {

+ // save new pair

+ if err := txn.Set(bTgKey, bXmppId); err != nil {

+ return err

+ }

+ if err := txn.Set(bXmppKey, bTgId); err != nil {

+ return err

+ }

+ // delete old xmpp->tg entry

+ return txn.Delete(bOldXmppKey)

+ })

+}

+

+// ReplaceTgId replaces an old Tg ID with new Tg ID and keeps Tg chat ID and XMPP ID intact

+func (db *IdsDB) ReplaceTgId(tgAccount, xmppAccount string, chatId, oldMsgId, newMsgId int64) error {

+ // read old XMPP ID

+ xmppId, err := db.GetByTgIds(tgAccount, xmppAccount, chatId, oldMsgId)

+ if err != nil {

+ return err

+ }

+

+ bPrefix := toKeyPrefix(tgAccount, xmppAccount)

+

+ bOldTgId := toTgByteString(chatId, oldMsgId)

+ bOldTgKey := toByteKey(bPrefix, bOldTgId, "tg")

+

+ bTgId := toTgByteString(chatId, newMsgId)

+ bXmppId := toXmppByteString(xmppId)

+ bTgKey := toByteKey(bPrefix, bTgId, "tg")

+ bXmppKey := toByteKey(bPrefix, bXmppId, "xmpp")

+

+ return db.db.Update(func(txn *badger.Txn) error {

+ // save new pair

+ if err := txn.Set(bTgKey, bXmppId); err != nil {

+ return err

+ }

+ if err := txn.Set(bXmppKey, bTgId); err != nil {

+ return err

+ }

+ // delete old tg->xmpp entry

+ return txn.Delete(bOldTgKey)

+ })

+}

+

+// Gc compacts the value log

+func (db *IdsDB) Gc() {

+ db.db.RunValueLogGC(0.7)

+}

+

+// Close closes a DB

+func (db *IdsDB) Close() {

+ db.db.Close()

+}