1 files changed, 39 insertions, 8 deletions

diff --git a/xmpp/component.go b/xmpp/component.go
index 109d3d5..4f5ad79 100644
--- a/xmpp/component.go
+++ b/xmpp/component.go
@@ -2,6 +2,7 @@ package xmpp
 
 import (
 	"github.com/pkg/errors"
+	"sync"
 	"time"
 
 	"dev.narayana.im/narayana/telegabber/config"
@@ -17,6 +18,7 @@ import (
 var tgConf config.TelegramConfig
 var sessions map[string]*telegram.Client
 var db *persistence.SessionsYamlDB
+var sessionLock sync.Mutex
 
 // NewComponent starts a new component and wraps it in
 // a stream manager that you should start yourself
@@ -69,12 +71,14 @@ func heartbeat(component *xmpp.Component) {
 	var err error
 	probeType := gateway.SPType("probe")
 
+	sessionLock.Lock()
 	for jid := range sessions {
 		err = gateway.SendPresence(component, jid, probeType)
 		if err != nil {
 			log.Error(err)
 		}
 	}
+	sessionLock.Unlock()
 
 	log.Info("Starting heartbeat queue")
 
@@ -89,6 +93,13 @@ func heartbeat(component *xmpp.Component) {
 				delete(gateway.Queue, key)
 			}
 		}
+
+		if gateway.DirtySessions {
+			gateway.DirtySessions = false
+			// no problem if a dirty flag gets set again here,
+			// it would be resolved on the next iteration
+			SaveSessions()
+		}
 	}
 }
 
@@ -104,7 +115,10 @@ func loadSessions(dbPath string, component *xmpp.Component) error {
 
 	db.Transaction(func() bool {
 		for jid, session := range db.Data.Sessions {
-			getTelegramInstance(jid, &session, component)
+			// copy the session struct, otherwise all of them would reference
+			// the same temporary range variable
+			currentSession := session
+			getTelegramInstance(jid, &currentSession, component)
 		}
 
 		return false
@@ -122,29 +136,46 @@ func getTelegramInstance(jid string, savedSession *persistence.Session, componen
 			log.Error(errors.Wrap(err, "TDlib initialization failure"))
 			return session, false
 		}
+		if savedSession.KeepOnline {
+			if err = session.Connect(""); err != nil {
+				log.Error(err)
+				return session, false
+			}
+		}
+		sessionLock.Lock()
 		sessions[jid] = session
+		sessionLock.Unlock()
 	}
 
 	return session, true
 }
 
+// SaveSessions dumps current sessions to the file
+func SaveSessions() {
+	sessionLock.Lock()
+	defer sessionLock.Unlock()
+	db.Transaction(func() bool {
+		for jid, session := range sessions {
+			db.Data.Sessions[jid] = *session.Session
+		}
+
+		return true
+	}, persistence.SessionMarshaller)
+}
+
 // Close gracefully terminates the component and saves active sessions
 func Close(component *xmpp.Component) {
 	log.Error("Disconnecting...")
 
+	sessionLock.Lock()
 	// close all sessions
 	for _, session := range sessions {
 		session.Disconnect("", true)
 	}
+	sessionLock.Unlock()
 
 	// save sessions
-	db.Transaction(func() bool {
-		for jid, session := range sessions {
-			db.Data.Sessions[jid] = *session.Session
-		}
-
-		return true
-	}, persistence.SessionMarshaller)
+	SaveSessions()
 
 	// close stream
 	component.Disconnect()