Add `/_dendrite/admin/downloadState/{serverName}/{roomID}` admin endpoint

6 files changed, 192 insertions, 12 deletions

diff --git a/roomserver/api/api.go b/roomserver/api/api.go
index 403bbe8b..a1373a62 100644
--- a/roomserver/api/api.go
+++ b/roomserver/api/api.go
@@ -150,6 +150,7 @@ type ClientRoomserverAPI interface {
 	PerformRoomUpgrade(ctx context.Context, req *PerformRoomUpgradeRequest, resp *PerformRoomUpgradeResponse) error
 	PerformAdminEvacuateRoom(ctx context.Context, req *PerformAdminEvacuateRoomRequest, res *PerformAdminEvacuateRoomResponse) error
 	PerformAdminEvacuateUser(ctx context.Context, req *PerformAdminEvacuateUserRequest, res *PerformAdminEvacuateUserResponse) error
+	PerformAdminDownloadState(ctx context.Context, req *PerformAdminDownloadStateRequest, res *PerformAdminDownloadStateResponse) error
 	PerformPeek(ctx context.Context, req *PerformPeekRequest, res *PerformPeekResponse) error
 	PerformUnpeek(ctx context.Context, req *PerformUnpeekRequest, res *PerformUnpeekResponse) error
 	PerformInvite(ctx context.Context, req *PerformInviteRequest, res *PerformInviteResponse) error
diff --git a/roomserver/api/api_trace.go b/roomserver/api/api_trace.go
index 8bef3537..342a3904 100644
--- a/roomserver/api/api_trace.go
+++ b/roomserver/api/api_trace.go
@@ -131,6 +131,16 @@ func (t *RoomserverInternalAPITrace) PerformAdminEvacuateUser(
 	return err
 }
 
+func (t *RoomserverInternalAPITrace) PerformAdminDownloadState(
+	ctx context.Context,
+	req *PerformAdminDownloadStateRequest,
+	res *PerformAdminDownloadStateResponse,
+) error {
+	err := t.Impl.PerformAdminDownloadState(ctx, req, res)
+	util.GetLogger(ctx).WithError(err).Infof("PerformAdminDownloadState req=%+v res=%+v", js(req), js(res))
+	return err
+}
+
 func (t *RoomserverInternalAPITrace) PerformInboundPeek(
 	ctx context.Context,
 	req *PerformInboundPeekRequest,
diff --git a/roomserver/api/perform.go b/roomserver/api/perform.go
index 1442a4b0..2536a4bb 100644
--- a/roomserver/api/perform.go
+++ b/roomserver/api/perform.go
@@ -237,3 +237,13 @@ type PerformAdminEvacuateUserResponse struct {
 	Affected []string `json:"affected"`
 	Error    *PerformError
 }
+
+type PerformAdminDownloadStateRequest struct {
+	RoomID     string                       `json:"room_id"`
+	UserID     string                       `json:"user_id"`
+	ServerName gomatrixserverlib.ServerName `json:"server_name"`
+}
+
+type PerformAdminDownloadStateResponse struct {
+	Error *PerformError `json:"error,omitempty"`
+}
diff --git a/roomserver/internal/perform/perform_admin.go b/roomserver/internal/perform/perform_admin.go
index 6a6d51b0..0406fc8f 100644
--- a/roomserver/internal/perform/perform_admin.go
+++ b/roomserver/internal/perform/perform_admin.go
@@ -236,3 +236,145 @@ func (r *Admin) PerformAdminEvacuateUser(
 	}
 	return nil
 }
+
+func (r *Admin) PerformAdminDownloadState(
+	ctx context.Context,
+	req *api.PerformAdminDownloadStateRequest,
+	res *api.PerformAdminDownloadStateResponse,
+) error {
+	roomInfo, err := r.DB.RoomInfo(ctx, req.RoomID)
+	if err != nil {
+		res.Error = &api.PerformError{
+			Code: api.PerformErrorBadRequest,
+			Msg:  fmt.Sprintf("r.DB.RoomInfo: %s", err),
+		}
+		return nil
+	}
+
+	if roomInfo == nil || roomInfo.IsStub() {
+		res.Error = &api.PerformError{
+			Code: api.PerformErrorBadRequest,
+			Msg:  fmt.Sprintf("room %q not found", req.RoomID),
+		}
+		return nil
+	}
+
+	fwdExtremities, _, depth, err := r.DB.LatestEventIDs(ctx, roomInfo.RoomNID)
+	if err != nil {
+		res.Error = &api.PerformError{
+			Code: api.PerformErrorBadRequest,
+			Msg:  fmt.Sprintf("r.DB.LatestEventIDs: %s", err),
+		}
+		return nil
+	}
+
+	authEventMap := map[string]*gomatrixserverlib.Event{}
+	stateEventMap := map[string]*gomatrixserverlib.Event{}
+
+	for _, fwdExtremity := range fwdExtremities {
+		var state gomatrixserverlib.RespState
+		state, err = r.Inputer.FSAPI.LookupState(ctx, req.ServerName, req.RoomID, fwdExtremity.EventID, roomInfo.RoomVersion)
+		if err != nil {
+			res.Error = &api.PerformError{
+				Code: api.PerformErrorBadRequest,
+				Msg:  fmt.Sprintf("r.Inputer.FSAPI.LookupState (%q): %s", fwdExtremity.EventID, err),
+			}
+			return nil
+		}
+		for _, authEvent := range state.AuthEvents.UntrustedEvents(roomInfo.RoomVersion) {
+			if err = authEvent.VerifyEventSignatures(ctx, r.Inputer.KeyRing); err != nil {
+				continue
+			}
+			authEventMap[authEvent.EventID()] = authEvent
+		}
+		for _, stateEvent := range state.StateEvents.UntrustedEvents(roomInfo.RoomVersion) {
+			if err = stateEvent.VerifyEventSignatures(ctx, r.Inputer.KeyRing); err != nil {
+				continue
+			}
+			stateEventMap[stateEvent.EventID()] = stateEvent
+		}
+	}
+
+	authEvents := make([]*gomatrixserverlib.HeaderedEvent, 0, len(authEventMap))
+	stateEvents := make([]*gomatrixserverlib.HeaderedEvent, 0, len(stateEventMap))
+	stateIDs := make([]string, 0, len(stateEventMap))
+
+	for _, authEvent := range authEventMap {
+		authEvents = append(authEvents, authEvent.Headered(roomInfo.RoomVersion))
+	}
+	for _, stateEvent := range stateEventMap {
+		stateEvents = append(stateEvents, stateEvent.Headered(roomInfo.RoomVersion))
+		stateIDs = append(stateIDs, stateEvent.EventID())
+	}
+
+	builder := &gomatrixserverlib.EventBuilder{
+		Type:    "org.matrix.dendrite.state_download",
+		Sender:  req.UserID,
+		RoomID:  req.RoomID,
+		Content: gomatrixserverlib.RawJSON("{}"),
+	}
+
+	eventsNeeded, err := gomatrixserverlib.StateNeededForEventBuilder(builder)
+	if err != nil {
+		res.Error = &api.PerformError{
+			Code: api.PerformErrorBadRequest,
+			Msg:  fmt.Sprintf("gomatrixserverlib.StateNeededForEventBuilder: %s", err),
+		}
+		return nil
+	}
+
+	queryRes := &api.QueryLatestEventsAndStateResponse{
+		RoomExists:   true,
+		RoomVersion:  roomInfo.RoomVersion,
+		LatestEvents: fwdExtremities,
+		StateEvents:  stateEvents,
+		Depth:        depth,
+	}
+
+	ev, err := eventutil.BuildEvent(ctx, builder, r.Cfg.Matrix, time.Now(), &eventsNeeded, queryRes)
+	if err != nil {
+		res.Error = &api.PerformError{
+			Code: api.PerformErrorBadRequest,
+			Msg:  fmt.Sprintf("eventutil.BuildEvent: %s", err),
+		}
+		return nil
+	}
+
+	inputReq := &api.InputRoomEventsRequest{
+		Asynchronous: false,
+	}
+	inputRes := &api.InputRoomEventsResponse{}
+
+	for _, authEvent := range append(authEvents, stateEvents...) {
+		inputReq.InputRoomEvents = append(inputReq.InputRoomEvents, api.InputRoomEvent{
+			Kind:  api.KindOutlier,
+			Event: authEvent,
+		})
+	}
+
+	inputReq.InputRoomEvents = append(inputReq.InputRoomEvents, api.InputRoomEvent{
+		Kind:          api.KindNew,
+		Event:         ev,
+		Origin:        r.Cfg.Matrix.ServerName,
+		HasState:      true,
+		StateEventIDs: stateIDs,
+		SendAsServer:  string(r.Cfg.Matrix.ServerName),
+	})
+
+	if err := r.Inputer.InputRoomEvents(ctx, inputReq, inputRes); err != nil {
+		res.Error = &api.PerformError{
+			Code: api.PerformErrorBadRequest,
+			Msg:  fmt.Sprintf("r.Inputer.InputRoomEvents: %s", err),
+		}
+		return nil
+	}
+
+	if inputRes.ErrMsg != "" {
+		res.Error = &api.PerformError{
+			Code: api.PerformErrorBadRequest,
+			Msg:  inputRes.ErrMsg,
+		}
+	}
+
+	return nil
+}
diff --git a/roomserver/inthttp/client.go b/roomserver/inthttp/client.go
index a1dfc6aa..1bd1b3fb 100644
--- a/roomserver/inthttp/client.go
+++ b/roomserver/inthttp/client.go
@@ -27,18 +27,19 @@ const (
 	RoomserverInputRoomEventsPath = "/roomserver/inputRoomEvents"
 
 	// Perform operations
-	RoomserverPerformInvitePath            = "/roomserver/performInvite"
-	RoomserverPerformPeekPath              = "/roomserver/performPeek"
-	RoomserverPerformUnpeekPath            = "/roomserver/performUnpeek"
-	RoomserverPerformRoomUpgradePath       = "/roomserver/performRoomUpgrade"
-	RoomserverPerformJoinPath              = "/roomserver/performJoin"
-	RoomserverPerformLeavePath             = "/roomserver/performLeave"
-	RoomserverPerformBackfillPath          = "/roomserver/performBackfill"
-	RoomserverPerformPublishPath           = "/roomserver/performPublish"
-	RoomserverPerformInboundPeekPath       = "/roomserver/performInboundPeek"
-	RoomserverPerformForgetPath            = "/roomserver/performForget"
-	RoomserverPerformAdminEvacuateRoomPath = "/roomserver/performAdminEvacuateRoom"
-	RoomserverPerformAdminEvacuateUserPath = "/roomserver/performAdminEvacuateUser"
+	RoomserverPerformInvitePath             = "/roomserver/performInvite"
+	RoomserverPerformPeekPath               = "/roomserver/performPeek"
+	RoomserverPerformUnpeekPath             = "/roomserver/performUnpeek"
+	RoomserverPerformRoomUpgradePath        = "/roomserver/performRoomUpgrade"
+	RoomserverPerformJoinPath               = "/roomserver/performJoin"
+	RoomserverPerformLeavePath              = "/roomserver/performLeave"
+	RoomserverPerformBackfillPath           = "/roomserver/performBackfill"
+	RoomserverPerformPublishPath            = "/roomserver/performPublish"
+	RoomserverPerformInboundPeekPath        = "/roomserver/performInboundPeek"
+	RoomserverPerformForgetPath             = "/roomserver/performForget"
+	RoomserverPerformAdminEvacuateRoomPath  = "/roomserver/performAdminEvacuateRoom"
+	RoomserverPerformAdminEvacuateUserPath  = "/roomserver/performAdminEvacuateUser"
+	RoomserverPerformAdminDownloadStatePath = "/roomserver/performAdminDownloadState"
 
 	// Query operations
 	RoomserverQueryLatestEventsAndStatePath    = "/roomserver/queryLatestEventsAndState"
@@ -261,6 +262,17 @@ func (h *httpRoomserverInternalAPI) PerformAdminEvacuateRoom(
 	)
 }
 
+func (h *httpRoomserverInternalAPI) PerformAdminDownloadState(
+	ctx context.Context,
+	request *api.PerformAdminDownloadStateRequest,
+	response *api.PerformAdminDownloadStateResponse,
+) error {
+	return httputil.CallInternalRPCAPI(
+		"PerformAdminDownloadState", h.roomserverURL+RoomserverPerformAdminDownloadStatePath,
+		h.httpClient, ctx, request, response,
+	)
+}
+
 func (h *httpRoomserverInternalAPI) PerformAdminEvacuateUser(
 	ctx context.Context,
 	request *api.PerformAdminEvacuateUserRequest,
diff --git a/roomserver/inthttp/server.go b/roomserver/inthttp/server.go
index 3b688174..4d37e90b 100644
--- a/roomserver/inthttp/server.go
+++ b/roomserver/inthttp/server.go
@@ -66,6 +66,11 @@ func AddRoutes(r api.RoomserverInternalAPI, internalAPIMux *mux.Router) {
 	)
 
 	internalAPIMux.Handle(
+		RoomserverPerformAdminDownloadStatePath,
+		httputil.MakeInternalRPCAPI("RoomserverPerformAdminDownloadState", r.PerformAdminDownloadState),
+	)
+
+	internalAPIMux.Handle(
 		RoomserverQueryPublishedRoomsPath,
 		httputil.MakeInternalRPCAPI("RoomserverQueryPublishedRooms", r.QueryPublishedRooms),
 	)