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() }