From 7ca0c189e172fd14ce1871c1af6d937420dca120 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Thu, 21 Jul 2022 11:53:09 +0100
Subject: [PATCH] Better invite handling

---
 federationapi/routing/invite.go               |  2 -
 roomserver/internal/perform/perform_invite.go | 57 ++++++++++++-------
 2 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/federationapi/routing/invite.go b/federationapi/routing/invite.go
index cde87a0ac..4b795018c 100644
--- a/federationapi/routing/invite.go
+++ b/federationapi/routing/invite.go
@@ -26,7 +26,6 @@ import (
 	"github.com/matrix-org/dendrite/setup/config"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/util"
-	"github.com/sirupsen/logrus"
 )
 
 // InviteV2 implements /_matrix/federation/v2/invite/{roomID}/{eventID}
@@ -144,7 +143,6 @@ func processInvite(
 	// Check that the event is signed by the server sending the request.
 	redacted, err := gomatrixserverlib.RedactEventJSON(event.JSON(), event.Version())
 	if err != nil {
-		logrus.WithError(err).Errorf("XXX: invite.go")
 		return util.JSONResponse{
 			Code: http.StatusBadRequest,
 			JSON: jsonerror.BadJSON("The event JSON could not be redacted"),
diff --git a/roomserver/internal/perform/perform_invite.go b/roomserver/internal/perform/perform_invite.go
index b7a5e274b..b2b495be8 100644
--- a/roomserver/internal/perform/perform_invite.go
+++ b/roomserver/internal/perform/perform_invite.go
@@ -24,6 +24,7 @@ import (
 	"github.com/matrix-org/dendrite/roomserver/internal/input"
 	"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/types"
 	"github.com/matrix-org/dendrite/setup/config"
 	"github.com/matrix-org/gomatrixserverlib"
@@ -38,6 +39,7 @@ type Inviter struct {
 	Inputer *input.Inputer
 }
 
+// nolint:gocyclo
 func (r *Inviter) PerformInvite(
 	ctx context.Context,
 	req *api.PerformInviteRequest,
@@ -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:  fmt.Sprintf("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())
@@ -146,7 +183,7 @@ func (r *Inviter) PerformInvite(
 	// about the invite so we can accept it, hence we return an output
 	// event to send to the Sync API.
 	if !isOriginLocal {
-		return outputUpdates, nil
+		return updateMembershipTableManually()
 	}
 
 	// The invite originated locally. Therefore we have a responsibility to
@@ -165,24 +202,6 @@ func (r *Inviter) PerformInvite(
 		return nil, nil
 	}
 
-	// Update the membership table for the invite.
-	updater, err := r.DB.MembershipUpdater(ctx, roomID, targetUserID, isTargetLocal, req.RoomVersion)
-	if err != nil {
-		return nil, fmt.Errorf("r.DB.MembershipUpdater: %w", err)
-	}
-	ev := &types.Event{
-		EventNID: 0,
-		Event:    event.Unwrap(),
-	}
-	outputUpdates, err = helpers.UpdateToInviteMembership(updater, ev, 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")
-
 	// If the invite originated from us and the target isn't local then we
 	// should try and send the invite over federation first. It might be
 	// that the remote user doesn't exist, in which case we can give up