Membership updater refactoring (#2541)

* Membership updater refactoring

* Pass in membership state

* Use membership check rather than referring to state directly

* Delete irrelevant membership states

* We don't need the leave event after all

* Tweaks

* Put a log entry in that I might stand a chance of finding

* Be less panicky

* Tweak invite handling

* Don't freak if we can't find the event NID

* Use event NID from `types.Event`

* Clean up

* Better invite handling

* Placate the almighty linter

* Blacklist a Sytest which is otherwise fine under Complement for reasons I don't understand

* Fix the sytest after all (thanks @S7evinK for the spot)

5 files changed, 117 insertions, 125 deletions

diff --git a/roomserver/internal/helpers/helpers.go b/roomserver/internal/helpers/helpers.go
index e67bbfca..2653027e 100644
--- a/roomserver/internal/helpers/helpers.go
+++ b/roomserver/internal/helpers/helpers.go
@@ -12,6 +12,7 @@ import (
 	"github.com/matrix-org/dendrite/roomserver/state"
 	"github.com/matrix-org/dendrite/roomserver/storage"
 	"github.com/matrix-org/dendrite/roomserver/storage/shared"
+	"github.com/matrix-org/dendrite/roomserver/storage/tables"
 	"github.com/matrix-org/dendrite/roomserver/types"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/util"
@@ -21,14 +22,14 @@ import (
 // Move these to a more sensible place.
 
 func UpdateToInviteMembership(
-	mu *shared.MembershipUpdater, add *gomatrixserverlib.Event, updates []api.OutputEvent,
+	mu *shared.MembershipUpdater, add *types.Event, updates []api.OutputEvent,
 	roomVersion gomatrixserverlib.RoomVersion,
 ) ([]api.OutputEvent, error) {
 	// We may have already sent the invite to the user, either because we are
 	// reprocessing this event, or because the we received this invite from a
 	// remote server via the federation invite API. In those cases we don't need
 	// to send the event.
-	needsSending, err := mu.SetToInvite(add)
+	needsSending, retired, err := mu.Update(tables.MembershipStateInvite, add)
 	if err != nil {
 		return nil, err
 	}
@@ -38,13 +39,23 @@ func UpdateToInviteMembership(
 		// room event stream. This ensures that the consumers only have to
 		// consider a single stream of events when determining whether a user
 		// is invited, rather than having to combine multiple streams themselves.
-		onie := api.OutputNewInviteEvent{
-			Event:       add.Headered(roomVersion),
-			RoomVersion: roomVersion,
-		}
 		updates = append(updates, api.OutputEvent{
-			Type:           api.OutputTypeNewInviteEvent,
-			NewInviteEvent: &onie,
+			Type: api.OutputTypeNewInviteEvent,
+			NewInviteEvent: &api.OutputNewInviteEvent{
+				Event:       add.Headered(roomVersion),
+				RoomVersion: roomVersion,
+			},
+		})
+	}
+	for _, eventID := range retired {
+		updates = append(updates, api.OutputEvent{
+			Type: api.OutputTypeRetireInviteEvent,
+			RetireInviteEvent: &api.OutputRetireInviteEvent{
+				EventID:          eventID,
+				Membership:       gomatrixserverlib.Join,
+				RetiredByEventID: add.EventID(),
+				TargetUserID:     *add.StateKey(),
+			},
 		})
 	}
 	return updates, nil
diff --git a/roomserver/internal/input/input_membership.go b/roomserver/internal/input/input_membership.go
index 3ce8791a..28a54623 100644
--- a/roomserver/internal/input/input_membership.go
+++ b/roomserver/internal/input/input_membership.go
@@ -21,6 +21,7 @@ import (
 	"github.com/matrix-org/dendrite/roomserver/api"
 	"github.com/matrix-org/dendrite/roomserver/internal/helpers"
 	"github.com/matrix-org/dendrite/roomserver/storage/shared"
+	"github.com/matrix-org/dendrite/roomserver/storage/tables"
 	"github.com/matrix-org/dendrite/roomserver/types"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/opentracing/opentracing-go"
@@ -60,20 +61,14 @@ func (r *Inputer) updateMemberships(
 	var updates []api.OutputEvent
 
 	for _, change := range changes {
-		var ae *gomatrixserverlib.Event
-		var re *gomatrixserverlib.Event
+		var ae *types.Event
+		var re *types.Event
 		targetUserNID := change.EventStateKeyNID
 		if change.removedEventNID != 0 {
-			ev, _ := helpers.EventMap(events).Lookup(change.removedEventNID)
-			if ev != nil {
-				re = ev.Event
-			}
+			re, _ = helpers.EventMap(events).Lookup(change.removedEventNID)
 		}
 		if change.addedEventNID != 0 {
-			ev, _ := helpers.EventMap(events).Lookup(change.addedEventNID)
-			if ev != nil {
-				ae = ev.Event
-			}
+			ae, _ = helpers.EventMap(events).Lookup(change.addedEventNID)
 		}
 		if updates, err = r.updateMembership(updater, targetUserNID, re, ae, updates); err != nil {
 			return nil, err
@@ -85,30 +80,27 @@ func (r *Inputer) updateMemberships(
 func (r *Inputer) updateMembership(
 	updater *shared.RoomUpdater,
 	targetUserNID types.EventStateKeyNID,
-	remove, add *gomatrixserverlib.Event,
+	remove, add *types.Event,
 	updates []api.OutputEvent,
 ) ([]api.OutputEvent, error) {
 	var err error
 	// Default the membership to Leave if no event was added or removed.
-	oldMembership := gomatrixserverlib.Leave
 	newMembership := gomatrixserverlib.Leave
-
-	if remove != nil {
-		oldMembership, err = remove.Membership()
-		if err != nil {
-			return nil, err
-		}
-	}
 	if add != nil {
 		newMembership, err = add.Membership()
 		if err != nil {
 			return nil, err
 		}
 	}
-	if oldMembership == newMembership && newMembership != gomatrixserverlib.Join {
-		// If the membership is the same then nothing changed and we can return
-		// immediately, unless it's a Join update (e.g. profile update).
-		return updates, nil
+
+	var targetLocal bool
+	if add != nil {
+		targetLocal = r.isLocalTarget(add)
+	}
+
+	mu, err := updater.MembershipUpdater(targetUserNID, targetLocal)
+	if err != nil {
+		return nil, err
 	}
 
 	// In an ideal world, we shouldn't ever have "add" be nil and "remove" be
@@ -120,17 +112,10 @@ func (r *Inputer) updateMembership(
 	// after a state reset, often thinking that the user was still joined to
 	// the room even though the room state said otherwise, and this would prevent
 	// the user from being able to attempt to rejoin the room without modifying
-	// the database. So instead what we'll do is we'll just update the membership
-	// table to say that the user is "leave" and we'll use the old event to
-	// avoid nil pointer exceptions on the code path that follows.
-	if add == nil {
-		add = remove
-		newMembership = gomatrixserverlib.Leave
-	}
-
-	mu, err := updater.MembershipUpdater(targetUserNID, r.isLocalTarget(add))
-	if err != nil {
-		return nil, err
+	// the database. So instead we're going to remove the membership from the
+	// database altogether, so that it doesn't create future problems.
+	if add == nil && remove != nil {
+		return nil, mu.Delete()
 	}
 
 	switch newMembership {
@@ -149,7 +134,7 @@ func (r *Inputer) updateMembership(
 	}
 }
 
-func (r *Inputer) isLocalTarget(event *gomatrixserverlib.Event) bool {
+func (r *Inputer) isLocalTarget(event *types.Event) bool {
 	isTargetLocalUser := false
 	if statekey := event.StateKey(); statekey != nil {
 		_, domain, _ := gomatrixserverlib.SplitID('@', *statekey)
@@ -159,81 +144,61 @@ func (r *Inputer) isLocalTarget(event *gomatrixserverlib.Event) bool {
 }
 
 func updateToJoinMembership(
-	mu *shared.MembershipUpdater, add *gomatrixserverlib.Event, updates []api.OutputEvent,
+	mu *shared.MembershipUpdater, add *types.Event, updates []api.OutputEvent,
 ) ([]api.OutputEvent, error) {
-	// If the user is already marked as being joined, we call SetToJoin to update
-	// the event ID then we can return immediately. Retired is ignored as there
-	// is no invite event to retire.
-	if mu.IsJoin() {
-		_, err := mu.SetToJoin(add.Sender(), add.EventID(), true)
-		if err != nil {
-			return nil, err
-		}
-		return updates, nil
-	}
 	// When we mark a user as being joined we will invalidate any invites that
 	// are active for that user. We notify the consumers that the invites have
 	// been retired using a special event, even though they could infer this
 	// by studying the state changes in the room event stream.
-	retired, err := mu.SetToJoin(add.Sender(), add.EventID(), false)
+	_, retired, err := mu.Update(tables.MembershipStateJoin, add)
 	if err != nil {
 		return nil, err
 	}
 	for _, eventID := range retired {
-		orie := api.OutputRetireInviteEvent{
-			EventID:          eventID,
-			Membership:       gomatrixserverlib.Join,
-			RetiredByEventID: add.EventID(),
-			TargetUserID:     *add.StateKey(),
-		}
 		updates = append(updates, api.OutputEvent{
-			Type:              api.OutputTypeRetireInviteEvent,
-			RetireInviteEvent: &orie,
+			Type: api.OutputTypeRetireInviteEvent,
+			RetireInviteEvent: &api.OutputRetireInviteEvent{
+				EventID:          eventID,
+				Membership:       gomatrixserverlib.Join,
+				RetiredByEventID: add.EventID(),
+				TargetUserID:     *add.StateKey(),
+			},
 		})
 	}
 	return updates, nil
 }
 
 func updateToLeaveMembership(
-	mu *shared.MembershipUpdater, add *gomatrixserverlib.Event,
+	mu *shared.MembershipUpdater, add *types.Event,
 	newMembership string, updates []api.OutputEvent,
 ) ([]api.OutputEvent, error) {
-	// If the user is already neither joined, nor invited to the room then we
-	// can return immediately.
-	if mu.IsLeave() {
-		return updates, nil
-	}
 	// When we mark a user as having left we will invalidate any invites that
 	// are active for that user. We notify the consumers that the invites have
 	// been retired using a special event, even though they could infer this
 	// by studying the state changes in the room event stream.
-	retired, err := mu.SetToLeave(add.Sender(), add.EventID())
+	_, retired, err := mu.Update(tables.MembershipStateLeaveOrBan, add)
 	if err != nil {
 		return nil, err
 	}
 	for _, eventID := range retired {
-		orie := api.OutputRetireInviteEvent{
-			EventID:          eventID,
-			Membership:       newMembership,
-			RetiredByEventID: add.EventID(),
-			TargetUserID:     *add.StateKey(),
-		}
 		updates = append(updates, api.OutputEvent{
-			Type:              api.OutputTypeRetireInviteEvent,
-			RetireInviteEvent: &orie,
+			Type: api.OutputTypeRetireInviteEvent,
+			RetireInviteEvent: &api.OutputRetireInviteEvent{
+				EventID:          eventID,
+				Membership:       newMembership,
+				RetiredByEventID: add.EventID(),
+				TargetUserID:     *add.StateKey(),
+			},
 		})
 	}
 	return updates, nil
 }
 
 func updateToKnockMembership(
-	mu *shared.MembershipUpdater, add *gomatrixserverlib.Event, updates []api.OutputEvent,
+	mu *shared.MembershipUpdater, add *types.Event, updates []api.OutputEvent,
 ) ([]api.OutputEvent, error) {
-	if mu.IsLeave() {
-		_, err := mu.SetToKnock(add)
-		if err != nil {
-			return nil, err
-		}
+	if _, _, err := mu.Update(tables.MembershipStateKnock, add); err != nil {
+		return nil, err
 	}
 	return updates, nil
 }
diff --git a/roomserver/internal/perform/perform_invite.go b/roomserver/internal/perform/perform_invite.go
index 644c954b..fce52930 100644
--- a/roomserver/internal/perform/perform_invite.go
+++ b/roomserver/internal/perform/perform_invite.go
@@ -39,11 +39,13 @@ type Inviter struct {
 	Inputer *input.Inputer
 }
 
+// nolint:gocyclo
 func (r *Inviter) PerformInvite(
 	ctx context.Context,
 	req *api.PerformInviteRequest,
 	res *api.PerformInviteResponse,
 ) ([]api.OutputEvent, error) {
+	var outputUpdates []api.OutputEvent
 	event := req.Event
 	if event.StateKey() == nil {
 		return nil, fmt.Errorf("invite must be a state event")
@@ -66,6 +68,13 @@ func (r *Inviter) PerformInvite(
 	}
 	isTargetLocal := domain == r.Cfg.Matrix.ServerName
 	isOriginLocal := event.Origin() == r.Cfg.Matrix.ServerName
+	if !isOriginLocal && !isTargetLocal {
+		res.Error = &api.PerformError{
+			Code: api.PerformErrorBadRequest,
+			Msg:  "The invite must be either from or to a local user",
+		}
+		return nil, nil
+	}
 
 	logger := util.GetLogger(ctx).WithFields(map[string]interface{}{
 		"inviter":  event.Sender(),
@@ -97,6 +106,34 @@ func (r *Inviter) PerformInvite(
 		}
 	}
 
+	updateMembershipTableManually := func() ([]api.OutputEvent, error) {
+		var updater *shared.MembershipUpdater
+		if updater, err = r.DB.MembershipUpdater(ctx, roomID, targetUserID, isTargetLocal, req.RoomVersion); err != nil {
+			return nil, fmt.Errorf("r.DB.MembershipUpdater: %w", err)
+		}
+		outputUpdates, err = helpers.UpdateToInviteMembership(updater, &types.Event{
+			EventNID: 0,
+			Event:    event.Unwrap(),
+		}, outputUpdates, req.Event.RoomVersion)
+		if err != nil {
+			return nil, fmt.Errorf("updateToInviteMembership: %w", err)
+		}
+		if err = updater.Commit(); err != nil {
+			return nil, fmt.Errorf("updater.Commit: %w", err)
+		}
+		logger.Debugf("updated membership to invite and sending invite OutputEvent")
+		return outputUpdates, nil
+	}
+
+	if (info == nil || info.IsStub) && !isOriginLocal && isTargetLocal {
+		// The invite came in over federation for a room that we don't know about
+		// yet. We need to handle this a bit differently to most invites because
+		// we don't know the room state, therefore the roomserver can't process
+		// an input event. Instead we will update the membership table with the
+		// new invite and generate an output event.
+		return updateMembershipTableManually()
+	}
+
 	var isAlreadyJoined bool
 	if info != nil {
 		_, isAlreadyJoined, _, err = r.DB.GetMembership(ctx, info.RoomNID, *event.StateKey())
@@ -140,31 +177,13 @@ func (r *Inviter) PerformInvite(
 		return nil, nil
 	}
 
+	// If the invite originated remotely then we can't send an
+	// InputRoomEvent for the invite as it will never pass auth checks
+	// due to lacking room state, but we still need to tell the client
+	// about the invite so we can accept it, hence we return an output
+	// event to send to the Sync API.
 	if !isOriginLocal {
-		// The invite originated over federation. Process the membership
-		// update, which will notify the sync API etc about the incoming
-		// invite. We do NOT send an InputRoomEvent for the invite as it
-		// will never pass auth checks due to lacking room state, but we
-		// still need to tell the client about the invite so we can accept
-		// it, hence we return an output event to send to the sync api.
-		var updater *shared.MembershipUpdater
-		updater, err = r.DB.MembershipUpdater(ctx, roomID, targetUserID, isTargetLocal, req.RoomVersion)
-		if err != nil {
-			return nil, fmt.Errorf("r.DB.MembershipUpdater: %w", err)
-		}
-
-		unwrapped := event.Unwrap()
-		var outputUpdates []api.OutputEvent
-		outputUpdates, err = helpers.UpdateToInviteMembership(updater, unwrapped, nil, req.Event.RoomVersion)
-		if err != nil {
-			return nil, fmt.Errorf("updateToInviteMembership: %w", err)
-		}
-
-		if err = updater.Commit(); err != nil {
-			return nil, fmt.Errorf("updater.Commit: %w", err)
-		}
-		logger.Debugf("updated membership to invite and sending invite OutputEvent")
-		return outputUpdates, nil
+		return updateMembershipTableManually()
 	}
 
 	// The invite originated locally. Therefore we have a responsibility to
@@ -229,12 +248,11 @@ func (r *Inviter) PerformInvite(
 			Code: api.PerformErrorNotAllowed,
 		}
 		logger.WithError(err).WithField("event_id", event.EventID()).Error("r.InputRoomEvents failed")
-		return nil, nil
 	}
 
 	// Don't notify the sync api of this event in the same way as a federated invite so the invitee
 	// gets the invite, as the roomserver will do this when it processes the m.room.member invite.
-	return nil, nil
+	return outputUpdates, nil
 }
 
 func buildInviteStrippedState(
diff --git a/roomserver/internal/perform/perform_join.go b/roomserver/internal/perform/perform_join.go
index c9e83919..1445b408 100644
--- a/roomserver/internal/perform/perform_join.go
+++ b/roomserver/internal/perform/perform_join.go
@@ -268,21 +268,19 @@ func (r *Joiner) performJoinRoomByID(
 	case nil:
 		// The room join is local. Send the new join event into the
 		// roomserver. First of all check that the user isn't already
-		// a member of the room.
-		alreadyJoined := false
-		for _, se := range buildRes.StateEvents {
-			if !se.StateKeyEquals(userID) {
-				continue
-			}
-			if membership, merr := se.Membership(); merr == nil {
-				alreadyJoined = (membership == gomatrixserverlib.Join)
-				break
-			}
+		// a member of the room. This is best-effort (as in we won't
+		// fail if we can't find the existing membership) because there
+		// is really no harm in just sending another membership event.
+		membershipReq := &api.QueryMembershipForUserRequest{
+			RoomID: req.RoomIDOrAlias,
+			UserID: userID,
 		}
+		membershipRes := &api.QueryMembershipForUserResponse{}
+		_ = r.Queryer.QueryMembershipForUser(ctx, membershipReq, membershipRes)
 
 		// If we haven't already joined the room then send an event
 		// into the room changing our membership status.
-		if !alreadyJoined {
+		if !membershipRes.RoomExists || !membershipRes.IsInRoom {
 			inputReq := rsAPI.InputRoomEventsRequest{
 				InputRoomEvents: []rsAPI.InputRoomEvent{
 					{
diff --git a/roomserver/internal/perform/perform_leave.go b/roomserver/internal/perform/perform_leave.go
index c5b62ac0..56e7240d 100644
--- a/roomserver/internal/perform/perform_leave.go
+++ b/roomserver/internal/perform/perform_leave.go
@@ -228,14 +228,14 @@ func (r *Leaver) performFederatedRejectInvite(
 		util.GetLogger(ctx).WithError(err).Errorf("failed to get MembershipUpdater, still retiring invite event")
 	}
 	if updater != nil {
-		if _, err = updater.SetToLeave(req.UserID, eventID); err != nil {
-			util.GetLogger(ctx).WithError(err).Errorf("failed to set membership to leave, still retiring invite event")
+		if err = updater.Delete(); err != nil {
+			util.GetLogger(ctx).WithError(err).Errorf("failed to delete membership, still retiring invite event")
 			if err = updater.Rollback(); err != nil {
-				util.GetLogger(ctx).WithError(err).Errorf("failed to rollback membership leave, still retiring invite event")
+				util.GetLogger(ctx).WithError(err).Errorf("failed to rollback deleting membership, still retiring invite event")
 			}
 		} else {
 			if err = updater.Commit(); err != nil {
-				util.GetLogger(ctx).WithError(err).Errorf("failed to commit membership update, still retiring invite event")
+				util.GetLogger(ctx).WithError(err).Errorf("failed to commit deleting membership, still retiring invite event")
 			}
 		}
 	}