Implement query to get state and auth chain (#352)

* Implement query to get state and auth chain

* Add routing for queryStateAndAuthChain

* Comments

* Fix fetching wrong set of events

* Add tests

* Shuffle and comment

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

@@ -155,6 +155,33 @@ 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 prev events for the event. Used to calculate the state at
+	// the event
+	PrevEventIDs []string `json:"prev_event_ids"`
+	// The list of auth events for the event. Used to calculate the auth chain
+	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 +225,15 @@ type RoomserverQueryAPI interface {
 		request *QueryServerAllowedToSeeEventRequest,
 		response *QueryServerAllowedToSeeEventResponse,
 	) error
+
+	// Query to get state and auth chain for a (potentially hypothetical) event.
+	// Takes lists of PrevEventIDs and AuthEventsIDs and uses them to calculate
+	// the state and auth chain to return.
+	QueryStateAndAuthChain(
+		ctx context.Context,
+		request *QueryStateAndAuthChainRequest,
+		response *QueryStateAndAuthChainResponse,
+	) error
 }
 
 // RoomserverQueryLatestEventsAndStatePath is the HTTP path for the QueryLatestEventsAndState API.
@@ -218,6 +254,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 +349,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

@@ -27,9 +27,19 @@ import (
 	"github.com/matrix-org/util"
 )
 
+// RoomserverQueryAPIEventDB has a convenience API to fetch events directly by
+// EventIDs.
+type RoomserverQueryAPIEventDB interface {
+	// Look up the Events for a list of event IDs. Does not error if event was
+	// not found.
+	// Returns an error if the retrieval went wrong.
+	EventsFromIDs(ctx context.Context, eventIDs []string) ([]types.Event, error)
+}
+
 // RoomserverQueryAPIDatabase has the storage APIs needed to implement the query API.
 type RoomserverQueryAPIDatabase interface {
 	state.RoomStateDatabase
+	RoomserverQueryAPIEventDB
 	// Look up the numeric ID for the room.
 	// Returns 0 if the room doesn't exists.
 	// Returns an error if there was a problem talking to the database.
@@ -418,6 +428,98 @@ 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 = getAuthChain(ctx, r.DB, 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 getAuthChain(
+	ctx context.Context, dB RoomserverQueryAPIEventDB, 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)
+
+	// Check if there's anything left to do
+	for len(eventsToFetch) > 0 {
+		// Convert eventIDs to events. First need to fetch NIDs
+		events, err := dB.EventsFromIDs(ctx, eventsToFetch)
+		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 {
+			fetchedEventMap[event.EventID()] = true
+			authEvents = append(authEvents, event.Event)
+
+			// Now we need to fetch any auth events that we haven't
+			// previously seen.
+			for _, authEventID := range event.AuthEventIDs() {
+				if !fetchedEventMap[authEventID] {
+					fetchedEventMap[authEventID] = true
+					eventsToFetch = append(eventsToFetch, authEventID)
+				}
+			}
+		}
+	}
+
+	return authEvents, nil
+}
+
 // SetupHTTP adds the RoomserverQueryAPI handlers to the http.ServeMux.
 // nolint: gocyclo
 func (r *RoomserverQueryAPI) SetupHTTP(servMux *http.ServeMux) {
@@ -505,4 +607,18 @@ func (r *RoomserverQueryAPI) SetupHTTP(servMux *http.ServeMux) {
 			return util.JSONResponse{Code: 200, JSON: &response}
 		}),
 	)
+	servMux.Handle(
+		api.RoomserverQueryStateAndAuthChainPath,
+		common.MakeInternalAPI("queryStateAndAuthChain", func(req *http.Request) util.JSONResponse {
+			var request api.QueryStateAndAuthChainRequest
+			var response api.QueryStateAndAuthChainResponse
+			if err := json.NewDecoder(req.Body).Decode(&request); err != nil {
+				return util.ErrorResponse(err)
+			}
+			if err := r.QueryStateAndAuthChain(req.Context(), &request, &response); err != nil {
+				return util.ErrorResponse(err)
+			}
+			return util.JSONResponse{Code: 200, JSON: &response}
+		}),
+	)
 }

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

@@ -0,0 +1,174 @@
+// Copyright 2017 Vector Creations Ltd
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package query
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	"sort"
+
+	"github.com/matrix-org/dendrite/roomserver/types"
+	"github.com/matrix-org/gomatrixserverlib"
+)
+
+// used to implement RoomserverQueryAPIEventDB to test getAuthChain
+type getEventDB struct {
+	eventMap map[string]gomatrixserverlib.Event
+}
+
+func createEventDB() *getEventDB {
+	return &getEventDB{
+		eventMap: make(map[string]gomatrixserverlib.Event),
+	}
+}
+
+// Adds a fake event to the storage with given auth events.
+func (db *getEventDB) addFakeEvent(eventID string, authIDs []string) error {
+	authEvents := []gomatrixserverlib.EventReference{}
+	for _, authID := range authIDs {
+		authEvents = append(authEvents, gomatrixserverlib.EventReference{
+			EventID: authID,
+		})
+	}
+
+	builder := map[string]interface{}{
+		"event_id":    eventID,
+		"auth_events": authEvents,
+	}
+
+	eventJSON, err := json.Marshal(&builder)
+	if err != nil {
+		return err
+	}
+
+	event, err := gomatrixserverlib.NewEventFromTrustedJSON(eventJSON, false)
+	if err != nil {
+		return err
+	}
+
+	db.eventMap[eventID] = event
+
+	return nil
+}
+
+// Adds multiple events at once, each entry in the map is an eventID and set of
+// auth events that are converted to an event and added.
+func (db *getEventDB) addFakeEvents(graph map[string][]string) error {
+	for eventID, authIDs := range graph {
+		err := db.addFakeEvent(eventID, authIDs)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// EventsFromIDs implements RoomserverQueryAPIEventDB
+func (db *getEventDB) EventsFromIDs(ctx context.Context, eventIDs []string) (res []types.Event, err error) {
+	for _, evID := range eventIDs {
+		res = append(res, types.Event{
+			EventNID: 0,
+			Event:    db.eventMap[evID],
+		})
+	}
+
+	return
+}
+
+// Returns if the slices are equal after sorting them.
+func compareUnsortedStringSlices(a []string, b []string) bool {
+	if len(a) != len(b) {
+		return false
+	}
+
+	sort.Strings(a)
+	sort.Strings(b)
+
+	for i := range a {
+		if a[i] != b[i] {
+			return false
+		}
+	}
+
+	return true
+}
+
+func TestGetAuthChainSingle(t *testing.T) {
+	db := createEventDB()
+
+	err := db.addFakeEvents(map[string][]string{
+		"a": {},
+		"b": {"a"},
+		"c": {"a", "b"},
+		"d": {"b", "c"},
+		"e": {"a", "d"},
+	})
+
+	if err != nil {
+		t.Fatalf("Failed to add events to db: %v", err)
+	}
+
+	result, err := getAuthChain(context.TODO(), db, []string{"e"})
+	if err != nil {
+		t.Fatalf("getAuthChain failed: %v", err)
+	}
+
+	var returnedIDs []string
+	for _, event := range result {
+		returnedIDs = append(returnedIDs, event.EventID())
+	}
+
+	expectedIDs := []string{"a", "b", "c", "d", "e"}
+
+	if !compareUnsortedStringSlices(expectedIDs, returnedIDs) {
+		t.Fatalf("returnedIDs got '%v', expected '%v'", returnedIDs, expectedIDs)
+	}
+}
+
+func TestGetAuthChainMultiple(t *testing.T) {
+	db := createEventDB()
+
+	err := db.addFakeEvents(map[string][]string{
+		"a": {},
+		"b": {"a"},
+		"c": {"a", "b"},
+		"d": {"b", "c"},
+		"e": {"a", "d"},
+		"f": {"a", "b", "c"},
+	})
+
+	if err != nil {
+		t.Fatalf("Failed to add events to db: %v", err)
+	}
+
+	result, err := getAuthChain(context.TODO(), db, []string{"e", "f"})
+	if err != nil {
+		t.Fatalf("getAuthChain failed: %v", err)
+	}
+
+	var returnedIDs []string
+	for _, event := range result {
+		returnedIDs = append(returnedIDs, event.EventID())
+	}
+
+	expectedIDs := []string{"a", "b", "c", "d", "e", "f"}
+
+	if !compareUnsortedStringSlices(expectedIDs, returnedIDs) {
+		t.Fatalf("returnedIDs got '%v', expected '%v'", returnedIDs, expectedIDs)
+	}
+}

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

@@ -651,6 +651,21 @@ func (d *Database) GetMembershipEventNIDsForRoom(
 	return d.statements.selectMembershipsFromRoom(ctx, roomNID)
 }
 
+// EventsFromIDs implements query.RoomserverQueryAPIEventDB
+func (d *Database) EventsFromIDs(ctx context.Context, eventIDs []string) ([]types.Event, error) {
+	nidMap, err := d.EventNIDs(ctx, eventIDs)
+	if err != nil {
+		return nil, err
+	}
+
+	var nids []types.EventNID
+	for _, nid := range nidMap {
+		nids = append(nids, nid)
+	}
+
+	return d.Events(ctx, nids)
+}
+
 type transaction struct {
 	ctx context.Context
 	txn *sql.Tx