From 8e76523b04e3ebc9546f2b019e86dcd3b516be5a Mon Sep 17 00:00:00 2001
From: S7evinK <2353100+S7evinK@users.noreply.github.com>
Date: Thu, 24 Mar 2022 22:13:19 +0100
Subject: [PATCH] Update database when rejecting federated invite (#2300)

* Actually set the DB entry to "leave"

* Try to rollback in case of error

* Fix linter issue
---
 roomserver/internal/perform/perform_leave.go | 24 +++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/roomserver/internal/perform/perform_leave.go b/roomserver/internal/perform/perform_leave.go
index 1e5fb9f1..5b4cd3c6 100644
--- a/roomserver/internal/perform/perform_leave.go
+++ b/roomserver/internal/perform/perform_leave.go
@@ -212,12 +212,34 @@ func (r *Leaver) performFederatedRejectInvite(
 		ServerNames: []gomatrixserverlib.ServerName{domain},
 	}
 	leaveRes := fsAPI.PerformLeaveResponse{}
-	if err := r.FSAPI.PerformLeave(ctx, &leaveReq, &leaveRes); err != nil {
+	if err = r.FSAPI.PerformLeave(ctx, &leaveReq, &leaveRes); err != nil {
 		// failures in PerformLeave should NEVER stop us from telling other components like the
 		// sync API that the invite was withdrawn. Otherwise we can end up with stuck invites.
 		util.GetLogger(ctx).WithError(err).Errorf("failed to PerformLeave, still retiring invite event")
 	}
 
+	info, err := r.DB.RoomInfo(ctx, req.RoomID)
+	if err != nil {
+		util.GetLogger(ctx).WithError(err).Errorf("failed to get RoomInfo, still retiring invite event")
+	}
+
+	updater, err := r.DB.MembershipUpdater(ctx, req.RoomID, req.UserID, true, info.RoomVersion)
+	if err != nil {
+		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.Rollback(); err != nil {
+				util.GetLogger(ctx).WithError(err).Errorf("failed to rollback membership leave, 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")
+			}
+		}
+	}
+
 	// Withdraw the invite, so that the sync API etc are
 	// notified that we rejected it.
 	return []api.OutputEvent{