From 078aa09a5c297a354cacf0bb6b8d6354cb4bc3d1 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Wed, 3 Mar 2021 11:47:58 +0000
Subject: [PATCH] Update QueryMembershipsForRoom

---
 roomserver/api/query.go            |  4 ++++
 roomserver/internal/query/query.go | 23 +++++++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/roomserver/api/query.go b/roomserver/api/query.go
index 657f1a6dd..af35f7e72 100644
--- a/roomserver/api/query.go
+++ b/roomserver/api/query.go
@@ -151,6 +151,10 @@ type QueryMembershipsForRoomRequest struct {
 	JoinedOnly bool `json:"joined_only"`
 	// ID of the room to fetch memberships from
 	RoomID string `json:"room_id"`
+	// Optional - ID of the user sending the request, for checking if the
+	// user is allowed to see the memberships. If not specified then all
+	// room memberships will be returned.
+	Sender string `json:"sender"`
 }
 
 // QueryMembershipsForRoomResponse is a response to QueryMembershipsForRoom
diff --git a/roomserver/internal/query/query.go b/roomserver/internal/query/query.go
index 7ff8d84e4..ec8f7e1b5 100644
--- a/roomserver/internal/query/query.go
+++ b/roomserver/internal/query/query.go
@@ -243,6 +243,29 @@ func (r *Queryer) QueryMembershipsForRoom(
 		return err
 	}
 
+	if request.Sender == "" {
+		var events []types.Event
+		var eventNIDs []types.EventNID
+		eventNIDs, err = r.DB.GetMembershipEventNIDsForRoom(ctx, info.RoomNID, request.JoinedOnly, false)
+		if err != nil {
+			return fmt.Errorf("r.DB.GetMembershipEventNIDsForRoom: %w", err)
+		}
+		events, err = r.DB.Events(ctx, eventNIDs)
+		if err != nil {
+			return fmt.Errorf("r.DB.Events: %w", err)
+		}
+		for _, event := range events {
+			clientEvent := gomatrixserverlib.ToClientEvent(event.Event, gomatrixserverlib.FormatAll)
+			response.JoinEvents = append(response.JoinEvents, clientEvent)
+		}
+		return nil
+	}
+
+	membershipEventNID, stillInRoom, isRoomforgotten, err := r.DB.GetMembership(ctx, info.RoomNID, request.Sender)
+	if err != nil {
+		return err
+	}
+
 	response.IsRoomForgotten = isRoomforgotten
 
 	if membershipEventNID == 0 {