From fbe99c65ec36bb0d8b3a508e0894a9830ff9c995 Mon Sep 17 00:00:00 2001
From: bodqhrohro <bodqhrohro@gmail.com>
Date: Mon, 11 Nov 2019 01:50:50 +0200
Subject: [PATCH] Implementation of Ruby YAML::Store and read/save of the
 session database

---
 .gitignore              |  1 +
 config.yml.example      |  2 +-
 persistence/sessions.go | 69 +++++++++++++++++++++++++++++++++++++++++
 xmpp/component.go       |  8 +++++
 yamldb/yamldb.go        | 47 ++++++++++++++++++++++++++++
 5 files changed, 126 insertions(+), 1 deletion(-)
 create mode 100644 persistence/sessions.go
 create mode 100644 yamldb/yamldb.go

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
+}