From 99612358a2a9ba1b16cc2b0b117fa0e8710f934f Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Fri, 3 Apr 2020 17:36:09 +0100
Subject: [PATCH] Perform state resolution at send_join

---
 federationapi/routing/join.go |  1 +
 roomserver/api/query.go       |  2 ++
 roomserver/query/query.go     |  8 ++++++
 roomserver/state/state.go     | 47 +++++++++++++++++++++++++++++++++++
 4 files changed, 58 insertions(+)

diff --git a/federationapi/routing/join.go b/federationapi/routing/join.go
index a39ff6394..7edfccb29 100644
--- a/federationapi/routing/join.go
+++ b/federationapi/routing/join.go
@@ -186,6 +186,7 @@ func SendJoin(
 		PrevEventIDs: event.PrevEventIDs(),
 		AuthEventIDs: event.AuthEventIDs(),
 		RoomID:       roomID,
+		ResolveState: true,
 	}, &stateAndAuthChainResponse)
 	if err != nil {
 		util.GetLogger(httpReq.Context()).WithError(err).Error("query.QueryStateAndAuthChain failed")
diff --git a/roomserver/api/query.go b/roomserver/api/query.go
index 9120da4bb..e6fbf5483 100644
--- a/roomserver/api/query.go
+++ b/roomserver/api/query.go
@@ -203,6 +203,8 @@ type QueryStateAndAuthChainRequest struct {
 	PrevEventIDs []string `json:"prev_event_ids"`
 	// The list of auth events for the event. Used to calculate the auth chain
 	AuthEventIDs []string `json:"auth_event_ids"`
+	// Should state resolution be ran on the result events?
+	ResolveState bool `json:"resolve_state"`
 }
 
 // QueryStateAndAuthChainResponse is a response to QueryStateAndAuthChain
diff --git a/roomserver/query/query.go b/roomserver/query/query.go
index b7cdf1507..ccd362774 100644
--- a/roomserver/query/query.go
+++ b/roomserver/query/query.go
@@ -736,6 +736,14 @@ func (r *RoomserverQueryAPI) QueryStateAndAuthChain(
 		return err
 	}
 
+	if request.ResolveState {
+		if stateEvents, err = state.ResolveConflictsAdhoc(
+			roomVersion, stateEvents, authEvents,
+		); err != nil {
+			return err
+		}
+	}
+
 	for _, event := range stateEvents {
 		response.StateEvents = append(response.StateEvents, event.Headered(roomVersion))
 	}
diff --git a/roomserver/state/state.go b/roomserver/state/state.go
index 94873dbeb..d0291a802 100644
--- a/roomserver/state/state.go
+++ b/roomserver/state/state.go
@@ -681,6 +681,53 @@ func (v StateResolution) calculateStateAfterManyEvents(
 	return
 }
 
+func ResolveConflictsAdhoc(
+	version gomatrixserverlib.RoomVersion,
+	events []gomatrixserverlib.Event,
+	authEvents []gomatrixserverlib.Event,
+) ([]gomatrixserverlib.Event, error) {
+	type stateKeyTuple struct {
+		Type     string
+		StateKey *string
+	}
+
+	eventMap := make(map[stateKeyTuple][]gomatrixserverlib.Event)
+	var conflicted, notConflicted, resolved []gomatrixserverlib.Event
+
+	for _, event := range events {
+		tuple := stateKeyTuple{event.Type(), event.StateKey()}
+		if _, ok := eventMap[tuple]; ok {
+			eventMap[tuple] = append(eventMap[tuple], event)
+		} else {
+			eventMap[tuple] = []gomatrixserverlib.Event{event}
+		}
+	}
+
+	for _, list := range eventMap {
+		if len(list) > 1 {
+			conflicted = append(conflicted, list...)
+		} else {
+			notConflicted = append(notConflicted, list...)
+		}
+	}
+
+	stateResAlgo, err := version.StateResAlgorithm()
+	if err != nil {
+		return nil, err
+	}
+	switch stateResAlgo {
+	case gomatrixserverlib.StateResV1:
+		resolved = gomatrixserverlib.ResolveStateConflicts(conflicted, authEvents)
+		resolved = append(resolved, notConflicted...)
+	case gomatrixserverlib.StateResV2:
+		resolved = gomatrixserverlib.ResolveStateConflictsV2(conflicted, notConflicted, authEvents, authEvents)
+	default:
+		return nil, errors.New("unsupported state resolution algorithm")
+	}
+
+	return resolved, nil
+}
+
 func (v StateResolution) resolveConflicts(
 	ctx context.Context, version gomatrixserverlib.RoomVersion,
 	notConflicted, conflicted []types.StateEntry,