From 79a9eaa20c489555033ae0c2f34535e427aee5ca Mon Sep 17 00:00:00 2001
From: Till Faelligen <2353100+S7evinK@users.noreply.github.com>
Date: Tue, 11 Oct 2022 13:35:10 +0200
Subject: [PATCH] Make get_missing_events test pass

---
 roomserver/internal/helpers/helpers.go          | 15 ++++++++-------
 roomserver/internal/perform/perform_backfill.go |  5 ++++-
 roomserver/internal/query/query.go              |  6 ++++--
 3 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/roomserver/internal/helpers/helpers.go b/roomserver/internal/helpers/helpers.go
index 3b83a0a6d..a6de8ac84 100644
--- a/roomserver/internal/helpers/helpers.go
+++ b/roomserver/internal/helpers/helpers.go
@@ -324,7 +324,7 @@ func slowGetHistoryVisibilityState(
 func ScanEventTree(
 	ctx context.Context, db storage.Database, info *types.RoomInfo, front []string, visited map[string]bool, limit int,
 	serverName gomatrixserverlib.ServerName,
-) ([]types.EventNID, error) {
+) ([]types.EventNID, map[string]struct{}, error) {
 	var resultNIDs []types.EventNID
 	var err error
 	var allowed bool
@@ -345,6 +345,7 @@ func ScanEventTree(
 
 	var checkedServerInRoom bool
 	var isServerInRoom bool
+	redactEventIDs := make(map[string]struct{})
 
 	// Loop through the event IDs to retrieve the requested events and go
 	// through the whole tree (up to the provided limit) using the events'
@@ -358,7 +359,7 @@ BFSLoop:
 		// Retrieve the events to process from the database.
 		events, err = db.EventsFromIDs(ctx, front)
 		if err != nil {
-			return resultNIDs, err
+			return resultNIDs, redactEventIDs, err
 		}
 
 		if !checkedServerInRoom && len(events) > 0 {
@@ -395,16 +396,16 @@ BFSLoop:
 						)
 						// drop the error, as we will often error at the DB level if we don't have the prev_event itself. Let's
 						// just return what we have.
-						return resultNIDs, nil
+						return resultNIDs, redactEventIDs, nil
 					}
 
 					// If the event hasn't been seen before and the HS
 					// requesting to retrieve it is allowed to do so, add it to
 					// the list of events to retrieve.
-					if allowed {
-						next = append(next, pre)
-					} else {
+					next = append(next, pre)
+					if !allowed {
 						util.GetLogger(ctx).WithField("server", serverName).WithField("event_id", pre).Info("Not allowed to see event")
+						redactEventIDs[pre] = struct{}{}
 					}
 				}
 			}
@@ -413,7 +414,7 @@ BFSLoop:
 		front = next
 	}
 
-	return resultNIDs, err
+	return resultNIDs, redactEventIDs, err
 }
 
 func QueryLatestEventsAndState(
diff --git a/roomserver/internal/perform/perform_backfill.go b/roomserver/internal/perform/perform_backfill.go
index 69a075733..57e121ea2 100644
--- a/roomserver/internal/perform/perform_backfill.go
+++ b/roomserver/internal/perform/perform_backfill.go
@@ -78,7 +78,7 @@ func (r *Backfiller) PerformBackfill(
 	}
 
 	// Scan the event tree for events to send back.
-	resultNIDs, err := helpers.ScanEventTree(ctx, r.DB, info, front, visited, request.Limit, request.ServerName)
+	resultNIDs, redactEventIDs, err := helpers.ScanEventTree(ctx, r.DB, info, front, visited, request.Limit, request.ServerName)
 	if err != nil {
 		return err
 	}
@@ -95,6 +95,9 @@ func (r *Backfiller) PerformBackfill(
 	}
 
 	for _, event := range loadedEvents {
+		if _, ok := redactEventIDs[event.EventID()]; ok {
+			event.Redact()
+		}
 		response.Events = append(response.Events, event.Headered(info.RoomVersion))
 	}
 
diff --git a/roomserver/internal/query/query.go b/roomserver/internal/query/query.go
index 7a424a334..cbdd0aadc 100644
--- a/roomserver/internal/query/query.go
+++ b/roomserver/internal/query/query.go
@@ -453,7 +453,7 @@ func (r *Queryer) QueryMissingEvents(
 		return fmt.Errorf("missing RoomInfo for room %s", events[0].RoomID())
 	}
 
-	resultNIDs, err := helpers.ScanEventTree(ctx, r.DB, info, front, visited, request.Limit, request.ServerName)
+	resultNIDs, redact, err := helpers.ScanEventTree(ctx, r.DB, info, front, visited, request.Limit, request.ServerName)
 	if err != nil {
 		return err
 	}
@@ -470,7 +470,9 @@ func (r *Queryer) QueryMissingEvents(
 			if verr != nil {
 				return verr
 			}
-
+			if _, ok := redact[event.EventID()]; ok {
+				event.Redact()
+			}
 			response.Events = append(response.Events, event.Headered(roomVersion))
 		}
 	}