From 04e879b581573b4cd4c77460ad561eddb6772304 Mon Sep 17 00:00:00 2001
From: Kegan Dougal <kegan@matrix.org>
Date: Tue, 24 Nov 2020 11:56:11 +0000
Subject: [PATCH] Add auth chain to federated response

---
 internal/mscs/msc2836/msc2836.go | 24 ++++++++++++++++++++++++
 roomserver/api/api.go            |  3 ++-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/internal/mscs/msc2836/msc2836.go b/internal/mscs/msc2836/msc2836.go
index 83e5e4d92..e19e7a9a6 100644
--- a/internal/mscs/msc2836/msc2836.go
+++ b/internal/mscs/msc2836/msc2836.go
@@ -198,6 +198,30 @@ func federatedEventRelationship(
 	if resErr != nil {
 		return *resErr
 	}
+	// add auth chain information
+	requiredAuthEventsSet := make(map[string]bool)
+	var requiredAuthEvents []string
+	for _, ev := range res.Events {
+		for _, a := range ev.AuthEventIDs() {
+			if requiredAuthEventsSet[a] {
+				continue
+			}
+			requiredAuthEvents = append(requiredAuthEvents, a)
+			requiredAuthEventsSet[a] = true
+		}
+	}
+	var queryRes roomserver.QueryAuthChainResponse
+	err = rsAPI.QueryAuthChain(ctx, &roomserver.QueryAuthChainRequest{
+		EventIDs: requiredAuthEvents,
+	}, &queryRes)
+	if err != nil {
+		// they may already have the auth events so don't fail this request
+		util.GetLogger(ctx).WithError(err).Error("Failed to QueryAuthChain")
+	}
+	res.AuthChain = make([]*gomatrixserverlib.Event, len(queryRes.AuthChain))
+	for i := range queryRes.AuthChain {
+		res.AuthChain[i] = queryRes.AuthChain[i].Unwrap()
+	}
 
 	return util.JSONResponse{
 		Code: 200,
diff --git a/roomserver/api/api.go b/roomserver/api/api.go
index fff61ed6f..200599cf2 100644
--- a/roomserver/api/api.go
+++ b/roomserver/api/api.go
@@ -127,7 +127,8 @@ type RoomserverInternalAPI interface {
 	) error
 
 	// QueryAuthChain returns the entire auth chain for the event IDs given.
-	// Omits without error any missing auth events.
+	// The response includes the events in the request.
+	// Omits without error for any missing auth events. There will be no duplicates.
 	QueryAuthChain(
 		ctx context.Context,
 		request *QueryAuthChainRequest,