easybridge/db.go

package main

import (
	"fmt"
	"sync"

	"github.com/hashicorp/golang-lru"
	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/mysql"
	_ "github.com/jinzhu/gorm/dialects/postgres"
	_ "github.com/jinzhu/gorm/dialects/sqlite"
	log "github.com/sirupsen/logrus"
	"golang.org/x/crypto/blake2b"

	"git.deuxfleurs.fr/Deuxfleurs/easybridge/connector"
	"git.deuxfleurs.fr/Deuxfleurs/easybridge/mxlib"
)

var db *gorm.DB
var dbCache *lru.TwoQueueCache

func InitDb() error {
	var err error

	db, err = gorm.Open(config.DbType, config.DbPath)
	if err != nil {
		return err
	}

	db.AutoMigrate(&DbAccountConfig{})

	db.AutoMigrate(&DbJoinedRoom{})
	db.Model(&DbJoinedRoom{}).AddIndex("idx_joined_room_user_protocol_account", "mx_user_id", "protocol", "account_name")

	db.AutoMigrate(&DbUserMap{})
	db.Model(&DbUserMap{}).AddIndex("idx_user_map_protocol_user", "protocol", "user_id")

	db.AutoMigrate(&DbRoomMap{})
	db.Model(&DbRoomMap{}).AddIndex("idx_room_map_protocol_room", "protocol", "room_id")

	db.AutoMigrate(&DbPmRoomMap{})
	db.Model(&DbPmRoomMap{}).AddIndex("idx_pm_room_map_protocol_user_account_user", "protocol", "user_id", "mx_user_id", "account_name")

	db.AutoMigrate(&DbKv{})

	dbCache, err = lru.New2Q(10000)
	if err != nil {
		return err
	}

	return nil
}

// Account configuration
type DbAccountConfig struct {
	gorm.Model

	MxUserID string `gorm:"type:text;index"`
	Name     string `gorm:"type:text"`
	Protocol string `gorm:"type:text"`
	Config   string `gorm:"type:text"`
}

// List of joined channels to be re-joined on reconnect
type DbJoinedRoom struct {
	ID uint `gorm:"primary_key"`

	// User id and account name
	MxUserID    string `gorm:"type:text"`
	Protocol    string `gorm:"type:text"`
	AccountName string `gorm:"type:text"`

	// Room ID
	RoomID connector.RoomID `gorm:"type:text"`
}

// User mapping between protocol user IDs and puppeted matrix ids
type DbUserMap struct {
	ID uint `gorm:"primary_key"`

	Protocol string           `gorm:"type:text"`
	UserID   connector.UserID `gorm:"type:text"`
	MxUserID string           `gorm:"type:text;index"`
}

// Room mapping between Matrix rooms and outside rooms
type DbRoomMap struct {
	ID uint `gorm:"primary_key"`

	// Network protocol
	Protocol string `gorm:"type:text"`

	// Room id on the bridged network
	RoomID connector.RoomID `gorm:"type:text"`

	// Bridged room matrix id
	MxRoomID string `gorm:"type:text;index"`
}

// Room mapping between Matrix rooms and private messages
type DbPmRoomMap struct {
	ID uint `gorm:"primary_key"`

	// User id and account name of the local end viewed on Matrix
	MxUserID    string `gorm:"type:text"`
	Protocol    string `gorm:"type:text"`
	AccountName string `gorm:"type:text"`

	// User id to reach them
	UserID connector.UserID `gorm:"type:text"`

	// Bridged room for PMs
	MxRoomID string `gorm:"type:text;index"`
}

// Key-value store for various things
type DbKv struct {
	Key   string `gorm:"type:text;primary_key"`
	Value string `gorm:"type:text"`
}

// ---- Simple locking mechanism
// Slot keys are strings that identify the object we are acting upon
// They define which lock to lock for a certain operation
// They are also used as keys in the LRU cache

var dbLocks [256]sync.Mutex

func dbLockSlot(key string) {
	slot := blake2b.Sum512([]byte(key))[0]
	dbLocks[slot].Lock()
}

func dbUnlockSlot(key string) {
	slot := blake2b.Sum512([]byte(key))[0]
	dbLocks[slot].Unlock()
}

// ---- Key-value store supporting atomic test-and-set

func dbKvSlotKey(key string) string {
	return "kv:" + key
}

func dbKvGet(key string) string {
	slot_key := dbKvSlotKey(key)

	if ent, ok := dbCache.Get(slot_key); ok {
		return ent.(string)
	}

	var entry DbKv
	if db.Where(&DbKv{Key: key}).First(&entry).RecordNotFound() {
		return ""
	} else {
		return entry.Value
	}
}

func dbKvPut(key string, value string) {
	slot_key := dbKvSlotKey(key)

	dbLockSlot(slot_key)
	defer dbUnlockSlot(slot_key)

	var entry DbKv
	db.Where(&DbKv{Key: key}).Assign(&DbKv{Value: value}).FirstOrCreate(&entry)
	dbCache.Add(slot_key, value)
}

func dbKvTestAndSet(key string, value string) bool {
	slot_key := dbKvSlotKey(key)

	dbLockSlot(slot_key)
	defer dbUnlockSlot(slot_key)

	// True if value was changed, false if was already set
	if dbKvGet(key) == value {
		return false
	}

	var entry DbKv
	db.Where(&DbKv{Key: key}).Assign(&DbKv{Value: value}).FirstOrCreate(&entry)
	dbCache.Add(slot_key, value)
	return true
}

// ----

func dbGetMxRoom(protocol string, roomId connector.RoomID) (string, error) {
	slot_key := fmt.Sprintf("room:%s/%s", protocol, roomId)

	dbLockSlot(slot_key)
	defer dbUnlockSlot(slot_key)

	if cached, ok := dbCache.Get(slot_key); ok {
		return cached.(string), nil
	}

	// Check if room exists in our mapping,
	// If not create it
	var room DbRoomMap
	must_create := db.First(&room, DbRoomMap{
		Protocol: protocol,
		RoomID:   roomId,
	}).RecordNotFound()

	if must_create {
		alias := roomAlias(protocol, roomId)
		// Lookup alias
		mx_room_id, err := mx.DirectoryRoom(fmt.Sprintf("#%s:%s", alias, config.MatrixDomain))

		// If no alias found, create room
		if err != nil {
			name := fmt.Sprintf("%s (%s)", roomId, protocol)

			mx_room_id, err = mx.CreateRoom(name, alias, []string{})
			if err != nil {
				log.Warnf("Could not create room for %s: %s", name, err)
				return "", err
			}
		}

		room = DbRoomMap{
			Protocol: protocol,
			RoomID:   roomId,
			MxRoomID: mx_room_id,
		}
		db.Create(&room)
	}

	log.Tracef("%s -> %s", slot_key, room.MxRoomID)
	dbCache.Add(slot_key, room.MxRoomID)

	return room.MxRoomID, nil
}

func dbPmRoomSlotKey(room *DbPmRoomMap) string {
	return fmt.Sprintf("pmroom:%s/%s/%s/%s",
		room.Protocol, room.MxUserID, room.AccountName, room.UserID)
}

func dbGetMxPmRoom(protocol string, them connector.UserID, themMxId string, usMxId string, usAccount string) (string, error) {
	map_key := &DbPmRoomMap{
		MxUserID:    usMxId,
		Protocol:    protocol,
		AccountName: usAccount,
		UserID:      them,
	}
	slot_key := dbPmRoomSlotKey(map_key)

	dbLockSlot(slot_key)
	defer dbUnlockSlot(slot_key)

	if cached, ok := dbCache.Get(slot_key); ok {
		return cached.(string), nil
	}

	var room DbPmRoomMap
	must_create := db.First(&room, map_key).RecordNotFound()

	if must_create {
		name := fmt.Sprintf("%s (%s)", them, protocol)

		mx_room_id, err := mx.CreateDirectRoomAs([]string{usMxId}, themMxId)
		if err != nil {
			log.Warnf("Could not create room for %s: %s", name, err)
			return "", err
		}

		room = DbPmRoomMap{
			MxUserID:    usMxId,
			Protocol:    protocol,
			AccountName: usAccount,
			UserID:      them,
			MxRoomID:    mx_room_id,
		}
		db.Create(&room)
	}

	log.Tracef("%s -> %s", slot_key, room.MxRoomID)
	dbCache.Add(slot_key, room.MxRoomID)

	return room.MxRoomID, nil
}

func dbDeletePmRoom(room *DbPmRoomMap) {
	slot_key := dbPmRoomSlotKey(room)

	dbLockSlot(slot_key)
	defer dbUnlockSlot(slot_key)

	db.Delete(room)
	dbCache.Remove(slot_key)
}

func dbGetMxUser(protocol string, userId connector.UserID) (string, error) {
	slot_key := fmt.Sprintf("user:%s/%s", protocol, userId)

	dbLockSlot(slot_key)
	defer dbUnlockSlot(slot_key)

	if cached, ok := dbCache.Get(slot_key); ok {
		return cached.(string), nil
	}

	var user DbUserMap

	must_create := db.First(&user, DbUserMap{
		Protocol: protocol,
		UserID:   userId,
	}).RecordNotFound()
	if must_create {
		username := userMxId(protocol, userId)

		err := mx.RegisterUser(username)
		if err != nil {
			if mxE, ok := err.(*mxlib.MxError); !ok || mxE.ErrCode != "M_USER_IN_USE" {
				log.Warnf("Could not register %s: %s", username, err)
				return "", err
			}
		}

		mxid := fmt.Sprintf("@%s:%s", username, config.MatrixDomain)
		mx.ProfileDisplayname(mxid, fmt.Sprintf("%s (%s)", userId, protocol))

		user = DbUserMap{
			Protocol: protocol,
			UserID:   userId,
			MxUserID: mxid,
		}
		db.Create(&user)
	}

	log.Tracef("%s -> %s", slot_key, user.MxUserID)
	dbCache.Add(slot_key, user.MxUserID)

	return user.MxUserID, nil
}

func dbIsPmRoom(mxRoomId string) *DbPmRoomMap {
	var pm_room DbPmRoomMap
	if db.First(&pm_room, DbPmRoomMap{MxRoomID: mxRoomId}).RecordNotFound() {
		return nil
	} else {
		return &pm_room
	}
}

func dbIsPublicRoom(mxRoomId string) *DbRoomMap {
	var room DbRoomMap
	if db.First(&room, DbRoomMap{MxRoomID: mxRoomId}).RecordNotFound() {
		return nil
	} else {
		return &room
	}
}