Fix data storage and other things

.gitignore

@@ -1,5 +1,8 @@
 .idea
-*.session
 
 *.yaml
 !example-config.yaml
+
+*.session
+*.json
+*.db

config/bridge.go

@@ -24,6 +24,7 @@ import (
 type BridgeConfig struct {
 	UsernameTemplate    string             `yaml:"username_template"`
 	DisplaynameTemplate string             `yaml:"displayname_template"`
+	StateStore          string             `yaml:"state_store_path"`
 	usernameTemplate    *template.Template `yaml:"-"`
 	displaynameTemplate *template.Template `yaml:"-"`
 }

config/config.go

@@ -78,6 +78,8 @@ func (config *Config) MakeAppService() (*appservice.AppService, error) {
 	as.LogConfig = config.Logging
 	as.HomeserverDomain = config.Homeserver.Domain
 	as.HomeserverURL = config.Homeserver.Address
+	as.Host.Hostname = config.AppService.Hostname
+	as.Host.Port = config.AppService.Port
 	var err error
 	as.Registration, err = config.GetRegistration()
 	return as, err

database/database.go

@@ -24,7 +24,7 @@ import (
 
 type Database struct {
 	*sql.DB
-	log *log.Sublogger
+	log log.Logger
 
 	User   *UserQuery
 	Portal *PortalQuery
@@ -39,23 +39,29 @@ func New(file string) (*Database, error) {
 
 	db := &Database{
 		DB:  conn,
-		log: log.CreateSublogger("Database", log.LevelDebug),
+		log: log.Sub("Database"),
 	}
 	db.User = &UserQuery{
 		db:  db,
-		log: log.CreateSublogger("Database/User", log.LevelDebug),
+		log: db.log.Sub("User"),
 	}
 	db.Portal = &PortalQuery{
 		db:  db,
-		log: log.CreateSublogger("Database/Portal", log.LevelDebug),
+		log: db.log.Sub("Portal"),
 	}
 	db.Puppet = &PuppetQuery{
 		db:  db,
-		log: log.CreateSublogger("Database/Puppet", log.LevelDebug),
+		log: db.log.Sub("Puppet"),
 	}
 	return db, nil
 }
 
+func (db *Database) CreateTables() {
+	db.User.CreateTable()
+	db.Portal.CreateTable()
+	db.Puppet.CreateTable()
+}
+
 type Scannable interface {
 	Scan(...interface{}) error
 }

database/portal.go

@@ -22,7 +22,7 @@ import (
 
 type PortalQuery struct {
 	db  *Database
-	log *log.Sublogger
+	log log.Logger
 }
 
 func (pq *PortalQuery) CreateTable() error {
@@ -74,7 +74,7 @@ func (pq *PortalQuery) get(query string, args ...interface{}) *Portal {
 
 type Portal struct {
 	db  *Database
-	log *log.Sublogger
+	log log.Logger
 
 	JID   string
 	MXID  string

database/puppet.go

@@ -22,7 +22,7 @@ import (
 
 type PuppetQuery struct {
 	db  *Database
-	log *log.Sublogger
+	log log.Logger
 }
 
 func (pq *PuppetQuery) CreateTable() error {
@@ -67,7 +67,7 @@ func (pq *PuppetQuery) Get(jid, receiver string) *Puppet {
 
 type Puppet struct {
 	db  *Database
-	log *log.Sublogger
+	log log.Logger
 
 	JID      string
 	Receiver string

database/user.go

@@ -23,7 +23,7 @@ import (
 
 type UserQuery struct {
 	db  *Database
-	log *log.Sublogger
+	log log.Logger
 }
 
 func (uq *UserQuery) CreateTable() error {
@@ -71,7 +71,7 @@ func (uq *UserQuery) Get(userID string) *User {
 
 type User struct {
 	db  *Database
-	log *log.Sublogger
+	log log.Logger
 
 	UserID         string
 	ManagementRoom string

example-config.yaml

@@ -46,6 +46,8 @@ bridge:
   # Displayname template for WhatsApp users.
   # {{.displayname}} is replaced with the display name of the WhatsApp user.
   displayname_template: "{{.Displayname}}"
+  # Path to the Matrix room state store.
+  state_store_path: ./mx-state.json
 
 # Logging config.
 logging:
@@ -61,4 +63,4 @@ logging:
   timestamp_format: Jan _2, 2006 15:04:05
   # Minimum severity for log messages.
   # Options: debug, info, warn, error, fatal
-  print_level: info
+  print_level: debug

main.go

@@ -64,7 +64,9 @@ type Bridge struct {
 	AppService *appservice.AppService
 	Config     *config.Config
 	DB         *database.Database
-	Log        *log.Logger
+	Log        log.Logger
+
+	StateStore *AutosavingStateStore
 
 	MatrixListener *MatrixListener
 
@@ -84,33 +86,71 @@ func NewBridge() *Bridge {
 
 func (bridge *Bridge) Init() {
 	var err error
+
 	bridge.AppService, err = bridge.Config.MakeAppService()
 	if err != nil {
 		fmt.Fprintln(os.Stderr, "Failed to initialize AppService:", err)
 		os.Exit(11)
 	}
 	bridge.AppService.Init()
-	bridge.Log = bridge.AppService.Log.Parent
-	log.DefaultLogger = bridge.Log
-	bridge.AppService.Log = log.CreateSublogger("Matrix", log.LevelDebug)
+	bridge.Log = bridge.AppService.Log
+	log.DefaultLogger = bridge.Log.(*log.BasicLogger)
+	bridge.AppService.Log = log.Sub("Matrix")
+
+	bridge.StateStore = NewAutosavingStateStore(bridge.Config.Bridge.StateStore)
+	err = bridge.StateStore.Load()
+	if err != nil {
+		bridge.Log.Fatalln("Failed to load state store:", err)
+		os.Exit(12)
+	}
+	bridge.AppService.StateStore = bridge.StateStore
 
 	bridge.DB, err = database.New(bridge.Config.AppService.Database.URI)
 	if err != nil {
 		bridge.Log.Fatalln("Failed to initialize database:", err)
-		os.Exit(12)
+		os.Exit(13)
 	}
 
 	bridge.MatrixListener = NewMatrixListener(bridge)
 }
 
 func (bridge *Bridge) Start() {
-	bridge.AppService.Start()
-	bridge.MatrixListener.Start()
+	bridge.DB.CreateTables()
+	go bridge.AppService.Start()
+	go bridge.MatrixListener.Start()
+	go bridge.UpdateBotProfile()
+}
+
+func (bridge *Bridge) UpdateBotProfile() {
+	botConfig := bridge.Config.AppService.Bot
+
+	var err error
+	if botConfig.Avatar == "remove" {
+		err = bridge.AppService.BotIntent().SetAvatarURL("")
+	} else if len(botConfig.Avatar) > 0 {
+		err = bridge.AppService.BotIntent().SetAvatarURL(botConfig.Avatar)
+	}
+	if err != nil {
+		bridge.Log.Warnln("Failed to update bot avatar:", err)
+	}
+
+	if botConfig.Displayname == "remove" {
+		err = bridge.AppService.BotIntent().SetDisplayName("")
+	} else if len(botConfig.Avatar) > 0 {
+		err = bridge.AppService.BotIntent().SetDisplayName(botConfig.Displayname)
+	}
+	if err != nil {
+		bridge.Log.Warnln("Failed to update bot displayname:", err)
+	}
 }
 
 func (bridge *Bridge) Stop() {
 	bridge.AppService.Stop()
 	bridge.MatrixListener.Stop()
+	err := bridge.StateStore.Save()
+	if err != nil {
+		bridge.Log.Warnln("Failed to save state store:", err)
+	}
 }
 
 func (bridge *Bridge) Main() {

matrix.go

@@ -25,7 +25,7 @@ import (
 type MatrixListener struct {
 	bridge *Bridge
 	as     *appservice.AppService
-	log    *log.Sublogger
+	log    log.Logger
 	stop   chan struct{}
 }
 
@@ -34,7 +34,7 @@ func NewMatrixListener(bridge *Bridge) *MatrixListener {
 		bridge: bridge,
 		as:     bridge.AppService,
 		stop:   make(chan struct{}, 1),
-		log:    bridge.Log.CreateSublogger("Matrix", log.LevelDebug),
+		log:    bridge.Log.Sub("Matrix Listener"),
 	}
 }
 
@@ -42,7 +42,7 @@ func (ml *MatrixListener) Start() {
 	for {
 		select {
 		case evt := <-ml.bridge.AppService.Events:
-			log.Debugln("Received Matrix event:", evt)
+			ml.log.Debugln("Received Matrix event:", evt)
 			switch evt.Type {
 			case gomatrix.StateMember:
 				ml.HandleMembership(evt)
@@ -56,47 +56,48 @@ func (ml *MatrixListener) Start() {
 }
 
 func (ml *MatrixListener) HandleBotInvite(evt *gomatrix.Event) {
-	cli := ml.as.BotClient()
+	intent := ml.as.BotIntent()
 
-	resp, err := cli.JoinRoom(evt.RoomID, "", nil)
+	resp, err := intent.JoinRoom(evt.RoomID, "", nil)
 	if err != nil {
 		ml.log.Debugln("Failed to join room", evt.RoomID, "with invite from", evt.Sender)
 		return
 	}
 
-	members, err := cli.JoinedMembers(resp.RoomID)
+	members, err := intent.JoinedMembers(resp.RoomID)
 	if err != nil {
 		ml.log.Debugln("Failed to get members in room", resp.RoomID, "after accepting invite from", evt.Sender)
-		cli.LeaveRoom(resp.RoomID)
+		intent.LeaveRoom(resp.RoomID)
 		return
 	}
 
 	if len(members.Joined) < 2 {
 		ml.log.Debugln("Leaving empty room", resp.RoomID, "after accepting invite from", evt.Sender)
-		cli.LeaveRoom(resp.RoomID)
+		intent.LeaveRoom(resp.RoomID)
 		return
 	}
 	for mxid, _ := range members.Joined {
-		if mxid == cli.UserID || mxid == evt.Sender {
+		if mxid == intent.UserID || mxid == evt.Sender {
 			continue
 		} else if true { // TODO check if mxid is WhatsApp puppet
 
 			continue
 		}
 		ml.log.Debugln("Leaving multi-user room", resp.RoomID, "after accepting invite from", evt.Sender)
-		cli.SendNotice(resp.RoomID, "This bridge is user-specific, please don't invite me into rooms with other users.")
-		cli.LeaveRoom(resp.RoomID)
+		intent.SendNotice(resp.RoomID, "This bridge is user-specific, please don't invite me into rooms with other users.")
+		intent.LeaveRoom(resp.RoomID)
 		return
 	}
 
 	user := ml.bridge.GetUser(evt.Sender)
 	user.ManagementRoom = resp.RoomID
 	user.Update()
-	cli.SendNotice(user.ManagementRoom, "This room has been registered as your bridge management/status room.")
+	intent.SendNotice(user.ManagementRoom, "This room has been registered as your bridge management/status room.")
 	ml.log.Debugln(resp.RoomID, "registered as a management room with", evt.Sender)
 }
 
 func (ml *MatrixListener) HandleMembership(evt *gomatrix.Event) {
+	ml.log.Debugln(evt.Content, evt.Content.Membership, evt.GetStateKey())
 	if evt.Content.Membership == "invite" && evt.GetStateKey() == ml.as.BotMXID() {
 		ml.HandleBotInvite(evt)
 	}

portal.go

@@ -76,7 +76,7 @@ func (user *User) NewPortal(dbPortal *database.Portal) *Portal {
 		Portal: dbPortal,
 		user:   user,
 		bridge: user.bridge,
-		log:    user.bridge.Log.CreateSublogger(fmt.Sprintf("Portal/%s/%s", user.UserID, dbPortal.JID), log.LevelDebug),
+		log:    user.log.Sub(fmt.Sprintf("Portal/%s", dbPortal.JID)),
 	}
 }
 
@@ -85,5 +85,5 @@ type Portal struct {
 
 	user   *User
 	bridge *Bridge
-	log    *log.Sublogger
+	log    log.Logger
 }

statestore.go

@@ -0,0 +1,67 @@
+// mautrix-whatsapp - A Matrix-WhatsApp puppeting bridge.
+// Copyright (C) 2018 Tulir Asokan
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+package main
+
+import (
+	"maunium.net/go/mautrix-appservice"
+	"encoding/json"
+	"io/ioutil"
+	"os"
+)
+
+type AutosavingStateStore struct {
+	*appservice.BasicStateStore
+	Path string
+}
+
+func NewAutosavingStateStore(path string) *AutosavingStateStore {
+	return &AutosavingStateStore{
+		BasicStateStore: appservice.NewBasicStateStore(),
+		Path:            path,
+	}
+}
+
+func (store *AutosavingStateStore) Save() error {
+	data, err := json.Marshal(store.BasicStateStore)
+	if err != nil {
+		return err
+	}
+
+	return ioutil.WriteFile(store.Path, data, 0600)
+}
+
+func (store *AutosavingStateStore) Load() error {
+	data, err := ioutil.ReadFile(store.Path)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return nil
+		}
+		return err
+	}
+
+	return json.Unmarshal(data, store.BasicStateStore)
+}
+
+func (store *AutosavingStateStore) MarkRegistered(userID string) {
+	store.BasicStateStore.MarkRegistered(userID)
+	store.Save()
+}
+
+func (store *AutosavingStateStore) SetMembership(roomID, userID, membership string) {
+	store.BasicStateStore.SetMembership(roomID, userID, membership)
+	store.Save()
+}

user.go

@@ -31,7 +31,7 @@ type User struct {
 	Conn *whatsapp.Conn
 
 	bridge *Bridge
-	log    *log.Sublogger
+	log    log.Logger
 
 	portalsByMXID map[string]*Portal
 	portalsByJID  map[string]*Portal
@@ -77,7 +77,7 @@ func (bridge *Bridge) NewUser(dbUser *database.User) *User {
 	return &User{
 		User:   dbUser,
 		bridge: bridge,
-		log:    bridge.Log.CreateSublogger(fmt.Sprintf("User/%s", dbUser.UserID), log.LevelDebug),
+		log:    bridge.Log.Sub("User").Sub(dbUser.UserID),
 	}
 }