From 15dd245badcb6e64aa65377255104b40f77897c5 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Wed, 21 Oct 2020 12:17:08 +0100
Subject: [PATCH] Add resolve-state helper

---
 cmd/resolve-state/main.go | 129 ++++++++++++++++++++++++++++++++++++++
 roomserver/state/state.go |   8 +--
 2 files changed, 130 insertions(+), 7 deletions(-)
 create mode 100644 cmd/resolve-state/main.go

diff --git a/cmd/resolve-state/main.go b/cmd/resolve-state/main.go
new file mode 100644
index 000000000..777df2da6
--- /dev/null
+++ b/cmd/resolve-state/main.go
@@ -0,0 +1,129 @@
+package main
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"os"
+	"strconv"
+
+	"github.com/matrix-org/dendrite/internal/caching"
+	"github.com/matrix-org/dendrite/internal/setup"
+	"github.com/matrix-org/dendrite/roomserver/state"
+	"github.com/matrix-org/dendrite/roomserver/storage"
+	"github.com/matrix-org/dendrite/roomserver/types"
+	"github.com/matrix-org/gomatrixserverlib"
+)
+
+// This is a utility for inspecting state snapshots and running state resolution
+// against real snapshots in an actual database.
+// It takes one or more state snapshot NIDs as arguments, along with a room version
+// to use for unmarshalling events, and will produce resolved output.
+
+var roomVersion = flag.String("roomversion", "5", "the room version to parse events as")
+
+// nolint:gocyclo
+func main() {
+	ctx := context.Background()
+	cfg := setup.ParseFlags(true)
+	args := os.Args[1:]
+
+	fmt.Println("Room version:", *roomVersion)
+
+	snapshotNIDs := []types.StateSnapshotNID{}
+	for _, arg := range args {
+		if i, err := strconv.Atoi(arg); err == nil {
+			snapshotNIDs = append(snapshotNIDs, types.StateSnapshotNID(i))
+		}
+	}
+
+	fmt.Println("Fetching", len(snapshotNIDs), "snapshot NIDs")
+
+	cache, err := caching.NewInMemoryLRUCache(true)
+	if err != nil {
+		panic(err)
+	}
+
+	roomserverDB, err := storage.Open(&cfg.RoomServer.Database, cache)
+	if err != nil {
+		panic(err)
+	}
+
+	blockNIDs, err := roomserverDB.StateBlockNIDs(ctx, snapshotNIDs)
+	if err != nil {
+		panic(err)
+	}
+
+	var stateEntries []types.StateEntryList
+	for _, list := range blockNIDs {
+		entries, err2 := roomserverDB.StateEntries(ctx, list.StateBlockNIDs)
+		if err2 != nil {
+			panic(err2)
+		}
+		stateEntries = append(stateEntries, entries...)
+	}
+
+	var eventNIDs []types.EventNID
+	for _, entry := range stateEntries {
+		for _, e := range entry.StateEntries {
+			eventNIDs = append(eventNIDs, e.EventNID)
+		}
+	}
+
+	fmt.Println("Fetching", len(eventNIDs), "state events")
+	eventEntries, err := roomserverDB.Events(ctx, eventNIDs)
+	if err != nil {
+		panic(err)
+	}
+
+	authEventIDMap := make(map[string]struct{})
+	eventPtrs := make([]*gomatrixserverlib.Event, len(eventEntries))
+	for i := range eventEntries {
+		eventPtrs[i] = &eventEntries[i].Event
+		for _, authEventID := range eventEntries[i].AuthEventIDs() {
+			authEventIDMap[authEventID] = struct{}{}
+		}
+	}
+
+	authEventIDs := make([]string, 0, len(authEventIDMap))
+	for authEventID := range authEventIDMap {
+		authEventIDs = append(authEventIDs, authEventID)
+	}
+
+	fmt.Println("Fetching", len(authEventIDs), "auth events")
+	authEventEntries, err := roomserverDB.EventsFromIDs(ctx, authEventIDs)
+	if err != nil {
+		panic(err)
+	}
+
+	authEventPtrs := make([]*gomatrixserverlib.Event, len(authEventEntries))
+	for i := range authEventEntries {
+		authEventPtrs[i] = &authEventEntries[i].Event
+	}
+
+	events := make([]gomatrixserverlib.Event, len(eventEntries))
+	authEvents := make([]gomatrixserverlib.Event, len(authEventEntries))
+	for i, ptr := range eventPtrs {
+		events[i] = *ptr
+	}
+	for i, ptr := range authEventPtrs {
+		authEvents[i] = *ptr
+	}
+
+	fmt.Println("Resolving state")
+	resolved, err := state.ResolveConflictsAdhoc(
+		gomatrixserverlib.RoomVersion(*roomVersion),
+		events,
+		authEvents,
+	)
+	if err != nil {
+		panic(err)
+	}
+
+	fmt.Println("Resolved state:")
+	for _, event := range resolved {
+		fmt.Println()
+		fmt.Printf("* %s %s %q\n", event.EventID(), event.Type(), *event.StateKey())
+		fmt.Printf("  %s\n", string(event.Content()))
+	}
+}
diff --git a/roomserver/state/state.go b/roomserver/state/state.go
index 2944f71c1..d23f14c84 100644
--- a/roomserver/state/state.go
+++ b/roomserver/state/state.go
@@ -526,13 +526,7 @@ func (v StateResolution) CalculateAndStoreStateBeforeEvent(
 	isRejected bool,
 ) (types.StateSnapshotNID, error) {
 	// Load the state at the prev events.
-	prevEventRefs := event.PrevEvents()
-	prevEventIDs := make([]string, len(prevEventRefs))
-	for i := range prevEventRefs {
-		prevEventIDs[i] = prevEventRefs[i].EventID
-	}
-
-	prevStates, err := v.db.StateAtEventIDs(ctx, prevEventIDs)
+	prevStates, err := v.db.StateAtEventIDs(ctx, event.PrevEventIDs())
 	if err != nil {
 		return 0, err
 	}