Remove QueryBulkStateContent from current state server (#1404)

* Remove QueryBulkStateContent from current state server

Expected fail due to db impl not existing

* Implement query bulk state content

* Fix up rejecting invites over federation

* Fix bulk content marshalling

5 files changed, 164 insertions, 4 deletions

diff --git a/roomserver/api/query.go b/roomserver/api/query.go
index d0d0474d..67a217c8 100644
--- a/roomserver/api/query.go
+++ b/roomserver/api/query.go
@@ -303,6 +303,39 @@ type QueryServerBannedFromRoomResponse struct {
 	Banned bool `json:"banned"`
 }
 
+// MarshalJSON stringifies the room ID and StateKeyTuple keys so they can be sent over the wire in HTTP API mode.
+func (r *QueryBulkStateContentResponse) MarshalJSON() ([]byte, error) {
+	se := make(map[string]string)
+	for roomID, tupleToEvent := range r.Rooms {
+		for tuple, event := range tupleToEvent {
+			// use 0x1F (unit separator) as the delimiter between room ID/type/state key,
+			se[fmt.Sprintf("%s\x1F%s\x1F%s", roomID, tuple.EventType, tuple.StateKey)] = event
+		}
+	}
+	return json.Marshal(se)
+}
+
+func (r *QueryBulkStateContentResponse) UnmarshalJSON(data []byte) error {
+	wireFormat := make(map[string]string)
+	err := json.Unmarshal(data, &wireFormat)
+	if err != nil {
+		return err
+	}
+	r.Rooms = make(map[string]map[gomatrixserverlib.StateKeyTuple]string)
+	for roomTuple, value := range wireFormat {
+		fields := strings.Split(roomTuple, "\x1F")
+		roomID := fields[0]
+		if r.Rooms[roomID] == nil {
+			r.Rooms[roomID] = make(map[gomatrixserverlib.StateKeyTuple]string)
+		}
+		r.Rooms[roomID][gomatrixserverlib.StateKeyTuple{
+			EventType: fields[1],
+			StateKey:  fields[2],
+		}] = value
+	}
+	return nil
+}
+
 // MarshalJSON stringifies the StateKeyTuple keys so they can be sent over the wire in HTTP API mode.
 func (r *QueryCurrentStateResponse) MarshalJSON() ([]byte, error) {
 	se := make(map[string]*gomatrixserverlib.HeaderedEvent, len(r.StateEvents))
diff --git a/roomserver/internal/query/query.go b/roomserver/internal/query/query.go
index f76c9316..b34ae770 100644
--- a/roomserver/internal/query/query.go
+++ b/roomserver/internal/query/query.go
@@ -140,6 +140,9 @@ func (r *Queryer) QueryMembershipForUser(
 	if err != nil {
 		return err
 	}
+	if info == nil {
+		return fmt.Errorf("QueryMembershipForUser: unknown room %s", request.RoomID)
+	}
 
 	membershipEventNID, stillInRoom, err := r.DB.GetMembership(ctx, info.RoomNID, request.UserID)
 	if err != nil {
diff --git a/roomserver/storage/interface.go b/roomserver/storage/interface.go
index c4119f7e..be724da6 100644
--- a/roomserver/storage/interface.go
+++ b/roomserver/storage/interface.go
@@ -17,9 +17,9 @@ package storage
 import (
 	"context"
 
-	"github.com/matrix-org/dendrite/currentstateserver/storage/tables"
 	"github.com/matrix-org/dendrite/roomserver/api"
 	"github.com/matrix-org/dendrite/roomserver/storage/shared"
+	"github.com/matrix-org/dendrite/roomserver/storage/tables"
 	"github.com/matrix-org/dendrite/roomserver/types"
 	"github.com/matrix-org/gomatrixserverlib"
 )
diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go
index a081603f..5c18c725 100644
--- a/roomserver/storage/shared/storage.go
+++ b/roomserver/storage/shared/storage.go
@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"sort"
 
-	csstables "github.com/matrix-org/dendrite/currentstateserver/storage/tables"
 	"github.com/matrix-org/dendrite/internal/caching"
 	"github.com/matrix-org/dendrite/internal/sqlutil"
 	"github.com/matrix-org/dendrite/roomserver/api"
@@ -799,8 +798,90 @@ func (d *Database) GetRoomsByMembership(ctx context.Context, userID, membership
 
 // GetBulkStateContent returns all state events which match a given room ID and a given state key tuple. Both must be satisfied for a match.
 // If a tuple has the StateKey of '*' and allowWildcards=true then all state events with the EventType should be returned.
-func (d *Database) GetBulkStateContent(ctx context.Context, roomIDs []string, tuples []gomatrixserverlib.StateKeyTuple, allowWildcards bool) ([]csstables.StrippedEvent, error) {
-	return nil, fmt.Errorf("not implemented yet")
+// nolint:gocyclo
+func (d *Database) GetBulkStateContent(ctx context.Context, roomIDs []string, tuples []gomatrixserverlib.StateKeyTuple, allowWildcards bool) ([]tables.StrippedEvent, error) {
+	eventTypes := make([]string, 0, len(tuples))
+	for _, tuple := range tuples {
+		eventTypes = append(eventTypes, tuple.EventType)
+	}
+	// we don't bother failing the request if we get asked for event types we don't know about, as all that would result in is no matches which
+	// isn't a failure.
+	eventTypeNIDMap, err := d.EventTypesTable.BulkSelectEventTypeNID(ctx, eventTypes)
+	if err != nil {
+		return nil, fmt.Errorf("GetBulkStateContent: failed to map event type nids: %w", err)
+	}
+	typeNIDSet := make(map[types.EventTypeNID]bool)
+	for _, nid := range eventTypeNIDMap {
+		typeNIDSet[nid] = true
+	}
+
+	allowWildcard := make(map[types.EventTypeNID]bool)
+	eventStateKeys := make([]string, 0, len(tuples))
+	for _, tuple := range tuples {
+		if allowWildcards && tuple.StateKey == "*" {
+			allowWildcard[eventTypeNIDMap[tuple.EventType]] = true
+			continue
+		}
+		eventStateKeys = append(eventStateKeys, tuple.StateKey)
+
+	}
+
+	eventStateKeyNIDMap, err := d.EventStateKeysTable.BulkSelectEventStateKeyNID(ctx, eventStateKeys)
+	if err != nil {
+		return nil, fmt.Errorf("GetBulkStateContent: failed to map state key nids: %w", err)
+	}
+	stateKeyNIDSet := make(map[types.EventStateKeyNID]bool)
+	for _, nid := range eventStateKeyNIDMap {
+		stateKeyNIDSet[nid] = true
+	}
+
+	var eventNIDs []types.EventNID
+	eventNIDToVer := make(map[types.EventNID]gomatrixserverlib.RoomVersion)
+	// TODO: This feels like this is going to be really slow...
+	for _, roomID := range roomIDs {
+		roomInfo, err2 := d.RoomInfo(ctx, roomID)
+		if err2 != nil {
+			return nil, fmt.Errorf("GetBulkStateContent: failed to load room info for room %s : %w", roomID, err2)
+		}
+		// for unknown rooms or rooms which we don't have the current state, skip them.
+		if roomInfo == nil || roomInfo.IsStub {
+			continue
+		}
+		entries, err2 := d.loadStateAtSnapshot(ctx, roomInfo.StateSnapshotNID)
+		if err2 != nil {
+			return nil, fmt.Errorf("GetBulkStateContent: failed to load state for room %s : %w", roomID, err2)
+		}
+		for _, entry := range entries {
+			if typeNIDSet[entry.EventTypeNID] {
+				if allowWildcard[entry.EventTypeNID] || stateKeyNIDSet[entry.EventStateKeyNID] {
+					eventNIDs = append(eventNIDs, entry.EventNID)
+					eventNIDToVer[entry.EventNID] = roomInfo.RoomVersion
+				}
+			}
+		}
+	}
+
+	events, err := d.EventJSONTable.BulkSelectEventJSON(ctx, eventNIDs)
+	if err != nil {
+		return nil, fmt.Errorf("GetBulkStateContent: failed to load event JSON for event nids: %w", err)
+	}
+	result := make([]tables.StrippedEvent, len(events))
+	for i := range events {
+		roomVer := eventNIDToVer[events[i].EventNID]
+		ev, err := gomatrixserverlib.NewEventFromTrustedJSON(events[i].EventJSON, false, roomVer)
+		if err != nil {
+			return nil, fmt.Errorf("GetBulkStateContent: failed to load event JSON for event NID %v : %w", events[i].EventNID, err)
+		}
+		hev := ev.Headered(roomVer)
+		result[i] = tables.StrippedEvent{
+			EventType:    ev.Type(),
+			RoomID:       ev.RoomID(),
+			StateKey:     *ev.StateKey(),
+			ContentValue: tables.ExtractContentValue(&hev),
+		}
+	}
+
+	return result, nil
 }
 
 // JoinedUsersSetInRooms returns all joined users in the rooms given, along with the count of how many times they appear.
diff --git a/roomserver/storage/tables/interface.go b/roomserver/storage/tables/interface.go
index a142f2b1..adb06212 100644
--- a/roomserver/storage/tables/interface.go
+++ b/roomserver/storage/tables/interface.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/matrix-org/dendrite/roomserver/types"
 	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/tidwall/gjson"
 )
 
 type EventJSONPair struct {
@@ -155,3 +156,45 @@ type Redactions interface {
 	// successfully redacted the event JSON.
 	MarkRedactionValidated(ctx context.Context, txn *sql.Tx, redactionEventID string, validated bool) error
 }
+
+// StrippedEvent represents a stripped event for returning extracted content values.
+type StrippedEvent struct {
+	RoomID       string
+	EventType    string
+	StateKey     string
+	ContentValue string
+}
+
+// ExtractContentValue from the given state event. For example, given an m.room.name event with:
+//    content: { name: "Foo" }
+// this returns "Foo".
+func ExtractContentValue(ev *gomatrixserverlib.HeaderedEvent) string {
+	content := ev.Content()
+	key := ""
+	switch ev.Type() {
+	case gomatrixserverlib.MRoomCreate:
+		key = "creator"
+	case gomatrixserverlib.MRoomCanonicalAlias:
+		key = "alias"
+	case gomatrixserverlib.MRoomHistoryVisibility:
+		key = "history_visibility"
+	case gomatrixserverlib.MRoomJoinRules:
+		key = "join_rule"
+	case gomatrixserverlib.MRoomMember:
+		key = "membership"
+	case gomatrixserverlib.MRoomName:
+		key = "name"
+	case "m.room.avatar":
+		key = "url"
+	case "m.room.topic":
+		key = "topic"
+	case "m.room.guest_access":
+		key = "guest_access"
+	}
+	result := gjson.GetBytes(content, key)
+	if !result.Exists() {
+		return ""
+	}
+	// this returns the empty string if this is not a string type
+	return result.Str
+}