Warp/server/internal/roomserver/federated_query.go

package roomserver

import (
	"context"
	"encoding/json"
	"fmt"
	"slices"
	"sort"

	"github.com/dariuszjaniak/warp/server/internal/matrix/event"
	"github.com/dariuszjaniak/warp/server/internal/matrix/roomversion"
	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// The read side of the federation API: other servers asking about a room
// this one is in.
//
// Every read is answered only for a server that is in the room itself, or
// for a room anybody may read. A server outside the room asking for its
// events is asking for history its users were never shown, and the answer
// is the same 403 whether or not the event exists.

// Bounds on the federation reads.
const (
	maxBackfillLimit      = 100
	maxMissingEventsLimit = 20
)

// requireServerInRoom refuses a server that may not read a room.
func (r *API) requireServerInRoom(ctx context.Context, roomID, origin string) (storage.Room, roomversion.Rules, error) {
	if err := r.requireFederation("answering another server"); err != nil {
		return storage.Room{}, roomversion.Rules{}, err
	}
	room, rules, err := r.roomAndRules(ctx, roomID)
	if err != nil {
		return storage.Room{}, roomversion.Rules{}, err
	}
	if err := r.requireAllowed(ctx, roomID, origin); err != nil {
		return storage.Room{}, roomversion.Rules{}, err
	}
	in, err := r.ServerInRoom(ctx, roomID, origin)
	if err != nil {
		return storage.Room{}, roomversion.Rules{}, err
	}
	if in {
		return room, rules, nil
	}
	// A world-readable room may be read by anybody, servers included.
	if stored, err := r.store.CurrentStateEvent(ctx, roomID, event.TypeHistoryVisibility, ""); err == nil {
		var content struct {
			HistoryVisibility string `json:"history_visibility"`
		}
		if json.Unmarshal(contentOf(stored), &content) == nil && content.HistoryVisibility == HistoryVisibilityWorldReadable {
			return room, rules, nil
		}
	}
	return storage.Room{}, roomversion.Rules{}, fmt.Errorf("%w: %s is not in %s", ErrForbidden, origin, roomID)
}

// contentOf reads the content of a stored event without parsing it as a
// PDU, for the one property a visibility check needs.
func contentOf(stored storage.StoredEvent) json.RawMessage {
	var probe struct {
		Content json.RawMessage `json:"content"`
	}
	_ = json.Unmarshal(stored.JSON, &probe)
	return probe.Content
}

// EventForServer answers GET /event: one event, as JSON.
func (r *API) EventForServer(ctx context.Context, origin, eventID string) (json.RawMessage, error) {
	stored, err := r.store.EventByID(ctx, eventID)
	if err != nil {
		return nil, err
	}
	if _, _, err := r.requireServerInRoom(ctx, stored.RoomID, origin); err != nil {
		return nil, err
	}
	return stored.JSON, nil
}

// StateIDsForServer answers GET /state_ids: the state before an event and
// its auth chain, as event IDs.
func (r *API) StateIDsForServer(ctx context.Context, origin, roomID, eventID string) ([]string, []string, error) {
	_, rules, err := r.requireServerInRoom(ctx, roomID, origin)
	if err != nil {
		return nil, nil, err
	}
	stateIDs, err := r.stateIDsBefore(ctx, roomID, eventID)
	if err != nil {
		return nil, nil, err
	}
	chain, err := r.authChainIDs(ctx, roomID, rules, stateIDs)
	if err != nil {
		return nil, nil, err
	}
	return stateIDs, chain, nil
}

// StateForServer answers GET /state: StateIDsForServer with the events.
func (r *API) StateForServer(ctx context.Context, origin, roomID, eventID string) ([]json.RawMessage, []json.RawMessage, error) {
	stateIDs, chainIDs, err := r.StateIDsForServer(ctx, origin, roomID, eventID)
	if err != nil {
		return nil, nil, err
	}
	state, err := r.rawEvents(ctx, roomID, stateIDs)
	if err != nil {
		return nil, nil, err
	}
	chain, err := r.rawEvents(ctx, roomID, chainIDs)
	if err != nil {
		return nil, nil, err
	}
	return state, chain, nil
}

// EventAuthForServer answers GET /event_auth: an event's auth chain.
func (r *API) EventAuthForServer(ctx context.Context, origin, roomID, eventID string) ([]json.RawMessage, error) {
	_, rules, err := r.requireServerInRoom(ctx, roomID, origin)
	if err != nil {
		return nil, err
	}
	stored, err := r.store.Event(ctx, roomID, eventID)
	if err != nil {
		return nil, err
	}
	e, err := parseEvent(stored, rules)
	if err != nil {
		return nil, err
	}
	ids, err := r.authChainIDs(ctx, roomID, rules, e.AuthEventIDs())
	if err != nil {
		return nil, err
	}
	return r.rawEvents(ctx, roomID, ids)
}

// stateIDsBefore returns the event IDs of the state before an event — the
// snapshot it was authorised against.
func (r *API) stateIDsBefore(ctx context.Context, roomID, eventID string) ([]string, error) {
	stored, err := r.store.Event(ctx, roomID, eventID)
	if err != nil {
		return nil, err
	}
	if stored.Outlier {
		// An outlier has no snapshot; the honest answer is that this
		// server does not know the state at it.
		return nil, fmt.Errorf("%w: the state at %s is not known here", ErrNoSuchEvent, eventID)
	}
	entries, err := r.store.StateBefore(ctx, roomID, eventID)
	if err != nil {
		return nil, err
	}
	return stateEntryIDs(entries), nil
}

// authChainIDs walks auth_events from a set of events to the create event
// and returns every event on the way, the starting events excluded unless
// one of the others cites them.
func (r *API) authChainIDs(ctx context.Context, roomID string, rules roomversion.Rules, from []string) ([]string, error) {
	seen := map[string]bool{}
	frontier := map[string]bool{}
	start := map[string]*event.Event{}

	starting, err := r.eventsByID(ctx, roomID, rules, from)
	if err != nil {
		return nil, err
	}
	for _, e := range starting {
		start[e.EventID()] = e
		for _, id := range e.AuthEventIDs() {
			frontier[id] = true
		}
	}
	for _, id := range from {
		if _, ok := start[id]; !ok {
			// A starting ID that is not an event here is somebody's auth
			// event reference rather than an event of the chain's own.
			frontier[id] = true
		}
	}

	for len(frontier) > 0 {
		ids := make([]string, 0, len(frontier))
		for id := range frontier {
			if !seen[id] {
				seen[id] = true
				ids = append(ids, id)
			}
		}
		frontier = map[string]bool{}
		if len(ids) == 0 {
			break
		}
		events, err := r.eventsByID(ctx, roomID, rules, ids)
		if err != nil {
			return nil, err
		}
		for _, e := range events {
			for _, id := range e.AuthEventIDs() {
				if !seen[id] {
					frontier[id] = true
				}
			}
		}
	}

	// Only the IDs that are events here: a reference to something this
	// server never had is not part of the answer.
	all := make([]string, 0, len(seen))
	for id := range seen {
		all = append(all, id)
	}
	present, err := r.store.EventsByID(ctx, roomID, all)
	if err != nil {
		return nil, err
	}
	out := make([]string, 0, len(present))
	for _, s := range present {
		out = append(out, s.EventID)
	}
	sort.Strings(out)
	return out, nil
}

// rawEvents reads events as the JSON they are held as.
func (r *API) rawEvents(ctx context.Context, roomID string, ids []string) ([]json.RawMessage, error) {
	if len(ids) == 0 {
		return []json.RawMessage{}, nil
	}
	stored, err := r.store.EventsByID(ctx, roomID, ids)
	if err != nil {
		return nil, err
	}
	out := make([]json.RawMessage, 0, len(stored))
	for _, s := range stored {
		out = append(out, s.JSON)
	}
	return out, nil
}

// BackfillForServer answers GET /backfill: up to limit events reachable
// backwards through prev_events from the given ones, them included,
// newest first.
func (r *API) BackfillForServer(ctx context.Context, origin, roomID string, from []string, limit int) ([]json.RawMessage, error) {
	_, rules, err := r.requireServerInRoom(ctx, roomID, origin)
	if err != nil {
		return nil, err
	}
	if limit <= 0 || limit > maxBackfillLimit {
		limit = maxBackfillLimit
	}
	events, err := r.walkBack(ctx, roomID, rules, from, nil, 0, limit)
	if err != nil {
		return nil, err
	}
	sort.SliceStable(events, func(i, j int) bool { return events[i].Depth() > events[j].Depth() })
	return jsonOf(events), nil
}

// MissingEventsForServer answers POST /get_missing_events: the events
// between earliest and latest, oldest first.
func (r *API) MissingEventsForServer(ctx context.Context, origin, roomID string, earliest, latest []string, limit int, minDepth int64) ([]json.RawMessage, error) {
	_, rules, err := r.requireServerInRoom(ctx, roomID, origin)
	if err != nil {
		return nil, err
	}
	if limit <= 0 || limit > maxMissingEventsLimit {
		limit = maxMissingEventsLimit
	}
	starting, err := r.eventsByID(ctx, roomID, rules, latest)
	if err != nil {
		return nil, err
	}
	var from []string
	for _, e := range starting {
		from = append(from, e.PrevEventIDs()...)
	}
	events, err := r.walkBack(ctx, roomID, rules, from, earliest, minDepth, limit)
	if err != nil {
		return nil, err
	}
	eventsByDepth(events)
	return jsonOf(events), nil
}

// walkBack walks prev_events breadth first from a set of events, stopping
// at the ones in stop, below minDepth, or after limit events.
//
// Outliers are skipped. Their prev_events were never fetched, and handing
// one out would invite the asker to come back for history this server
// does not have.
func (r *API) walkBack(ctx context.Context, roomID string, rules roomversion.Rules, from, stop []string, minDepth int64, limit int) ([]*event.Event, error) {
	seen := map[string]bool{}
	for _, id := range stop {
		seen[id] = true
	}
	var out []*event.Event
	frontier := slices.Clone(from)
	for len(frontier) > 0 && len(out) < limit {
		var ids []string
		for _, id := range frontier {
			if !seen[id] {
				seen[id] = true
				ids = append(ids, id)
			}
		}
		frontier = nil
		if len(ids) == 0 {
			break
		}
		stored, err := r.store.EventsByID(ctx, roomID, ids)
		if err != nil {
			return nil, err
		}
		var level []*event.Event
		for _, s := range stored {
			if s.Outlier {
				continue
			}
			e, err := parseEvent(s, rules)
			if err != nil {
				return nil, err
			}
			if e.Depth() < minDepth {
				continue
			}
			level = append(level, e)
		}
		// Newest first within a level, so that the limit keeps the events
		// nearest to where the walk started.
		sort.SliceStable(level, func(i, j int) bool { return level[i].Depth() > level[j].Depth() })
		for _, e := range level {
			if len(out) >= limit {
				break
			}
			out = append(out, e)
			frontier = append(frontier, e.PrevEventIDs()...)
		}
	}
	return out, nil
}

func jsonOf(events []*event.Event) []json.RawMessage {
	out := make([]json.RawMessage, 0, len(events))
	for _, e := range events {
		out = append(out, e.JSON())
	}
	return out
}

// QueryAliasForServer answers GET /query/directory for an alias of this
// server.
func (r *API) QueryAliasForServer(ctx context.Context, alias string) (ResolvedAlias, error) {
	if err := r.requireFederation("answering another server"); err != nil {
		return ResolvedAlias{}, err
	}
	if !r.isLocalAlias(alias) {
		return ResolvedAlias{}, fmt.Errorf("%w: %s is not an alias of this server", ErrNoSuchAlias, alias)
	}
	stored, err := r.store.Alias(ctx, alias)
	if err != nil {
		return ResolvedAlias{}, err
	}
	servers, err := r.JoinedServers(ctx, stored.RoomID)
	if err != nil {
		return ResolvedAlias{}, err
	}
	if !slices.Contains(servers, r.serverName) {
		servers = append([]string{r.serverName}, servers...)
	}
	return ResolvedAlias{RoomID: stored.RoomID, Servers: servers}, nil
}