From 9fcaadcb2cc45f3255b9451ef6b51d6935dbc291 Mon Sep 17 00:00:00 2001
From: Daniele Sluijters <daenney@users.noreply.github.com>
Date: Thu, 17 Dec 2020 21:49:51 +0100
Subject: [PATCH] store: Add AccountDataStore

The AccountDataStore leverages a user's account data to load/save the
next_batch token, while retaining the InMemoryStore behaviour for saving
filters and room state.

This permits a client to know which messages it has seen/processed when
(re)syncing, for example after a bot is restarted. This isn't possible
with the InMemoryStore since the token would be lost when the
process is restarted resulting in answering to messages we might have
already handled before.

An otherwise stateless client can now use the home server as its
persistent data store for the next_batch token, instead of having to
use a file or database of its own.
---
 store.go | 79 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)

diff --git a/store.go b/store.go
index fcdaa665..6fc9622c 100644
--- a/store.go
+++ b/store.go
@@ -81,3 +81,82 @@ func NewInMemoryStore() *InMemoryStore {
 		Rooms:     make(map[id.RoomID]*Room),
 	}
 }
+
+// AccountDataStore uses account data to store the next batch token, and
+// reuses the InMemoryStore for all other operations.
+type AccountDataStore struct {
+	*InMemoryStore
+	eventType string
+	client    *Client
+}
+
+type accountData struct {
+	NextBatch string `json:"next_batch"`
+}
+
+// SaveNextBatch to account data.
+func (s *AccountDataStore) SaveNextBatch(userID id.UserID, nextBatchToken string) {
+	if userID.String() != s.client.UserID.String() {
+		panic("AccountDataStore must only be used with bots")
+	}
+
+	data := accountData{
+		NextBatch: nextBatchToken,
+	}
+
+	err := s.client.SetAccountData(s.eventType, data)
+	if err != nil {
+		if s.client.Logger != nil {
+			s.client.Logger.Debugfln("failed to save next batch token to account data: %s", err.Error())
+		}
+	}
+}
+
+// LoadNextBatch from account data.
+func (s *AccountDataStore) LoadNextBatch(userID id.UserID) string {
+	if userID.String() != s.client.UserID.String() {
+		panic("AccountDataStore must only be used with bots")
+	}
+
+	data := &accountData{}
+
+	err := s.client.GetAccountData(s.eventType, data)
+	if err != nil {
+		if s.client.Logger != nil {
+			s.client.Logger.Debugfln("failed to load next batch token to account data: %s", err.Error())
+		}
+		return ""
+	}
+
+	return data.NextBatch
+}
+
+// NewAccountDataStore returns a new AccountDataStore, which stores
+// the next_batch token as a custom event in account data in the
+// homeserver.
+//
+// AccountDataStore is only appropriate for bots, not appservices.
+//
+// eventType should be a reversed DNS name like tld.domain.sub.internal and
+// must be unique for a client. The data stored in it is considered internal
+// and must not be modified through outside means. You should also add a filter
+// for account data changes of this event type, to avoid ending up in a sync
+// loop:
+//
+//	mautrix.Filter{
+//		AccountData: mautrix.FilterPart{
+//			Limit: 20,
+//			NotTypes: []event.Type{
+//				event.NewEventType(eventType),
+//			},
+//		},
+//	}
+//	mautrix.Client.CreateFilter(...)
+//
+func NewAccountDataStore(eventType string, client *Client) *AccountDataStore {
+	return &AccountDataStore{
+		InMemoryStore: NewInMemoryStore(),
+		eventType:     eventType,
+		client:        client,
+	}
+}