Implement query to get state and auth chain

src/github.com/matrix-org/dendrite/roomserver/api/query.go

@@ -155,6 +155,32 @@ type QueryServerAllowedToSeeEventResponse struct {
 	AllowedToSeeEvent bool `json:"can_see_event"`
 }
 
+// QueryStateAndAuthChainRequest is a request to QueryStateAndAuthChain
+type QueryStateAndAuthChainRequest struct {
+	// The room ID to query the state in.
+	RoomID string `json:"room_id"`
+	// The list of events to return state after.
+	PrevEventIDs []string `json:"prev_event_ids"`
+	// The list of auth events to get the auth chain for.
+	AuthEventIDs []string `json:"auth_event_ids"`
+}
+
+// QueryStateAndAuthChainResponse is a response to QueryStateAndAuthChain
+type QueryStateAndAuthChainResponse struct {
+	// Copy of the request for debugging.
+	QueryStateAndAuthChainRequest
+	// Does the room exist on this roomserver?
+	// If the room doesn't exist this will be false and StateEvents will be empty.
+	RoomExists bool `json:"room_exists"`
+	// Do all the previous events exist on this roomserver?
+	// If some of previous events do not exist this will be false and StateEvents will be empty.
+	PrevEventsExist bool `json:"prev_events_exist"`
+	// The state and auth chain events that were requested.
+	// The lists will be in an arbitrary order.
+	StateEvents     []gomatrixserverlib.Event `json:"state_events"`
+	AuthChainEvents []gomatrixserverlib.Event `json:"auth_chain_events"`
+}
+
 // RoomserverQueryAPI is used to query information from the room server.
 type RoomserverQueryAPI interface {
 	// Query the latest events and state for a room from the room server.
@@ -198,6 +224,13 @@ type RoomserverQueryAPI interface {
 		request *QueryServerAllowedToSeeEventRequest,
 		response *QueryServerAllowedToSeeEventResponse,
 	) error
+
+	// Query to get state and auth chain
+	QueryStateAndAuthChain(
+		ctx context.Context,
+		request *QueryStateAndAuthChainRequest,
+		response *QueryStateAndAuthChainResponse,
+	) error
 }
 
 // RoomserverQueryLatestEventsAndStatePath is the HTTP path for the QueryLatestEventsAndState API.
@@ -218,6 +251,9 @@ const RoomserverQueryInvitesForUserPath = "/api/roomserver/queryInvitesForUser"
 // RoomserverQueryServerAllowedToSeeEventPath is the HTTP path for the QueryServerAllowedToSeeEvent API
 const RoomserverQueryServerAllowedToSeeEventPath = "/api/roomserver/queryServerAllowedToSeeEvent"
 
+// RoomserverQueryStateAndAuthChainPath is the HTTP path for the QueryStateAndAuthChain API
+const RoomserverQueryStateAndAuthChainPath = "/api/roomserver/queryStateAndAuthChain"
+
 // NewRoomserverQueryAPIHTTP creates a RoomserverQueryAPI implemented by talking to a HTTP POST API.
 // If httpClient is nil then it uses the http.DefaultClient
 func NewRoomserverQueryAPIHTTP(roomserverURL string, httpClient *http.Client) RoomserverQueryAPI {
@@ -310,6 +346,19 @@ func (h *httpRoomserverQueryAPI) QueryServerAllowedToSeeEvent(
 	return postJSON(ctx, span, h.httpClient, apiURL, request, response)
 }
 
+// QueryStateAndAuthChain implements RoomserverQueryAPI
+func (h *httpRoomserverQueryAPI) QueryStateAndAuthChain(
+	ctx context.Context,
+	request *QueryStateAndAuthChainRequest,
+	response *QueryStateAndAuthChainResponse,
+) error {
+	span, ctx := opentracing.StartSpanFromContext(ctx, "QueryStateAndAuthChain")
+	defer span.Finish()
+
+	apiURL := h.roomserverURL + RoomserverQueryStateAndAuthChainPath
+	return postJSON(ctx, span, h.httpClient, apiURL, request, response)
+}
+
 func postJSON(
 	ctx context.Context, span opentracing.Span, httpClient *http.Client,
 	apiURL string, request, response interface{},

src/github.com/matrix-org/dendrite/roomserver/query/query.go

@@ -418,6 +418,108 @@ func (r *RoomserverQueryAPI) QueryServerAllowedToSeeEvent(
 	return nil
 }
 
+// QueryStateAndAuthChain implements api.RoomserverQueryAPI
+func (r *RoomserverQueryAPI) QueryStateAndAuthChain(
+	ctx context.Context,
+	request *api.QueryStateAndAuthChainRequest,
+	response *api.QueryStateAndAuthChainResponse,
+) error {
+	response.QueryStateAndAuthChainRequest = *request
+	roomNID, err := r.DB.RoomNID(ctx, request.RoomID)
+	if err != nil {
+		return err
+	}
+	if roomNID == 0 {
+		return nil
+	}
+	response.RoomExists = true
+
+	prevStates, err := r.DB.StateAtEventIDs(ctx, request.PrevEventIDs)
+	if err != nil {
+		switch err.(type) {
+		case types.MissingEventError:
+			return nil
+		default:
+			return err
+		}
+	}
+	response.PrevEventsExist = true
+
+	// Look up the currrent state for the requested tuples.
+	stateEntries, err := state.LoadCombinedStateAfterEvents(
+		ctx, r.DB, prevStates,
+	)
+	if err != nil {
+		return err
+	}
+
+	stateEvents, err := r.loadStateEvents(ctx, stateEntries)
+	if err != nil {
+		return err
+	}
+
+	response.StateEvents = stateEvents
+	response.AuthChainEvents, err = r.getAuthChain(ctx, request.AuthEventIDs)
+	return err
+}
+
+// getAuthChain fetches the auth chain for the given auth events.
+// An auth chain is the list of all events that are referenced in the
+// auth_events section, and all their auth_events, recursively.
+// The returned set of events contain the given events.
+// Will *not* error if we don't have all auth events.
+func (r *RoomserverQueryAPI) getAuthChain(
+	ctx context.Context, authEventIDs []string,
+) ([]gomatrixserverlib.Event, error) {
+	var authEvents []gomatrixserverlib.Event
+
+	// List of event ids to fetch. These will be added to the result and
+	// their auth events will be fetched (if they haven't been previously)
+	eventsToFetch := authEventIDs
+
+	// Set of events we've already fetched.
+	fetchedEventMap := make(map[string]bool)
+	for _, eventID := range authEventIDs {
+		fetchedEventMap[eventID] = true
+	}
+
+	// Check if there's anything left to do
+	for len(eventsToFetch) > 0 {
+		// Convert eventIDs to events. First need to fetch NIDs
+		nidMap, err := r.DB.EventNIDs(ctx, authEventIDs)
+		if err != nil {
+			return nil, err
+		}
+
+		var nids []types.EventNID
+		for _, nid := range nidMap {
+			nids = append(nids, nid)
+		}
+
+		events, err := r.DB.Events(ctx, nids)
+		if err != nil {
+			return nil, err
+		}
+
+		// Work out a) which events we should add to the returned list of
+		// events and b) which of the auth events we haven't seen yet and
+		// add them to the list of events to fetch.
+		eventsToFetch = eventsToFetch[:0]
+		for _, event := range events {
+			for _, authEventID := range event.AuthEventIDs() {
+				if !fetchedEventMap[authEventID] {
+					fetchedEventMap[authEventID] = true
+					eventsToFetch = append(eventsToFetch, authEventID)
+				}
+			}
+
+			authEvents = append(authEvents, event.Event)
+		}
+	}
+
+	return authEvents, nil
+}
+
 // SetupHTTP adds the RoomserverQueryAPI handlers to the http.ServeMux.
 // nolint: gocyclo
 func (r *RoomserverQueryAPI) SetupHTTP(servMux *http.ServeMux) {