From cb5c62e5e571373778cebfd1fe89573b947a9b0c Mon Sep 17 00:00:00 2001
From: Kegan Dougal <kegan@matrix.org>
Date: Fri, 15 Jan 2021 15:37:48 +0000
Subject: [PATCH] Add exclude_rooms impl

---
 setup/mscs/msc2946/msc2946.go | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/setup/mscs/msc2946/msc2946.go b/setup/mscs/msc2946/msc2946.go
index 7572bca05..e99900e88 100644
--- a/setup/mscs/msc2946/msc2946.go
+++ b/setup/mscs/msc2946/msc2946.go
@@ -143,6 +143,15 @@ type walker struct {
 	mu                 sync.Mutex
 }
 
+func (w *walker) roomIsExcluded(roomID string) bool {
+	for _, exclRoom := range w.req.ExcludeRooms {
+		if exclRoom == roomID {
+			return true
+		}
+	}
+	return false
+}
+
 func (w *walker) alreadySent(id string) bool {
 	w.mu.Lock()
 	defer w.mu.Unlock()
@@ -195,7 +204,7 @@ func (w *walker) walk() *gomatrixserverlib.MSC2946SpacesResponse {
 
 		// If this room has not ever been in `rooms` (across multiple requests), extract the
 		// `PublicRoomsChunk` for this room.
-		if !w.alreadySent(roomID) {
+		if !w.alreadySent(roomID) && !w.roomIsExcluded(roomID) {
 			pubRoom := w.publicRoomsChunk(roomID)
 			roomType := ""
 			create := w.stateEvent(roomID, gomatrixserverlib.MRoomCreate, "")
@@ -209,6 +218,7 @@ func (w *walker) walk() *gomatrixserverlib.MSC2946SpacesResponse {
 				NumRefs:    refs.len(),
 				RoomType:   roomType,
 			})
+			w.markSent(roomID)
 		}
 
 		uniqueRooms := make(set)
@@ -239,6 +249,11 @@ func (w *walker) walk() *gomatrixserverlib.MSC2946SpacesResponse {
 				if w.alreadySent(ev.EventID()) {
 					continue
 				}
+				// Skip the room if it's part of exclude_rooms but ONLY IF the source matches, as we still
+				// want to catch arrows which point to excluded rooms.
+				if w.roomIsExcluded(ev.RoomID()) {
+					continue
+				}
 				res.Events = append(res.Events, ev.Event)
 				uniqueRooms[ev.RoomID()] = true
 				uniqueRooms[SpaceTarget(ev)] = true