Handle state with input event as new events (#1415)

* SendEventWithState events as new

* Use cumulative state IDs for final event

* Error wrapping in calculateAndSetState

* Handle overwriting same event type and state key

* Hacky way to spot historical events

* Don't exclude from sync

* Don't generate output events when rewriting forward extremities

* Update output event check

* Historical output events

* Define output room event type

* Notify key changes on state

* Don't send our membership event twice

* Deduplicate state entries

* Tweaks

* Remove unnecessary nolint

* Fix current state upsert in sync API

* Send auth events as outliers, state events as rewrite

* Sync API don't consume state events

* Process events actually

* Improve outlier check

* Fix local room check

* Remove extra room check, it seems to break the whole damn world

* Fix federated join check

* Fix nil pointer exception

* Better comments on DeduplicateStateEntries

* Reflow forced federated joins

* Don't force federated join for possibly even local invites

* Comment SendEventWithState better

* Rewrite room state in sync API storage

* Add TODO

* Clean up all room data when receiving create event

* Don't generate output events for rewrites, but instead notify that state is rewritten on the final new event

* Rename to PurgeRoom

* Exclude backfilled messages from /sync

* Split out rewriting state from updating state from state res

Co-authored-by: Kegan Dougal <kegan@matrix.org>

10 files changed, 454 insertions, 27 deletions

diff --git a/roomserver/api/input.go b/roomserver/api/input.go
index 73c4994a..651c0e9f 100644
--- a/roomserver/api/input.go
+++ b/roomserver/api/input.go
@@ -33,6 +33,10 @@ const (
 	// KindBackfill event extend the contiguous graph going backwards.
 	// They always have state.
 	KindBackfill = 3
+	// KindRewrite events are used when rewriting the head of the room
+	// graph with entirely new state. The output events generated will
+	// be state events rather than timeline events.
+	KindRewrite = 4
 )
 
 // DoNotSendToOtherServers tells us not to send the event to other matrix
diff --git a/roomserver/api/output.go b/roomserver/api/output.go
index 013ebdc8..d57f3b04 100644
--- a/roomserver/api/output.go
+++ b/roomserver/api/output.go
@@ -68,6 +68,17 @@ type OutputEvent struct {
 	NewPeek *OutputNewPeek `json:"new_peek,omitempty"`
 }
 
+// Type of the OutputNewRoomEvent.
+type OutputRoomEventType int
+
+const (
+	// The event is a timeline event and likely just happened.
+	OutputRoomTimeline OutputRoomEventType = iota
+
+	// The event is a state event and quite possibly happened in the past.
+	OutputRoomState
+)
+
 // An OutputNewRoomEvent is written when the roomserver receives a new event.
 // It contains the full matrix room event and enough information for a
 // consumer to construct the current state of the room and the state before the
@@ -80,6 +91,9 @@ type OutputEvent struct {
 type OutputNewRoomEvent struct {
 	// The Event.
 	Event gomatrixserverlib.HeaderedEvent `json:"event"`
+	// Does the event completely rewrite the room state? If so, then AddsStateEventIDs
+	// will contain the entire room state.
+	RewritesState bool `json:"rewrites_state"`
 	// The latest events in the room after this event.
 	// This can be used to set the prev events for new events in the room.
 	// This also can be used to get the full current state after this event.
diff --git a/roomserver/api/wrapper.go b/roomserver/api/wrapper.go
index 82a4a571..e5339311 100644
--- a/roomserver/api/wrapper.go
+++ b/roomserver/api/wrapper.go
@@ -80,6 +80,107 @@ func SendEventWithState(
 	return SendInputRoomEvents(ctx, rsAPI, ires)
 }
 
+// SendEventWithRewrite writes an event with KindNew to the roomserver along
+// with a number of rewrite and outlier events for state and auth events
+// respectively.
+func SendEventWithRewrite(
+	ctx context.Context, rsAPI RoomserverInternalAPI, state *gomatrixserverlib.RespState,
+	event gomatrixserverlib.HeaderedEvent, haveEventIDs map[string]bool,
+) error {
+	isCurrentState := map[string]struct{}{}
+	for _, se := range state.StateEvents {
+		isCurrentState[se.EventID()] = struct{}{}
+	}
+
+	authAndStateEvents, err := state.Events()
+	if err != nil {
+		return err
+	}
+
+	var ires []InputRoomEvent
+	var stateIDs []string
+
+	// This function generates three things:
+	// A - A set of "rewrite" events, which will form the newly rewritten
+	//     state before the event, which includes every rewrite event that
+	//     came before it in its state
+	// B - A set of "outlier" events, which are auth events but not part
+	//     of the rewritten state
+	// C - A "new" event, which include all of the rewrite events in its
+	//     state
+	for _, authOrStateEvent := range authAndStateEvents {
+		if authOrStateEvent.StateKey() == nil {
+			continue
+		}
+		if haveEventIDs[authOrStateEvent.EventID()] {
+			continue
+		}
+		if event.StateKey() == nil {
+			continue
+		}
+
+		// We will handle an event as if it's an outlier if one of the
+		// following conditions is true:
+		storeAsOutlier := false
+		if authOrStateEvent.Type() == event.Type() && *authOrStateEvent.StateKey() == *event.StateKey() {
+			// The event is a state event but the input event is going to
+			// replace it, therefore it can't be added to the state or we'll
+			// get duplicate state keys in the state block. We'll send it
+			// as an outlier because we don't know if something will be
+			// referring to it as an auth event, but need it to be stored
+			// just in case.
+			storeAsOutlier = true
+		} else if _, ok := isCurrentState[authOrStateEvent.EventID()]; !ok {
+			// The event is an auth event and isn't a part of the state set.
+			// We'll send it as an outlier because we need it to be stored
+			// in case something is referring to it as an auth event.
+			storeAsOutlier = true
+		}
+
+		if storeAsOutlier {
+			ires = append(ires, InputRoomEvent{
+				Kind:         KindOutlier,
+				Event:        authOrStateEvent.Headered(event.RoomVersion),
+				AuthEventIDs: authOrStateEvent.AuthEventIDs(),
+			})
+			continue
+		}
+
+		// If the event isn't an outlier then we'll instead send it as a
+		// rewrite event, so that it'll form part of the rewritten state.
+		// These events will go through the membership and latest event
+		// updaters and we will generate output events, but they will be
+		// flagged as non-current (i.e. didn't just happen) events.
+		// Each of these rewrite events includes all of the rewrite events
+		// that came before in their StateEventIDs.
+		ires = append(ires, InputRoomEvent{
+			Kind:          KindRewrite,
+			Event:         authOrStateEvent.Headered(event.RoomVersion),
+			AuthEventIDs:  authOrStateEvent.AuthEventIDs(),
+			HasState:      true,
+			StateEventIDs: stateIDs,
+		})
+
+		// Add the event ID into the StateEventIDs of all subsequent
+		// rewrite events, and the new event.
+		stateIDs = append(stateIDs, authOrStateEvent.EventID())
+	}
+
+	// Send the final event as a new event, which will generate
+	// a timeline output event for it. All of the rewrite events
+	// that came before will be sent as StateEventIDs, forming a
+	// new clean state before the event.
+	ires = append(ires, InputRoomEvent{
+		Kind:          KindNew,
+		Event:         event,
+		AuthEventIDs:  event.AuthEventIDs(),
+		HasState:      true,
+		StateEventIDs: stateIDs,
+	})
+
+	return SendInputRoomEvents(ctx, rsAPI, ires)
+}
+
 // SendInputRoomEvents to the roomserver.
 func SendInputRoomEvents(
 	ctx context.Context, rsAPI RoomserverInternalAPI, ires []InputRoomEvent,
diff --git a/roomserver/internal/helpers/auth.go b/roomserver/internal/helpers/auth.go
index 060f0a0e..524a5451 100644
--- a/roomserver/internal/helpers/auth.go
+++ b/roomserver/internal/helpers/auth.go
@@ -36,7 +36,7 @@ func CheckAuthEvents(
 	if err != nil {
 		return nil, err
 	}
-	// TODO: check for duplicate state keys here.
+	authStateEntries = types.DeduplicateStateEntries(authStateEntries)
 
 	// Work out which of the state events we actually need.
 	stateNeeded := gomatrixserverlib.StateNeededForAuth([]gomatrixserverlib.Event{event.Unwrap()})
diff --git a/roomserver/internal/input/input_events.go b/roomserver/internal/input/input_events.go
index 6ee679da..daf1afcd 100644
--- a/roomserver/internal/input/input_events.go
+++ b/roomserver/internal/input/input_events.go
@@ -86,7 +86,7 @@ func (r *Inputer) processRoomEvent(
 			"event_id": event.EventID(),
 			"type":     event.Type(),
 			"room":     event.RoomID(),
-		}).Info("Stored outlier")
+		}).Debug("Stored outlier")
 		return event.EventID(), nil
 	}
 
@@ -107,6 +107,15 @@ func (r *Inputer) processRoomEvent(
 		}
 	}
 
+	if input.Kind == api.KindRewrite {
+		logrus.WithFields(logrus.Fields{
+			"event_id": event.EventID(),
+			"type":     event.Type(),
+			"room":     event.RoomID(),
+		}).Debug("Stored rewrite")
+		return event.EventID(), nil
+	}
+
 	if err = r.updateLatestEvents(
 		ctx,                 // context
 		roomInfo,            // room info for the room being updated
@@ -114,6 +123,7 @@ func (r *Inputer) processRoomEvent(
 		event,               // event
 		input.SendAsServer,  // send as server
 		input.TransactionID, // transaction ID
+		input.HasState,      // rewrites state?
 	); err != nil {
 		return "", fmt.Errorf("r.updateLatestEvents: %w", err)
 	}
@@ -167,19 +177,25 @@ func (r *Inputer) calculateAndSetState(
 		// Check that those state events are in the database and store the state.
 		var entries []types.StateEntry
 		if entries, err = r.DB.StateEntriesForEventIDs(ctx, input.StateEventIDs); err != nil {
-			return err
+			return fmt.Errorf("r.DB.StateEntriesForEventIDs: %w", err)
 		}
+		entries = types.DeduplicateStateEntries(entries)
 
 		if stateAtEvent.BeforeStateSnapshotNID, err = r.DB.AddState(ctx, roomInfo.RoomNID, nil, entries); err != nil {
-			return err
+			return fmt.Errorf("r.DB.AddState: %w", err)
 		}
 	} else {
 		stateAtEvent.Overwrite = false
 
 		// We haven't been told what the state at the event is so we need to calculate it from the prev_events
 		if stateAtEvent.BeforeStateSnapshotNID, err = roomState.CalculateAndStoreStateBeforeEvent(ctx, event); err != nil {
-			return err
+			return fmt.Errorf("roomState.CalculateAndStoreStateBeforeEvent: %w", err)
 		}
 	}
-	return r.DB.SetState(ctx, stateAtEvent.EventNID, stateAtEvent.BeforeStateSnapshotNID)
+
+	err = r.DB.SetState(ctx, stateAtEvent.EventNID, stateAtEvent.BeforeStateSnapshotNID)
+	if err != nil {
+		return fmt.Errorf("r.DB.SetState: %w", err)
+	}
+	return nil
 }
diff --git a/roomserver/internal/input/input_latest_events.go b/roomserver/internal/input/input_latest_events.go
index 67a7d8a4..5c2a1de6 100644
--- a/roomserver/internal/input/input_latest_events.go
+++ b/roomserver/internal/input/input_latest_events.go
@@ -54,6 +54,7 @@ func (r *Inputer) updateLatestEvents(
 	event gomatrixserverlib.Event,
 	sendAsServer string,
 	transactionID *api.TransactionID,
+	rewritesState bool,
 ) (err error) {
 	updater, err := r.DB.GetLatestEventsForUpdate(ctx, *roomInfo)
 	if err != nil {
@@ -71,6 +72,7 @@ func (r *Inputer) updateLatestEvents(
 		event:         event,
 		sendAsServer:  sendAsServer,
 		transactionID: transactionID,
+		rewritesState: rewritesState,
 	}
 
 	if err = u.doUpdateLatestEvents(); err != nil {
@@ -93,6 +95,7 @@ type latestEventsUpdater struct {
 	stateAtEvent  types.StateAtEvent
 	event         gomatrixserverlib.Event
 	transactionID *api.TransactionID
+	rewritesState bool
 	// Which server to send this event as.
 	sendAsServer string
 	// The eventID of the event that was processed before this one.
@@ -178,7 +181,8 @@ func (u *latestEventsUpdater) doUpdateLatestEvents() error {
 		return fmt.Errorf("u.api.updateMemberships: %w", err)
 	}
 
-	update, err := u.makeOutputNewRoomEvent()
+	var update *api.OutputEvent
+	update, err = u.makeOutputNewRoomEvent()
 	if err != nil {
 		return fmt.Errorf("u.makeOutputNewRoomEvent: %w", err)
 	}
@@ -305,6 +309,7 @@ func (u *latestEventsUpdater) makeOutputNewRoomEvent() (*api.OutputEvent, error)
 
 	ore := api.OutputNewRoomEvent{
 		Event:           u.event.Headered(u.roomInfo.RoomVersion),
+		RewritesState:   u.rewritesState,
 		LastSentEventID: u.lastEventIDSent,
 		LatestEventIDs:  latestEventIDs,
 		TransactionID:   u.transactionID,
@@ -337,6 +342,11 @@ func (u *latestEventsUpdater) makeOutputNewRoomEvent() (*api.OutputEvent, error)
 			return nil, fmt.Errorf("failed to load add_state_events from db: %w", err)
 		}
 	}
+	// State is rewritten if the input room event HasState and we actually produced a delta on state events.
+	// Without this check, /get_missing_events which produce events with associated (but not complete) state
+	// will incorrectly purge the room and set it to no state. TODO: This is likely flakey, as if /gme produced
+	// a state conflict res which just so happens to include 2+ events we might purge the room state downstream.
+	ore.RewritesState = len(ore.AddsStateEventIDs) > 1
 
 	return &api.OutputEvent{
 		Type:         api.OutputTypeNewRoomEvent,
diff --git a/roomserver/internal/perform/perform_join.go b/roomserver/internal/perform/perform_join.go
index 3d194227..f76806c7 100644
--- a/roomserver/internal/perform/perform_join.go
+++ b/roomserver/internal/perform/perform_join.go
@@ -183,33 +183,33 @@ func (r *Joiner) performJoinRoomByID(
 		return "", fmt.Errorf("eb.SetContent: %w", err)
 	}
 
-	// First work out if this is in response to an existing invite
-	// from a federated server. If it is then we avoid the situation
-	// where we might think we know about a room in the following
-	// section but don't know the latest state as all of our users
-	// have left.
+	// Force a federated join if we aren't in the room and we've been
+	// given some server names to try joining by.
 	serverInRoom, _ := helpers.IsServerCurrentlyInRoom(ctx, r.DB, r.ServerName, req.RoomIDOrAlias)
+	forceFederatedJoin := len(req.ServerNames) > 0 && !serverInRoom
+
+	// Force a federated join if we're dealing with a pending invite
+	// and we aren't in the room.
 	isInvitePending, inviteSender, _, err := helpers.IsInvitePending(ctx, r.DB, req.RoomIDOrAlias, req.UserID)
-	if err == nil && isInvitePending && !serverInRoom {
-		// Check if there's an invite pending.
+	if err == nil && isInvitePending {
 		_, inviterDomain, ierr := gomatrixserverlib.SplitID('@', inviteSender)
 		if ierr != nil {
 			return "", fmt.Errorf("gomatrixserverlib.SplitID: %w", err)
 		}
 
-		// Check that the domain isn't ours. If it's local then we don't
-		// need to do anything as our own copy of the room state will be
-		// up-to-date.
+		// If we were invited by someone from another server then we can
+		// assume they are in the room so we can join via them.
 		if inviterDomain != r.Cfg.Matrix.ServerName {
-			// Add the server of the person who invited us to the server list,
-			// as they should be a fairly good bet.
 			req.ServerNames = append(req.ServerNames, inviterDomain)
-
-			// Perform a federated room join.
-			return req.RoomIDOrAlias, r.performFederatedJoinRoomByID(ctx, req)
+			forceFederatedJoin = true
 		}
 	}
 
+	// If we should do a forced federated join then do that.
+	if forceFederatedJoin {
+		return req.RoomIDOrAlias, r.performFederatedJoinRoomByID(ctx, req)
+	}
+
 	// Try to construct an actual join event from the template.
 	// If this succeeds then it is a sign that the room already exists
 	// locally on the homeserver.
diff --git a/roomserver/roomserver_test.go b/roomserver/roomserver_test.go
index 786d4f31..5a67a1be 100644
--- a/roomserver/roomserver_test.go
+++ b/roomserver/roomserver_test.go
@@ -1,12 +1,15 @@
 package roomserver
 
 import (
+	"bytes"
 	"context"
+	"crypto/ed25519"
 	"encoding/json"
 	"fmt"
 	"os"
 	"reflect"
 	"testing"
+	"time"
 
 	"github.com/Shopify/sarama"
 	"github.com/matrix-org/dendrite/internal/caching"
@@ -80,7 +83,73 @@ func deleteDatabase() {
 	}
 }
 
-func mustLoadEvents(t *testing.T, ver gomatrixserverlib.RoomVersion, events []json.RawMessage) []gomatrixserverlib.HeaderedEvent {
+type fledglingEvent struct {
+	Type     string
+	StateKey *string
+	Content  interface{}
+	Sender   string
+	RoomID   string
+}
+
+func mustCreateEvents(t *testing.T, roomVer gomatrixserverlib.RoomVersion, events []fledglingEvent) (result []gomatrixserverlib.HeaderedEvent) {
+	t.Helper()
+	depth := int64(1)
+	seed := make([]byte, ed25519.SeedSize) // zero seed
+	key := ed25519.NewKeyFromSeed(seed)
+	var prevs []string
+	roomState := make(map[gomatrixserverlib.StateKeyTuple]string) // state -> event ID
+	for _, ev := range events {
+		eb := gomatrixserverlib.EventBuilder{
+			Sender:     ev.Sender,
+			Depth:      depth,
+			Type:       ev.Type,
+			StateKey:   ev.StateKey,
+			RoomID:     ev.RoomID,
+			PrevEvents: prevs,
+		}
+		err := eb.SetContent(ev.Content)
+		if err != nil {
+			t.Fatalf("mustCreateEvent: failed to marshal event content %+v", ev.Content)
+		}
+		stateNeeded, err := gomatrixserverlib.StateNeededForEventBuilder(&eb)
+		if err != nil {
+			t.Fatalf("mustCreateEvent: failed to work out auth_events : %s", err)
+		}
+		var authEvents []string
+		for _, tuple := range stateNeeded.Tuples() {
+			eventID := roomState[tuple]
+			if eventID != "" {
+				authEvents = append(authEvents, eventID)
+			}
+		}
+		eb.AuthEvents = authEvents
+		signedEvent, err := eb.Build(time.Now(), testOrigin, "ed25519:test", key, roomVer)
+		if err != nil {
+			t.Fatalf("mustCreateEvent: failed to sign event: %s", err)
+		}
+		depth++
+		prevs = []string{signedEvent.EventID()}
+		if ev.StateKey != nil {
+			roomState[gomatrixserverlib.StateKeyTuple{
+				EventType: ev.Type,
+				StateKey:  *ev.StateKey,
+			}] = signedEvent.EventID()
+		}
+		result = append(result, signedEvent.Headered(roomVer))
+	}
+	return
+}
+
+func eventsJSON(events []gomatrixserverlib.Event) []json.RawMessage {
+	result := make([]json.RawMessage, len(events))
+	for i := range events {
+		result[i] = events[i].JSON()
+	}
+	return result
+}
+
+func mustLoadRawEvents(t *testing.T, ver gomatrixserverlib.RoomVersion, events []json.RawMessage) []gomatrixserverlib.HeaderedEvent {
+	t.Helper()
 	hs := make([]gomatrixserverlib.HeaderedEvent, len(events))
 	for i := range events {
 		e, err := gomatrixserverlib.NewEventFromTrustedJSON(events[i], false, ver)
@@ -93,7 +162,8 @@ func mustLoadEvents(t *testing.T, ver gomatrixserverlib.RoomVersion, events []js
 	return hs
 }
 
-func mustSendEvents(t *testing.T, ver gomatrixserverlib.RoomVersion, events []json.RawMessage) (api.RoomserverInternalAPI, *dummyProducer, []gomatrixserverlib.HeaderedEvent) {
+func mustCreateRoomserverAPI(t *testing.T) (api.RoomserverInternalAPI, *dummyProducer) {
+	t.Helper()
 	cfg := &config.Dendrite{}
 	cfg.Defaults()
 	cfg.Global.ServerName = testOrigin
@@ -112,9 +182,14 @@ func mustSendEvents(t *testing.T, ver gomatrixserverlib.RoomVersion, events []js
 		Cfg:           cfg,
 	}
 
-	rsAPI := NewInternalAPI(base, &test.NopJSONVerifier{})
-	hevents := mustLoadEvents(t, ver, events)
-	if err = api.SendEvents(ctx, rsAPI, hevents, testOrigin, nil); err != nil {
+	return NewInternalAPI(base, &test.NopJSONVerifier{}), dp
+}
+
+func mustSendEvents(t *testing.T, ver gomatrixserverlib.RoomVersion, events []json.RawMessage) (api.RoomserverInternalAPI, *dummyProducer, []gomatrixserverlib.HeaderedEvent) {
+	t.Helper()
+	rsAPI, dp := mustCreateRoomserverAPI(t)
+	hevents := mustLoadRawEvents(t, ver, events)
+	if err := api.SendEvents(ctx, rsAPI, hevents, testOrigin, nil); err != nil {
 		t.Errorf("failed to SendEvents: %s", err)
 	}
 	return rsAPI, dp, hevents
@@ -170,3 +245,163 @@ func TestOutputRedactedEvent(t *testing.T) {
 		}
 	}
 }
+
+// This tests that rewriting state via KindRewrite works correctly.
+// This creates a small room with a create/join/name state, then replays it
+// with a new room name. We expect the output events to contain the original events,
+// followed by a single OutputNewRoomEvent with RewritesState set to true with the
+// rewritten state events (with the 2nd room name).
+func TestOutputRewritesState(t *testing.T) {
+	roomID := "!foo:" + string(testOrigin)
+	alice := "@alice:" + string(testOrigin)
+	emptyKey := ""
+	originalEvents := mustCreateEvents(t, gomatrixserverlib.RoomVersionV6, []fledglingEvent{
+		{
+			RoomID: roomID,
+			Sender: alice,
+			Content: map[string]interface{}{
+				"creator":      alice,
+				"room_version": "6",
+			},
+			StateKey: &emptyKey,
+			Type:     gomatrixserverlib.MRoomCreate,
+		},
+		{
+			RoomID: roomID,
+			Sender: alice,
+			Content: map[string]interface{}{
+				"membership": "join",
+			},
+			StateKey: &alice,
+			Type:     gomatrixserverlib.MRoomMember,
+		},
+		{
+			RoomID: roomID,
+			Sender: alice,
+			Content: map[string]interface{}{
+				"body": "hello world",
+			},
+			StateKey: nil,
+			Type:     "m.room.message",
+		},
+		{
+			RoomID: roomID,
+			Sender: alice,
+			Content: map[string]interface{}{
+				"name": "Room Name",
+			},
+			StateKey: &emptyKey,
+			Type:     "m.room.name",
+		},
+	})
+	rewriteEvents := mustCreateEvents(t, gomatrixserverlib.RoomVersionV6, []fledglingEvent{
+		{
+			RoomID: roomID,
+			Sender: alice,
+			Content: map[string]interface{}{
+				"creator": alice,
+			},
+			StateKey: &emptyKey,
+			Type:     gomatrixserverlib.MRoomCreate,
+		},
+		{
+			RoomID: roomID,
+			Sender: alice,
+			Content: map[string]interface{}{
+				"membership": "join",
+			},
+			StateKey: &alice,
+			Type:     gomatrixserverlib.MRoomMember,
+		},
+		{
+			RoomID: roomID,
+			Sender: alice,
+			Content: map[string]interface{}{
+				"name": "Room Name 2",
+			},
+			StateKey: &emptyKey,
+			Type:     "m.room.name",
+		},
+		{
+			RoomID: roomID,
+			Sender: alice,
+			Content: map[string]interface{}{
+				"body": "hello world 2",
+			},
+			StateKey: nil,
+			Type:     "m.room.message",
+		},
+	})
+	deleteDatabase()
+	rsAPI, producer := mustCreateRoomserverAPI(t)
+	defer deleteDatabase()
+	err := api.SendEvents(context.Background(), rsAPI, originalEvents, testOrigin, nil)
+	if err != nil {
+		t.Fatalf("failed to send original events: %s", err)
+	}
+	// assert we got them produced, this is just a sanity check and isn't the intention of this test
+	if len(producer.producedMessages) != len(originalEvents) {
+		t.Fatalf("SendEvents didn't result in same number of produced output events: got %d want %d", len(producer.producedMessages), len(originalEvents))
+	}
+	producer.producedMessages = nil // we aren't actually interested in these events, just the rewrite ones
+
+	var inputEvents []api.InputRoomEvent
+	// slowly build up the state IDs again, we're basically telling the roomserver what to store as a snapshot
+	var stateIDs []string
+	// skip the last event, we'll use this to tie together the rewrite as the KindNew event
+	for i := 0; i < len(rewriteEvents)-1; i++ {
+		ev := rewriteEvents[i]
+		inputEvents = append(inputEvents, api.InputRoomEvent{
+			Kind:          api.KindRewrite,
+			Event:         ev,
+			AuthEventIDs:  ev.AuthEventIDs(),
+			HasState:      true,
+			StateEventIDs: stateIDs,
+		})
+		if ev.StateKey() != nil {
+			stateIDs = append(stateIDs, ev.EventID())
+		}
+	}
+	lastEv := rewriteEvents[len(rewriteEvents)-1]
+	inputEvents = append(inputEvents, api.InputRoomEvent{
+		Kind:          api.KindNew,
+		Event:         lastEv,
+		AuthEventIDs:  lastEv.AuthEventIDs(),
+		HasState:      true,
+		StateEventIDs: stateIDs,
+	})
+	if err := api.SendInputRoomEvents(context.Background(), rsAPI, inputEvents); err != nil {
+		t.Fatalf("SendInputRoomEvents returned error for rewrite events: %s", err)
+	}
+	// we should just have one output event with the entire state of the room in it
+	if len(producer.producedMessages) != 1 {
+		t.Fatalf("Rewritten events got output, want only 1 got %d", len(producer.producedMessages))
+	}
+	outputEvent := producer.producedMessages[0]
+	if !outputEvent.NewRoomEvent.RewritesState {
+		t.Errorf("RewritesState flag not set on output event")
+	}
+	if !reflect.DeepEqual(stateIDs, outputEvent.NewRoomEvent.AddsStateEventIDs) {
+		t.Errorf("Output event is missing room state event IDs, got %v want %v", outputEvent.NewRoomEvent.AddsStateEventIDs, stateIDs)
+	}
+	if !bytes.Equal(outputEvent.NewRoomEvent.Event.JSON(), lastEv.JSON()) {
+		t.Errorf(
+			"Output event isn't the latest KindNew event:\ngot  %s\nwant %s",
+			string(outputEvent.NewRoomEvent.Event.JSON()),
+			string(lastEv.JSON()),
+		)
+	}
+	if len(outputEvent.NewRoomEvent.AddStateEvents) != len(stateIDs) {
+		t.Errorf("Output event is missing room state events themselves, got %d want %d", len(outputEvent.NewRoomEvent.AddStateEvents), len(stateIDs))
+	}
+	// make sure the state got overwritten, check the room name
+	hasRoomName := false
+	for _, ev := range outputEvent.NewRoomEvent.AddStateEvents {
+		if ev.Type() == "m.room.name" {
+			hasRoomName = string(ev.Content()) == `{"name":"Room Name 2"}`
+		}
+	}
+	if !hasRoomName {
+		t.Errorf("Output event did not overwrite room state")
+	}
+}
diff --git a/roomserver/types/types.go b/roomserver/types/types.go
index 60f4b0fd..f5b45763 100644
--- a/roomserver/types/types.go
+++ b/roomserver/types/types.go
@@ -16,6 +16,8 @@
 package types
 
 import (
+	"sort"
+
 	"github.com/matrix-org/gomatrixserverlib"
 )
 
@@ -72,6 +74,25 @@ func (a StateEntry) LessThan(b StateEntry) bool {
 	return a.EventNID < b.EventNID
 }
 
+// Deduplicate takes a set of state entries and ensures that there are no
+// duplicate (event type, state key) tuples. If there are then we dedupe
+// them, making sure that the latest/highest NIDs are always chosen.
+func DeduplicateStateEntries(a []StateEntry) []StateEntry {
+	if len(a) < 2 {
+		return a
+	}
+	sort.SliceStable(a, func(i, j int) bool {
+		return a[i].LessThan(a[j])
+	})
+	for i := 0; i < len(a)-1; i++ {
+		if a[i].StateKeyTuple == a[i+1].StateKeyTuple {
+			a = append(a[:i], a[i+1:]...)
+			i--
+		}
+	}
+	return a
+}
+
 // StateAtEvent is the state before and after a matrix event.
 type StateAtEvent struct {
 	// Should this state overwrite the latest events and memberships of the room?
diff --git a/roomserver/types/types_test.go b/roomserver/types/types_test.go
new file mode 100644
index 00000000..b1e84b82
--- /dev/null
+++ b/roomserver/types/types_test.go
@@ -0,0 +1,26 @@
+package types
+
+import (
+	"testing"
+)
+
+func TestDeduplicateStateEntries(t *testing.T) {
+	entries := []StateEntry{
+		{StateKeyTuple{1, 1}, 1},
+		{StateKeyTuple{1, 1}, 2},
+		{StateKeyTuple{1, 1}, 3},
+		{StateKeyTuple{2, 2}, 4},
+		{StateKeyTuple{2, 3}, 5},
+		{StateKeyTuple{3, 3}, 6},
+	}
+	expected := []EventNID{3, 4, 5, 6}
+	entries = DeduplicateStateEntries(entries)
+	if len(entries) != 4 {
+		t.Fatalf("Expected 4 entries, got %d entries", len(entries))
+	}
+	for i, v := range entries {
+		if v.EventNID != expected[i] {
+			t.Fatalf("Expected position %d to be %d but got %d", i, expected[i], v.EventNID)
+		}
+	}
+}