Move currentstateserver API to roomserver (#1387)

* Move currentstateserver API to roomserver

Stub out DB functions for now, nothing uses the roomserver version yet.

* Allow it to startup

* Implement some current-state-server storage interface functions

* Add missing package

4 files changed, 258 insertions, 0 deletions

diff --git a/roomserver/api/api.go b/roomserver/api/api.go
index 0fe30b8b..96bdc767 100644
--- a/roomserver/api/api.go
+++ b/roomserver/api/api.go
@@ -106,6 +106,20 @@ type RoomserverInternalAPI interface {
 		response *QueryStateAndAuthChainResponse,
 	) error
 
+	// QueryCurrentState retrieves the requested state events. If state events are not found, they will be missing from
+	// the response.
+	QueryCurrentState(ctx context.Context, req *QueryCurrentStateRequest, res *QueryCurrentStateResponse) error
+	// QueryRoomsForUser retrieves a list of room IDs matching the given query.
+	QueryRoomsForUser(ctx context.Context, req *QueryRoomsForUserRequest, res *QueryRoomsForUserResponse) error
+	// QueryBulkStateContent does a bulk query for state event content in the given rooms.
+	QueryBulkStateContent(ctx context.Context, req *QueryBulkStateContentRequest, res *QueryBulkStateContentResponse) error
+	// QuerySharedUsers returns a list of users who share at least 1 room in common with the given user.
+	QuerySharedUsers(ctx context.Context, req *QuerySharedUsersRequest, res *QuerySharedUsersResponse) error
+	// QueryKnownUsers returns a list of users that we know about from our joined rooms.
+	QueryKnownUsers(ctx context.Context, req *QueryKnownUsersRequest, res *QueryKnownUsersResponse) error
+	// QueryServerBannedFromRoom returns whether a server is banned from a room by server ACLs.
+	QueryServerBannedFromRoom(ctx context.Context, req *QueryServerBannedFromRoomRequest, res *QueryServerBannedFromRoomResponse) error
+
 	// Query a given amount (or less) of events prior to a given set of events.
 	PerformBackfill(
 		ctx context.Context,
diff --git a/roomserver/api/api_trace.go b/roomserver/api/api_trace.go
index 9b53aa88..25da2e8e 100644
--- a/roomserver/api/api_trace.go
+++ b/roomserver/api/api_trace.go
@@ -236,6 +236,47 @@ func (t *RoomserverInternalAPITrace) RemoveRoomAlias(
 	return err
 }
 
+func (t *RoomserverInternalAPITrace) QueryCurrentState(ctx context.Context, req *QueryCurrentStateRequest, res *QueryCurrentStateResponse) error {
+	err := t.Impl.QueryCurrentState(ctx, req, res)
+	util.GetLogger(ctx).WithError(err).Infof("QueryCurrentState req=%+v res=%+v", js(req), js(res))
+	return err
+}
+
+// QueryRoomsForUser retrieves a list of room IDs matching the given query.
+func (t *RoomserverInternalAPITrace) QueryRoomsForUser(ctx context.Context, req *QueryRoomsForUserRequest, res *QueryRoomsForUserResponse) error {
+	err := t.Impl.QueryRoomsForUser(ctx, req, res)
+	util.GetLogger(ctx).WithError(err).Infof("QueryRoomsForUser req=%+v res=%+v", js(req), js(res))
+	return err
+}
+
+// QueryBulkStateContent does a bulk query for state event content in the given rooms.
+func (t *RoomserverInternalAPITrace) QueryBulkStateContent(ctx context.Context, req *QueryBulkStateContentRequest, res *QueryBulkStateContentResponse) error {
+	err := t.Impl.QueryBulkStateContent(ctx, req, res)
+	util.GetLogger(ctx).WithError(err).Infof("QueryBulkStateContent req=%+v res=%+v", js(req), js(res))
+	return err
+}
+
+// QuerySharedUsers returns a list of users who share at least 1 room in common with the given user.
+func (t *RoomserverInternalAPITrace) QuerySharedUsers(ctx context.Context, req *QuerySharedUsersRequest, res *QuerySharedUsersResponse) error {
+	err := t.Impl.QuerySharedUsers(ctx, req, res)
+	util.GetLogger(ctx).WithError(err).Infof("QuerySharedUsers req=%+v res=%+v", js(req), js(res))
+	return err
+}
+
+// QueryKnownUsers returns a list of users that we know about from our joined rooms.
+func (t *RoomserverInternalAPITrace) QueryKnownUsers(ctx context.Context, req *QueryKnownUsersRequest, res *QueryKnownUsersResponse) error {
+	err := t.Impl.QueryKnownUsers(ctx, req, res)
+	util.GetLogger(ctx).WithError(err).Infof("QueryKnownUsers req=%+v res=%+v", js(req), js(res))
+	return err
+}
+
+// QueryServerBannedFromRoom returns whether a server is banned from a room by server ACLs.
+func (t *RoomserverInternalAPITrace) QueryServerBannedFromRoom(ctx context.Context, req *QueryServerBannedFromRoomRequest, res *QueryServerBannedFromRoomResponse) error {
+	err := t.Impl.QueryServerBannedFromRoom(ctx, req, res)
+	util.GetLogger(ctx).WithError(err).Infof("QueryServerBannedFromRoom req=%+v res=%+v", js(req), js(res))
+	return err
+}
+
 func js(thing interface{}) string {
 	b, err := json.Marshal(thing)
 	if err != nil {
diff --git a/roomserver/api/query.go b/roomserver/api/query.go
index 4e1d09c3..d0d0474d 100644
--- a/roomserver/api/query.go
+++ b/roomserver/api/query.go
@@ -17,6 +17,11 @@
 package api
 
 import (
+	"encoding/json"
+	"fmt"
+	"strings"
+
+	"github.com/matrix-org/dendrite/clientapi/auth/authtypes"
 	"github.com/matrix-org/gomatrixserverlib"
 )
 
@@ -225,3 +230,102 @@ type QueryPublishedRoomsResponse struct {
 	// The list of published rooms.
 	RoomIDs []string
 }
+
+type QuerySharedUsersRequest struct {
+	UserID         string
+	ExcludeRoomIDs []string
+	IncludeRoomIDs []string
+}
+
+type QuerySharedUsersResponse struct {
+	UserIDsToCount map[string]int
+}
+
+type QueryRoomsForUserRequest struct {
+	UserID string
+	// The desired membership of the user. If this is the empty string then no rooms are returned.
+	WantMembership string
+}
+
+type QueryRoomsForUserResponse struct {
+	RoomIDs []string
+}
+
+type QueryBulkStateContentRequest struct {
+	// Returns state events in these rooms
+	RoomIDs []string
+	// If true, treats the '*' StateKey as "all state events of this type" rather than a literal value of '*'
+	AllowWildcards bool
+	// The state events to return. Only a small subset of tuples are allowed in this request as only certain events
+	// have their content fields extracted. Specifically, the tuple Type must be one of:
+	//   m.room.avatar
+	//   m.room.create
+	//   m.room.canonical_alias
+	//   m.room.guest_access
+	//   m.room.history_visibility
+	//   m.room.join_rules
+	//   m.room.member
+	//   m.room.name
+	//   m.room.topic
+	// Any other tuple type will result in the query failing.
+	StateTuples []gomatrixserverlib.StateKeyTuple
+}
+type QueryBulkStateContentResponse struct {
+	// map of room ID -> tuple -> content_value
+	Rooms map[string]map[gomatrixserverlib.StateKeyTuple]string
+}
+
+type QueryCurrentStateRequest struct {
+	RoomID      string
+	StateTuples []gomatrixserverlib.StateKeyTuple
+}
+
+type QueryCurrentStateResponse struct {
+	StateEvents map[gomatrixserverlib.StateKeyTuple]*gomatrixserverlib.HeaderedEvent
+}
+
+type QueryKnownUsersRequest struct {
+	UserID       string `json:"user_id"`
+	SearchString string `json:"search_string"`
+	Limit        int    `json:"limit"`
+}
+
+type QueryKnownUsersResponse struct {
+	Users []authtypes.FullyQualifiedProfile `json:"profiles"`
+}
+
+type QueryServerBannedFromRoomRequest struct {
+	ServerName gomatrixserverlib.ServerName `json:"server_name"`
+	RoomID     string                       `json:"room_id"`
+}
+
+type QueryServerBannedFromRoomResponse struct {
+	Banned bool `json:"banned"`
+}
+
+// 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))
+	for k, v := range r.StateEvents {
+		// use 0x1F (unit separator) as the delimiter between type/state key,
+		se[fmt.Sprintf("%s\x1F%s", k.EventType, k.StateKey)] = v
+	}
+	return json.Marshal(se)
+}
+
+func (r *QueryCurrentStateResponse) UnmarshalJSON(data []byte) error {
+	res := make(map[string]*gomatrixserverlib.HeaderedEvent)
+	err := json.Unmarshal(data, &res)
+	if err != nil {
+		return err
+	}
+	r.StateEvents = make(map[gomatrixserverlib.StateKeyTuple]*gomatrixserverlib.HeaderedEvent, len(res))
+	for k, v := range res {
+		fields := strings.Split(k, "\x1F")
+		r.StateEvents[gomatrixserverlib.StateKeyTuple{
+			EventType: fields[0],
+			StateKey:  fields[1],
+		}] = v
+	}
+	return nil
+}
diff --git a/roomserver/api/wrapper.go b/roomserver/api/wrapper.go
index 16f5e8e1..82a4a571 100644
--- a/roomserver/api/wrapper.go
+++ b/roomserver/api/wrapper.go
@@ -133,3 +133,102 @@ func GetEvent(ctx context.Context, rsAPI RoomserverInternalAPI, eventID string)
 	}
 	return &res.Events[0]
 }
+
+// GetStateEvent returns the current state event in the room or nil.
+func GetStateEvent(ctx context.Context, rsAPI RoomserverInternalAPI, roomID string, tuple gomatrixserverlib.StateKeyTuple) *gomatrixserverlib.HeaderedEvent {
+	var res QueryCurrentStateResponse
+	err := rsAPI.QueryCurrentState(ctx, &QueryCurrentStateRequest{
+		RoomID:      roomID,
+		StateTuples: []gomatrixserverlib.StateKeyTuple{tuple},
+	}, &res)
+	if err != nil {
+		util.GetLogger(ctx).WithError(err).Error("Failed to QueryCurrentState")
+		return nil
+	}
+	ev, ok := res.StateEvents[tuple]
+	if ok {
+		return ev
+	}
+	return nil
+}
+
+// IsServerBannedFromRoom returns whether the server is banned from a room by server ACLs.
+func IsServerBannedFromRoom(ctx context.Context, rsAPI RoomserverInternalAPI, roomID string, serverName gomatrixserverlib.ServerName) bool {
+	req := &QueryServerBannedFromRoomRequest{
+		ServerName: serverName,
+		RoomID:     roomID,
+	}
+	res := &QueryServerBannedFromRoomResponse{}
+	if err := rsAPI.QueryServerBannedFromRoom(ctx, req, res); err != nil {
+		util.GetLogger(ctx).WithError(err).Error("Failed to QueryServerBannedFromRoom")
+		return true
+	}
+	return res.Banned
+}
+
+// PopulatePublicRooms extracts PublicRoom information for all the provided room IDs. The IDs are not checked to see if they are visible in the
+// published room directory.
+// due to lots of switches
+// nolint:gocyclo
+func PopulatePublicRooms(ctx context.Context, roomIDs []string, rsAPI RoomserverInternalAPI) ([]gomatrixserverlib.PublicRoom, error) {
+	avatarTuple := gomatrixserverlib.StateKeyTuple{EventType: "m.room.avatar", StateKey: ""}
+	nameTuple := gomatrixserverlib.StateKeyTuple{EventType: "m.room.name", StateKey: ""}
+	canonicalTuple := gomatrixserverlib.StateKeyTuple{EventType: gomatrixserverlib.MRoomCanonicalAlias, StateKey: ""}
+	topicTuple := gomatrixserverlib.StateKeyTuple{EventType: "m.room.topic", StateKey: ""}
+	guestTuple := gomatrixserverlib.StateKeyTuple{EventType: "m.room.guest_access", StateKey: ""}
+	visibilityTuple := gomatrixserverlib.StateKeyTuple{EventType: gomatrixserverlib.MRoomHistoryVisibility, StateKey: ""}
+	joinRuleTuple := gomatrixserverlib.StateKeyTuple{EventType: gomatrixserverlib.MRoomJoinRules, StateKey: ""}
+
+	var stateRes QueryBulkStateContentResponse
+	err := rsAPI.QueryBulkStateContent(ctx, &QueryBulkStateContentRequest{
+		RoomIDs:        roomIDs,
+		AllowWildcards: true,
+		StateTuples: []gomatrixserverlib.StateKeyTuple{
+			nameTuple, canonicalTuple, topicTuple, guestTuple, visibilityTuple, joinRuleTuple, avatarTuple,
+			{EventType: gomatrixserverlib.MRoomMember, StateKey: "*"},
+		},
+	}, &stateRes)
+	if err != nil {
+		util.GetLogger(ctx).WithError(err).Error("QueryBulkStateContent failed")
+		return nil, err
+	}
+	chunk := make([]gomatrixserverlib.PublicRoom, len(roomIDs))
+	i := 0
+	for roomID, data := range stateRes.Rooms {
+		pub := gomatrixserverlib.PublicRoom{
+			RoomID: roomID,
+		}
+		joinCount := 0
+		var joinRule, guestAccess string
+		for tuple, contentVal := range data {
+			if tuple.EventType == gomatrixserverlib.MRoomMember && contentVal == "join" {
+				joinCount++
+				continue
+			}
+			switch tuple {
+			case avatarTuple:
+				pub.AvatarURL = contentVal
+			case nameTuple:
+				pub.Name = contentVal
+			case topicTuple:
+				pub.Topic = contentVal
+			case canonicalTuple:
+				pub.CanonicalAlias = contentVal
+			case visibilityTuple:
+				pub.WorldReadable = contentVal == "world_readable"
+			// need both of these to determine whether guests can join
+			case joinRuleTuple:
+				joinRule = contentVal
+			case guestTuple:
+				guestAccess = contentVal
+			}
+		}
+		if joinRule == gomatrixserverlib.Public && guestAccess == "can_join" {
+			pub.GuestCanJoin = true
+		}
+		pub.JoinedMembersCount = joinCount
+		chunk[i] = pub
+		i++
+	}
+	return chunk, nil
+}