diff --git a/.gitignore b/.gitignore index 772ab8b..e2dbbf2 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ config.yml telegabber sessions/ +session.dat diff --git a/config.yml.example b/config.yml.example index a31cffd..5179508 100644 --- a/config.yml.example +++ b/config.yml.example @@ -21,4 +21,4 @@ :host: '127.0.0.1' :port: 8899 :password: 'password' - :db: 'sessions.dat' + :db: 'session.dat' diff --git a/persistence/sessions.go b/persistence/sessions.go new file mode 100644 index 0000000..50e3f3e --- /dev/null +++ b/persistence/sessions.go @@ -0,0 +1,69 @@ +package persistence + +import ( + "github.com/pkg/errors" + "io/ioutil" + + "dev.narayana.im/narayana/telegabber/yamldb" + + "gopkg.in/yaml.v2" +) + +// SessionsYamlDB wraps YamlDB with Session +type SessionsYamlDB struct { + yamldb.YamlDB + Data *SessionsMap +} + +// SessionsMap is for :sessions: subtree +type SessionsMap struct { + Sessions map[string]Session `yaml:":sessions"` +} + +// Session is a key-values subtree +type Session struct { + Login string `yaml:":login"` +} + +var sessionDB SessionsYamlDB + +// Marshaller implementation for YamlDB +func Marshaller() ([]byte, error) { + return yaml.Marshal(sessionDB.Data) +} + +// LoadSessions restores TDlib sessions from the previous run +func LoadSessions(path string) (SessionsYamlDB, error) { + var sessionData SessionsMap + + sessionDB, err := initYamlDB(path, &sessionData) + if err != nil { + return sessionDB, errors.Wrap(err, "Sessions restore error") + } + + sessionDB.Transaction(func() { + }, Marshaller) + + return sessionDB, nil +} + +func initYamlDB(path string, dataPtr *SessionsMap) (SessionsYamlDB, error) { + file, err := ioutil.ReadFile(path) + if err == nil { + err = yaml.Unmarshal(file, dataPtr) + if err != nil { + return SessionsYamlDB{}, errors.Wrap(err, "YamlDB is corrupted") + } + } else { + // DB file does not exist, create an empty DB + dataPtr.Sessions = make(map[string]Session) + } + + return SessionsYamlDB{ + YamlDB: yamldb.YamlDB{ + Path: path, + PathNew: path + ".new", + }, + Data: dataPtr, + }, nil +} diff --git a/xmpp/component.go b/xmpp/component.go index 036bc3c..fcef91b 100644 --- a/xmpp/component.go +++ b/xmpp/component.go @@ -2,6 +2,7 @@ package xmpp import ( "dev.narayana.im/narayana/telegabber/config" + "dev.narayana.im/narayana/telegabber/persistence" "dev.narayana.im/narayana/telegabber/telegram" "gosrc.io/xmpp" @@ -10,11 +11,13 @@ import ( var jid *xmpp.Jid var tgConf config.TelegramConfig var sessions map[string]telegram.Client +var db persistence.SessionsYamlDB // NewComponent starts a new component and wraps it in // a stream manager that you should start yourself func NewComponent(conf config.XMPPConfig, tc config.TelegramConfig) (*xmpp.StreamManager, error) { var err error + jid, err = xmpp.NewJid(conf.Jid) if err != nil { return nil, err @@ -23,6 +26,11 @@ func NewComponent(conf config.XMPPConfig, tc config.TelegramConfig) (*xmpp.Strea tgConf = tc sessions = make(map[string]telegram.Client) + db, err = persistence.LoadSessions(conf.Db) + if err != nil { + return nil, err + } + options := xmpp.ComponentOptions{ Address: conf.Host + ":" + conf.Port, Domain: conf.Jid, diff --git a/yamldb/yamldb.go b/yamldb/yamldb.go new file mode 100644 index 0000000..eec23be --- /dev/null +++ b/yamldb/yamldb.go @@ -0,0 +1,47 @@ +package yamldb + +import ( + "github.com/pkg/errors" + "io/ioutil" + "os" + "sync" + + log "github.com/sirupsen/logrus" +) + +// YamlDB represents a YAML file database instance +type YamlDB struct { + Path string + PathNew string + lock sync.Mutex +} + +// Transaction executes the given callback and safely saves +// the data after they are modified within the callback +func (db *YamlDB) Transaction(callback func(), marshaller func() ([]byte, error)) error { + var err error + + log.Debug("Enter transaction") + db.lock.Lock() + defer func() { + db.lock.Unlock() + log.Debug("Exit transaction") + }() + + callback() + + yamlData, err := marshaller() + if err != nil { + return errors.Wrap(err, "Data marshalling error") + } + err = ioutil.WriteFile(db.PathNew, yamlData, 0644) + if err != nil { + return errors.Wrap(err, "YamlDB write failure") + } + err = os.Rename(db.PathNew, db.Path) + if err != nil { + return errors.Wrap(err, "Couldn't rewrite an old YamlDB file") + } + + return nil +}