Warp/server/internal/storage/sqlstore/rooms.go

package sqlstore

import (
	"context"
	"crypto/sha256"
	"database/sql"
	"encoding/base64"
	"errors"
	"fmt"
	"slices"
	"strings"

	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// Rooms is the shared SQL implementation of storage.RoomStore.
//
// Like the user store it is written once with `?` placeholders and rebound
// per driver. The statements stay inside the syntax both engines agree on:
// ON CONFLICT ... DO UPDATE with `excluded`, LIMIT/OFFSET, and no JSON
// functions at all — what a room's events mean is the roomserver's
// business, and the only reason this package ever looks inside one is that
// it does not.
type Rooms struct {
	db       *sql.DB
	rebind   func(string) string
	isUnique func(error) bool
}

// NewRooms builds a room store.
func NewRooms(db *sql.DB, rebind func(string) string, isUnique func(error) bool) *Rooms {
	if rebind == nil {
		rebind = func(query string) string { return query }
	}
	if isUnique == nil {
		isUnique = func(error) bool { return false }
	}
	return &Rooms{db: db, rebind: rebind, isUnique: isUnique}
}

// Page sizes for a room listing. The public room directory is reachable
// without a token on some deployments, so an unbounded ?limit= there would
// be a way to ask the server to serialise every room it knows.
const (
	defaultRoomPageSize = 100
	maxRoomPageSize     = 1000
)

// The counters in roomserver_sequences.
const (
	seqStreamOrdering = "stream_ordering"
	seqStateSnapshot  = "state_snapshot"

	// seqBackfillOrdering numbers backfilled history. It starts at 0 and
	// only ever goes down; see WriteBackfill.
	seqBackfillOrdering = "backfill_ordering"
)

// q rebinds a statement for this driver.
func (s *Rooms) q(query string) string { return s.rebind(query) }

// tx runs fn in a transaction, rolling back on any error.
func (s *Rooms) tx(ctx context.Context, fn func(*sql.Tx) error) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin transaction: %w", err)
	}
	defer func() { _ = tx.Rollback() }()

	if err := fn(tx); err != nil {
		return err
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("commit transaction: %w", err)
	}
	return nil
}

const insertRoomSQL = `
	INSERT INTO roomserver_rooms (room_id, room_version, creator, is_published, created_at_ms)
	VALUES (?, ?, ?, ?, ?)`

// CreateRoom implements storage.RoomStore.
func (s *Rooms) CreateRoom(ctx context.Context, room storage.NewRoom) error {
	_, err := s.db.ExecContext(ctx, s.q(insertRoomSQL),
		room.ID, room.Version, room.Creator, room.Published, room.CreatedAtMS)
	if err != nil {
		if s.isUnique(err) {
			return fmt.Errorf("create room %q: %w", room.ID, storage.ErrRoomExists)
		}
		return fmt.Errorf("create room %q: %w", room.ID, err)
	}
	return nil
}

const selectRoomSQL = `
	SELECT room_id, room_version, creator, is_published, created_at_ms
	FROM roomserver_rooms WHERE room_id = ?`

// Room implements storage.RoomStore.
func (s *Rooms) Room(ctx context.Context, roomID string) (storage.Room, error) {
	var room storage.Room
	err := s.db.QueryRowContext(ctx, s.q(selectRoomSQL), roomID).Scan(
		&room.ID, &room.Version, &room.Creator, &room.Published, &room.CreatedAtMS)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return storage.Room{}, fmt.Errorf("room %q: %w", roomID, storage.ErrNoSuchRoom)
	case err != nil:
		return storage.Room{}, fmt.Errorf("read room %q: %w", roomID, err)
	}
	return room, nil
}

// Write implements storage.RoomStore.
//
// The order inside the transaction is the order the pieces depend on each
// other: the stream position first, because the event row needs it; the
// state snapshot next, because the event row points at it; then the event,
// then everything that refers to the event by ID.
func (s *Rooms) Write(ctx context.Context, w storage.RoomWrite) (int64, error) {
	var ordering int64
	err := s.tx(ctx, func(tx *sql.Tx) error {
		var err error
		if ordering, err = s.nextSequence(ctx, tx, seqStreamOrdering); err != nil {
			return err
		}

		// An outlier has no state snapshot: that is what makes it one.
		var snapshotID int64
		if !w.Outlier {
			if snapshotID, err = s.stateSnapshot(ctx, tx, w.Event.RoomID, w.StateBefore); err != nil {
				return err
			}
		}

		if err := s.insertEvent(ctx, tx, w.Event, ordering, snapshotID, w.Outlier); err != nil {
			return err
		}
		if w.Event.Type == memberEventType && w.Event.StateKey != nil {
			if err := s.replaceInviteState(ctx, tx, w.Event.RoomID, *w.Event.StateKey, w.InviteState); err != nil {
				return err
			}
		}
		if w.Redacts != "" {
			if err := s.applyRedaction(ctx, tx, w.Event.RoomID, w.Redacts,
				w.Event.EventID, w.RedactedJSON); err != nil {
				return err
			}
		}
		if err := s.applyStateDelta(ctx, tx, w); err != nil {
			return err
		}
		if err := s.replaceExtremities(ctx, tx, w.Event.RoomID, w.Extremities); err != nil {
			return err
		}
		if w.Profile != nil {
			if err := s.updateProfile(ctx, tx, w.Event.RoomID, *w.Profile); err != nil {
				return err
			}
		}
		// The member count is recomputed whenever current state moved,
		// because that is the only thing that can change it.
		if len(w.StateDelta) > 0 || w.StateDeltaIsWholeState {
			if err := s.recountMembers(ctx, tx, w.Event.RoomID); err != nil {
				return err
			}
		}
		if w.Transaction != nil {
			if err := s.recordTransaction(ctx, tx, *w.Transaction,
				w.Event.EventID, w.Event.RoomID, w.Event.OriginServerTS); err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		return 0, fmt.Errorf("write event %s: %w", w.Event.EventID, err)
	}
	return ordering, nil
}

// nextSequence hands out the next value of one of this component's
// counters; see sequences.go for why they are rows rather than sequences.
func (s *Rooms) nextSequence(ctx context.Context, tx *sql.Tx, name string) (int64, error) {
	return nextSequence(ctx, tx, s.q, roomSequences, name)
}

const (
	selectSnapshotSQL = `
		SELECT snapshot_id FROM roomserver_state_snapshots
		WHERE room_id = ? AND state_hash = ?`
	insertSnapshotSQL = `
		INSERT INTO roomserver_state_snapshots (snapshot_id, room_id, state_hash)
		VALUES (?, ?, ?)`
	insertStateEntrySQL = `
		INSERT INTO roomserver_state_entries (snapshot_id, event_type, state_key, event_id)
		VALUES (?, ?, ?, ?)`
)

// stateSnapshot stores a state set, re-using the row of an identical one.
//
// Deduplication is what makes a per-event state snapshot affordable: every
// message in a room whose state is not changing resolves to the snapshot
// the previous one already stored, so the table grows with state changes
// rather than with traffic. An empty state set is snapshot 0 and stores
// nothing, which is the create event's case.
func (s *Rooms) stateSnapshot(ctx context.Context, tx *sql.Tx, roomID string, entries []storage.StateEntry) (int64, error) {
	if len(entries) == 0 {
		return 0, nil
	}
	hash := stateHash(entries)

	var id int64
	err := tx.QueryRowContext(ctx, s.q(selectSnapshotSQL), roomID, hash).Scan(&id)
	switch {
	case err == nil:
		return id, nil
	case !errors.Is(err, sql.ErrNoRows):
		return 0, fmt.Errorf("look up state snapshot: %w", err)
	}

	if id, err = s.nextSequence(ctx, tx, seqStateSnapshot); err != nil {
		return 0, err
	}
	if _, err := tx.ExecContext(ctx, s.q(insertSnapshotSQL), id, roomID, hash); err != nil {
		return 0, fmt.Errorf("insert state snapshot: %w", err)
	}
	for _, e := range entries {
		if _, err := tx.ExecContext(ctx, s.q(insertStateEntrySQL),
			id, e.Type, e.StateKey, e.EventID); err != nil {
			return 0, fmt.Errorf("insert state entry %s/%s: %w", e.Type, e.StateKey, err)
		}
	}
	return id, nil
}

// stateHash identifies a state set by its contents.
//
// The entries are sorted first so that the same state always hashes the
// same whatever order it was assembled in, and each field is written with
// its length so that no two different sets can be run together into the
// same bytes.
func stateHash(entries []storage.StateEntry) string {
	sorted := slices.Clone(entries)
	slices.SortFunc(sorted, func(a, b storage.StateEntry) int {
		if c := strings.Compare(a.Type, b.Type); c != 0 {
			return c
		}
		return strings.Compare(a.StateKey, b.StateKey)
	})

	h := sha256.New()
	for _, e := range sorted {
		for _, field := range [...]string{e.Type, e.StateKey, e.EventID} {
			fmt.Fprintf(h, "%d:%s", len(field), field)
		}
	}
	return base64.RawURLEncoding.EncodeToString(h.Sum(nil))
}

// memberEventType is the one event type the store has to recognise: its
// rows carry a membership, and its arrival clears a stored invite.
const memberEventType = "m.room.member"

const (
	insertEventSQL = `
		INSERT INTO roomserver_events (
			event_id, room_id, stream_ordering, event_type, state_key,
			sender, depth, origin_server_ts, event_json, state_before_id, outlier)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`
	selectOutlierSQL  = `SELECT outlier FROM roomserver_events WHERE event_id = ?`
	promoteOutlierSQL = `
		UPDATE roomserver_events
		SET stream_ordering = ?, state_before_id = ?, outlier = ?
		WHERE event_id = ? AND outlier = ?`
)

// insertEvent writes the event row.
//
// An event this server already holds is not written twice. The one
// exception is an outlier arriving again as an ordinary event, which is
// promoted in place: it gets the new stream position and the state
// snapshot it now has, and joins the timeline. Its JSON is left as it was —
// the event ID says the two copies are the same event, and the stored one
// may already have been redacted.
func (s *Rooms) insertEvent(ctx context.Context, tx *sql.Tx, e storage.StoredEvent, ordering, snapshotID int64, outlier bool) error {
	var existingOutlier bool
	err := tx.QueryRowContext(ctx, s.q(selectOutlierSQL), e.EventID).Scan(&existingOutlier)
	switch {
	case err == nil:
		if existingOutlier && !outlier {
			if _, err := tx.ExecContext(ctx, s.q(promoteOutlierSQL),
				ordering, snapshotID, false, e.EventID, true); err != nil {
				return fmt.Errorf("promote outlier: %w", err)
			}
		}
		// The same event twice is not a failure worth propagating: an
		// event ID is a hash of its contents, so the second copy says
		// exactly what the first one did.
		return nil
	case !errors.Is(err, sql.ErrNoRows):
		return fmt.Errorf("look up event: %w", err)
	}

	_, err = tx.ExecContext(ctx, s.q(insertEventSQL),
		e.EventID, e.RoomID, ordering, e.Type, e.StateKey,
		e.Sender, e.Depth, e.OriginServerTS, string(e.JSON), snapshotID, outlier)
	if err != nil {
		if s.isUnique(err) {
			return nil
		}
		return fmt.Errorf("insert event: %w", err)
	}
	return nil
}

const (
	deleteInviteStateSQL = `DELETE FROM roomserver_invite_state WHERE room_id = ? AND user_id = ?`
	insertInviteStateSQL = `
		INSERT INTO roomserver_invite_state (room_id, user_id, state_json) VALUES (?, ?, ?)`
)

// replaceInviteState clears a user's stored invite state and, for an
// invite that came with some, stores the new one.
func (s *Rooms) replaceInviteState(ctx context.Context, tx *sql.Tx, roomID, userID string, state []byte) error {
	if _, err := tx.ExecContext(ctx, s.q(deleteInviteStateSQL), roomID, userID); err != nil {
		return fmt.Errorf("clear invite state: %w", err)
	}
	if len(state) == 0 {
		return nil
	}
	if _, err := tx.ExecContext(ctx, s.q(insertInviteStateSQL), roomID, userID, string(state)); err != nil {
		return fmt.Errorf("store invite state: %w", err)
	}
	return nil
}

const selectInviteStateSQL = `
	SELECT state_json FROM roomserver_invite_state WHERE room_id = ? AND user_id = ?`

// InviteState implements storage.RoomStore.
func (s *Rooms) InviteState(ctx context.Context, roomID, userID string) ([]byte, error) {
	var state string
	err := s.db.QueryRowContext(ctx, s.q(selectInviteStateSQL), roomID, userID).Scan(&state)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return nil, nil
	case err != nil:
		return nil, fmt.Errorf("read the invite state of %q in %q: %w", userID, roomID, err)
	}
	return []byte(state), nil
}

// WriteOutliers implements storage.RoomStore.
//
// The stream positions are taken as one block, in one UPDATE, rather than
// one at a time: a join to a large room hands over tens of thousands of
// events, and a counter bumped per event would be that many statements
// holding the counter's row lock.
func (s *Rooms) WriteOutliers(ctx context.Context, roomID string, events []storage.StoredEvent) error {
	if len(events) == 0 {
		return nil
	}
	err := s.tx(ctx, func(tx *sql.Tx) error {
		top, err := advanceSequence(ctx, tx, s.q, roomSequences, seqStreamOrdering, int64(len(events)))
		if err != nil {
			return err
		}
		first := top - int64(len(events)) + 1
		for i, e := range events {
			if e.RoomID == "" {
				e.RoomID = roomID
			}
			if err := s.insertEvent(ctx, tx, e, first+int64(i), 0, true); err != nil {
				return fmt.Errorf("outlier %s: %w", e.EventID, err)
			}
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("write %d outliers to %q: %w", len(events), roomID, err)
	}
	return nil
}

// WriteBackfill implements storage.RoomStore.
//
// The positions come from their own counter, advanced downwards by the
// number of events actually written, in one UPDATE: the events already in
// the timeline are sorted out first, inside the transaction, so that the
// block taken is exactly as big as what goes into it. A hole below the
// room would be harmless — nothing streams through backfilled history, it
// is only ever read by bounds — but there is no reason to leave one.
//
// The batch takes the block oldest first: if the counter stood at -10 and
// three events are written, they become -13, -12 and -11.
func (s *Rooms) WriteBackfill(ctx context.Context, roomID string, events []storage.BackfillEvent) error {
	if len(events) == 0 {
		return nil
	}
	err := s.tx(ctx, func(tx *sql.Tx) error {
		// Which of the batch is new, which is an outlier to promote, and
		// which is already in the timeline and left alone.
		var write []storage.BackfillEvent
		promote := map[string]bool{}
		for _, b := range events {
			var outlier bool
			err := tx.QueryRowContext(ctx, s.q(selectOutlierSQL), b.Event.EventID).Scan(&outlier)
			switch {
			case errors.Is(err, sql.ErrNoRows):
				write = append(write, b)
			case err != nil:
				return fmt.Errorf("look up event %s: %w", b.Event.EventID, err)
			case outlier:
				write = append(write, b)
				promote[b.Event.EventID] = true
			}
		}
		if len(write) == 0 {
			return nil
		}

		bottom, err := advanceSequence(ctx, tx, s.q, roomSequences, seqBackfillOrdering, -int64(len(write)))
		if err != nil {
			return err
		}
		for i, b := range write {
			e := b.Event
			if e.RoomID == "" {
				e.RoomID = roomID
			}
			ordering := bottom + int64(i)
			snapshotID, err := s.stateSnapshot(ctx, tx, e.RoomID, b.StateBefore)
			if err != nil {
				return err
			}
			if promote[e.EventID] {
				if _, err := tx.ExecContext(ctx, s.q(promoteOutlierSQL),
					ordering, snapshotID, false, e.EventID, true); err != nil {
					return fmt.Errorf("promote outlier %s: %w", e.EventID, err)
				}
				continue
			}
			if _, err := tx.ExecContext(ctx, s.q(insertEventSQL),
				e.EventID, e.RoomID, ordering, e.Type, e.StateKey,
				e.Sender, e.Depth, e.OriginServerTS, string(e.JSON), snapshotID, false); err != nil {
				return fmt.Errorf("insert event %s: %w", e.EventID, err)
			}
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("backfill %d events into %q: %w", len(events), roomID, err)
	}
	return nil
}

const applyRedactionSQL = `
	UPDATE roomserver_events SET event_json = ?, redacted_by = ?
	WHERE room_id = ? AND event_id = ? AND redacted_by = ''`

// applyRedaction replaces a stored event with its redacted form.
//
// The condition on redacted_by makes a second redaction of the same event a
// no-op rather than a rewrite: the first one already stripped the content,
// and the event may only be redacted once.
func (s *Rooms) applyRedaction(ctx context.Context, tx *sql.Tx, roomID, targetID, redactionID string, redacted []byte) error {
	_, err := tx.ExecContext(ctx, s.q(applyRedactionSQL),
		string(redacted), redactionID, roomID, targetID)
	if err != nil {
		return fmt.Errorf("apply redaction to %s: %w", targetID, err)
	}
	return nil
}

const (
	upsertCurrentStateSQL = `
		INSERT INTO roomserver_current_state (room_id, event_type, state_key, event_id, membership)
		VALUES (?, ?, ?, ?, ?)
		ON CONFLICT (room_id, event_type, state_key) DO UPDATE
			SET event_id = excluded.event_id, membership = excluded.membership`
	deleteCurrentStateSQL = `DELETE FROM roomserver_current_state WHERE room_id = ?`
)

// applyStateDelta moves the room's current state.
//
// The usual case is a single tuple to upsert. A whole-state write — which
// only state resolution produces — clears the table for the room first,
// because tuples the resolution dropped have to go.
func (s *Rooms) applyStateDelta(ctx context.Context, tx *sql.Tx, w storage.RoomWrite) error {
	if len(w.StateDelta) == 0 && !w.StateDeltaIsWholeState {
		return nil
	}
	if w.StateDeltaIsWholeState {
		if _, err := tx.ExecContext(ctx, s.q(deleteCurrentStateSQL), w.Event.RoomID); err != nil {
			return fmt.Errorf("clear current state: %w", err)
		}
	}
	for _, e := range w.StateDelta {
		if _, err := tx.ExecContext(ctx, s.q(upsertCurrentStateSQL),
			w.Event.RoomID, e.Type, e.StateKey, e.EventID, e.Membership); err != nil {
			return fmt.Errorf("set current state %s/%s: %w", e.Type, e.StateKey, err)
		}
	}
	return nil
}

const (
	deleteExtremitiesSQL = `DELETE FROM roomserver_extremities WHERE room_id = ?`
	insertExtremitySQL   = `
		INSERT INTO roomserver_extremities (room_id, event_id) VALUES (?, ?)
		ON CONFLICT (room_id, event_id) DO NOTHING`
)

func (s *Rooms) replaceExtremities(ctx context.Context, tx *sql.Tx, roomID string, extremities []string) error {
	if len(extremities) == 0 {
		return nil
	}
	if _, err := tx.ExecContext(ctx, s.q(deleteExtremitiesSQL), roomID); err != nil {
		return fmt.Errorf("clear forward extremities: %w", err)
	}
	for _, id := range extremities {
		if _, err := tx.ExecContext(ctx, s.q(insertExtremitySQL), roomID, id); err != nil {
			return fmt.Errorf("insert forward extremity %s: %w", id, err)
		}
	}
	return nil
}

const updateProfileSQL = `
	UPDATE roomserver_rooms SET
		name = ?, topic = ?, avatar_url = ?, canonical_alias = ?,
		join_rule = ?, history_visibility = ?, guest_access = ?, room_type = ?
	WHERE room_id = ?`

func (s *Rooms) updateProfile(ctx context.Context, tx *sql.Tx, roomID string, p storage.RoomProfile) error {
	_, err := tx.ExecContext(ctx, s.q(updateProfileSQL),
		p.Name, p.Topic, p.AvatarURL, p.CanonicalAlias,
		p.JoinRule, p.HistoryVisibility, p.GuestAccess, p.RoomType, roomID)
	if err != nil {
		return fmt.Errorf("update room profile: %w", err)
	}
	return nil
}

const recountMembersSQL = `
	UPDATE roomserver_rooms SET joined_members = (
		SELECT COUNT(*) FROM roomserver_current_state
		WHERE room_id = ? AND membership = 'join'
	) WHERE room_id = ?`

func (s *Rooms) recountMembers(ctx context.Context, tx *sql.Tx, roomID string) error {
	if _, err := tx.ExecContext(ctx, s.q(recountMembersSQL), roomID, roomID); err != nil {
		return fmt.Errorf("recount joined members: %w", err)
	}
	return nil
}

const insertTransactionSQL = `
	INSERT INTO roomserver_transactions (localpart, device_id, txn_id, event_id, room_id, sent_at_ms)
	VALUES (?, ?, ?, ?, ?, ?)
	ON CONFLICT (localpart, device_id, txn_id) DO NOTHING`

func (s *Rooms) recordTransaction(ctx context.Context, tx *sql.Tx, ref storage.TransactionRef, eventID, roomID string, atMS int64) error {
	_, err := tx.ExecContext(ctx, s.q(insertTransactionSQL),
		ref.Localpart, ref.DeviceID, ref.TxnID, eventID, roomID, atMS)
	if err != nil {
		return fmt.Errorf("record transaction %q: %w", ref.TxnID, err)
	}
	return nil
}

// eventColumns is the select list every event read shares, in the order
// scanEvent expects.
const eventColumns = `
	e.room_id, e.event_id, e.event_type, e.state_key, e.sender,
	e.depth, e.origin_server_ts, e.event_json, e.stream_ordering, e.redacted_by, e.outlier`

// scanEvent reads one row of eventColumns.
func scanEvent(scan func(...any) error) (storage.StoredEvent, error) {
	var (
		e        storage.StoredEvent
		stateKey sql.NullString
		body     string
	)
	err := scan(&e.RoomID, &e.EventID, &e.Type, &stateKey, &e.Sender,
		&e.Depth, &e.OriginServerTS, &body, &e.StreamOrdering, &e.RedactedBy, &e.Outlier)
	if err != nil {
		return storage.StoredEvent{}, err
	}
	if stateKey.Valid {
		e.StateKey = &stateKey.String
	}
	e.JSON = []byte(body)
	return e, nil
}

// scanEvents drains a result set of eventColumns rows.
func scanEvents(rows *sql.Rows) ([]storage.StoredEvent, error) {
	defer func() { _ = rows.Close() }()

	var out []storage.StoredEvent
	for rows.Next() {
		e, err := scanEvent(rows.Scan)
		if err != nil {
			return nil, err
		}
		out = append(out, e)
	}
	return out, rows.Err()
}

var selectEventSQL = `
	SELECT ` + eventColumns + `
	FROM roomserver_events e WHERE e.room_id = ? AND e.event_id = ?`

// Event implements storage.RoomStore.
func (s *Rooms) Event(ctx context.Context, roomID, eventID string) (storage.StoredEvent, error) {
	e, err := scanEvent(s.db.QueryRowContext(ctx, s.q(selectEventSQL), roomID, eventID).Scan)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return storage.StoredEvent{}, fmt.Errorf("event %q: %w", eventID, storage.ErrNoSuchEvent)
	case err != nil:
		return storage.StoredEvent{}, fmt.Errorf("read event %q: %w", eventID, err)
	}
	return e, nil
}

// EventsByID implements storage.RoomStore.
func (s *Rooms) EventsByID(ctx context.Context, roomID string, eventIDs []string) ([]storage.StoredEvent, error) {
	if len(eventIDs) == 0 {
		return nil, nil
	}
	query := `SELECT ` + eventColumns + `
		FROM roomserver_events e
		WHERE e.room_id = ? AND e.event_id IN (` + placeholders(len(eventIDs)) + `)`

	args := make([]any, 0, len(eventIDs)+1)
	args = append(args, roomID)
	for _, id := range eventIDs {
		args = append(args, id)
	}

	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("read events by ID: %w", err)
	}
	events, err := scanEvents(rows)
	if err != nil {
		return nil, fmt.Errorf("read events by ID: %w", err)
	}
	return events, nil
}

const selectCurrentStateSQL = `
	SELECT event_type, state_key, event_id, membership
	FROM roomserver_current_state WHERE room_id = ?
	ORDER BY event_type, state_key`

// CurrentState implements storage.RoomStore.
func (s *Rooms) CurrentState(ctx context.Context, roomID string) ([]storage.StateEntry, error) {
	rows, err := s.db.QueryContext(ctx, s.q(selectCurrentStateSQL), roomID)
	if err != nil {
		return nil, fmt.Errorf("read current state of %q: %w", roomID, err)
	}
	entries, err := scanStateEntries(rows)
	if err != nil {
		return nil, fmt.Errorf("read current state of %q: %w", roomID, err)
	}
	return entries, nil
}

func scanStateEntries(rows *sql.Rows) ([]storage.StateEntry, error) {
	defer func() { _ = rows.Close() }()

	var out []storage.StateEntry
	for rows.Next() {
		var e storage.StateEntry
		if err := rows.Scan(&e.Type, &e.StateKey, &e.EventID, &e.Membership); err != nil {
			return nil, err
		}
		out = append(out, e)
	}
	return out, rows.Err()
}

var selectCurrentStateEventsSQL = `
	SELECT ` + eventColumns + `
	FROM roomserver_current_state cs
	JOIN roomserver_events e ON e.event_id = cs.event_id
	WHERE cs.room_id = ?
	ORDER BY cs.event_type, cs.state_key`

// CurrentStateEvents implements storage.RoomStore.
func (s *Rooms) CurrentStateEvents(ctx context.Context, roomID string) ([]storage.StoredEvent, error) {
	rows, err := s.db.QueryContext(ctx, s.q(selectCurrentStateEventsSQL), roomID)
	if err != nil {
		return nil, fmt.Errorf("read current state events of %q: %w", roomID, err)
	}
	events, err := scanEvents(rows)
	if err != nil {
		return nil, fmt.Errorf("read current state events of %q: %w", roomID, err)
	}
	return events, nil
}

var selectCurrentStateEventSQL = `
	SELECT ` + eventColumns + `
	FROM roomserver_current_state cs
	JOIN roomserver_events e ON e.event_id = cs.event_id
	WHERE cs.room_id = ? AND cs.event_type = ? AND cs.state_key = ?`

// CurrentStateEvent implements storage.RoomStore.
func (s *Rooms) CurrentStateEvent(ctx context.Context, roomID, eventType, stateKey string) (storage.StoredEvent, error) {
	e, err := scanEvent(s.db.QueryRowContext(ctx,
		s.q(selectCurrentStateEventSQL), roomID, eventType, stateKey).Scan)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return storage.StoredEvent{}, fmt.Errorf("state %s/%s of %q: %w",
			eventType, stateKey, roomID, storage.ErrNoSuchEvent)
	case err != nil:
		return storage.StoredEvent{}, fmt.Errorf("read state %s/%s of %q: %w",
			eventType, stateKey, roomID, err)
	}
	return e, nil
}

const selectStateBeforeSQL = `
	SELECT se.event_type, se.state_key, se.event_id, ''
	FROM roomserver_events e
	JOIN roomserver_state_entries se ON se.snapshot_id = e.state_before_id
	WHERE e.room_id = ? AND e.event_id = ?
	ORDER BY se.event_type, se.state_key`

// StateBefore implements storage.RoomStore.
//
// The membership column is not stored on a snapshot entry: it exists on
// current state to make a membership lookup an index hit, and a historical
// snapshot is only ever read as a whole and then parsed anyway.
func (s *Rooms) StateBefore(ctx context.Context, roomID, eventID string) ([]storage.StateEntry, error) {
	rows, err := s.db.QueryContext(ctx, s.q(selectStateBeforeSQL), roomID, eventID)
	if err != nil {
		return nil, fmt.Errorf("read the state before %q: %w", eventID, err)
	}
	entries, err := scanStateEntries(rows)
	if err != nil {
		return nil, fmt.Errorf("read the state before %q: %w", eventID, err)
	}
	return entries, nil
}

var selectExtremitiesSQL = `
	SELECT ` + eventColumns + `
	FROM roomserver_extremities x
	JOIN roomserver_events e ON e.event_id = x.event_id
	WHERE x.room_id = ?
	ORDER BY e.stream_ordering`

// Extremities implements storage.RoomStore.
func (s *Rooms) Extremities(ctx context.Context, roomID string) ([]storage.StoredEvent, error) {
	rows, err := s.db.QueryContext(ctx, s.q(selectExtremitiesSQL), roomID)
	if err != nil {
		return nil, fmt.Errorf("read the forward extremities of %q: %w", roomID, err)
	}
	events, err := scanEvents(rows)
	if err != nil {
		return nil, fmt.Errorf("read the forward extremities of %q: %w", roomID, err)
	}
	return events, nil
}

// Members implements storage.RoomStore.
func (s *Rooms) Members(ctx context.Context, roomID string, memberships []string) ([]storage.StoredEvent, error) {
	query := `SELECT ` + eventColumns + `
		FROM roomserver_current_state cs
		JOIN roomserver_events e ON e.event_id = cs.event_id
		WHERE cs.room_id = ? AND cs.event_type = 'm.room.member'`
	args := []any{roomID}
	if len(memberships) > 0 {
		query += ` AND cs.membership IN (` + placeholders(len(memberships)) + `)`
		for _, m := range memberships {
			args = append(args, m)
		}
	}
	query += ` ORDER BY cs.state_key`

	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("read the members of %q: %w", roomID, err)
	}
	events, err := scanEvents(rows)
	if err != nil {
		return nil, fmt.Errorf("read the members of %q: %w", roomID, err)
	}
	return events, nil
}

// MembershipEntries implements storage.RoomStore.
func (s *Rooms) MembershipEntries(ctx context.Context, roomID string, memberships []string) ([]storage.MembershipEntry, error) {
	query := `
		SELECT state_key, membership, event_id FROM roomserver_current_state
		WHERE room_id = ? AND event_type = 'm.room.member'`
	args := []any{roomID}
	if len(memberships) > 0 {
		query += ` AND membership IN (` + placeholders(len(memberships)) + `)`
		for _, m := range memberships {
			args = append(args, m)
		}
	}
	query += ` ORDER BY state_key`

	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("read the memberships of %q: %w", roomID, err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.MembershipEntry
	for rows.Next() {
		var m storage.MembershipEntry
		if err := rows.Scan(&m.UserID, &m.Membership, &m.EventID); err != nil {
			return nil, fmt.Errorf("read the memberships of %q: %w", roomID, err)
		}
		out = append(out, m)
	}
	return out, rows.Err()
}

const selectMembershipSQL = `
	SELECT state_key, membership, event_id FROM roomserver_current_state
	WHERE room_id = ? AND event_type = 'm.room.member' AND state_key = ?`

// Membership implements storage.RoomStore.
func (s *Rooms) Membership(ctx context.Context, roomID, userID string) (storage.MembershipEntry, error) {
	var m storage.MembershipEntry
	err := s.db.QueryRowContext(ctx, s.q(selectMembershipSQL), roomID, userID).
		Scan(&m.UserID, &m.Membership, &m.EventID)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		// Never a member is not an error: it is the same answer as having
		// left, which is what the rules above this ask for.
		return storage.MembershipEntry{UserID: userID}, nil
	case err != nil:
		return storage.MembershipEntry{}, fmt.Errorf("read the membership of %q in %q: %w", userID, roomID, err)
	}
	return m, nil
}

// RoomsForUser implements storage.RoomStore.
func (s *Rooms) RoomsForUser(ctx context.Context, userID string, memberships []string) ([]string, error) {
	query := `
		SELECT room_id FROM roomserver_current_state
		WHERE event_type = 'm.room.member' AND state_key = ?`
	args := []any{userID}
	if len(memberships) > 0 {
		query += ` AND membership IN (` + placeholders(len(memberships)) + `)`
		for _, m := range memberships {
			args = append(args, m)
		}
	}
	query += ` ORDER BY room_id`

	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("read the rooms of %q: %w", userID, err)
	}
	defer func() { _ = rows.Close() }()

	var out []string
	for rows.Next() {
		var roomID string
		if err := rows.Scan(&roomID); err != nil {
			return nil, fmt.Errorf("read the rooms of %q: %w", userID, err)
		}
		out = append(out, roomID)
	}
	return out, rows.Err()
}

const selectTransactionSQL = `
	SELECT event_id FROM roomserver_transactions
	WHERE localpart = ? AND device_id = ? AND txn_id = ?`

// TransactionEvent implements storage.RoomStore.
func (s *Rooms) TransactionEvent(ctx context.Context, ref storage.TransactionRef) (string, error) {
	var eventID string
	err := s.db.QueryRowContext(ctx, s.q(selectTransactionSQL),
		ref.Localpart, ref.DeviceID, ref.TxnID).Scan(&eventID)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return "", fmt.Errorf("transaction %q: %w", ref.TxnID, storage.ErrNoSuchEvent)
	case err != nil:
		return "", fmt.Errorf("read transaction %q: %w", ref.TxnID, err)
	}
	return eventID, nil
}

const insertAliasSQL = `
	INSERT INTO roomserver_aliases (alias, room_id, creator, created_at_ms)
	VALUES (?, ?, ?, ?)`

// CreateAlias implements storage.RoomStore.
func (s *Rooms) CreateAlias(ctx context.Context, alias storage.RoomAlias) error {
	_, err := s.db.ExecContext(ctx, s.q(insertAliasSQL),
		alias.Alias, alias.RoomID, alias.Creator, alias.CreatedAtMS)
	if err != nil {
		if s.isUnique(err) {
			return fmt.Errorf("create alias %q: %w", alias.Alias, storage.ErrAliasExists)
		}
		return fmt.Errorf("create alias %q: %w", alias.Alias, err)
	}
	return nil
}

const deleteAliasSQL = `DELETE FROM roomserver_aliases WHERE alias = ?`

// DeleteAlias implements storage.RoomStore.
func (s *Rooms) DeleteAlias(ctx context.Context, alias string) error {
	result, err := s.db.ExecContext(ctx, s.q(deleteAliasSQL), alias)
	if err != nil {
		return fmt.Errorf("delete alias %q: %w", alias, err)
	}
	return expectOneRow(result, alias, storage.ErrNoSuchAlias)
}

const selectAliasSQL = `
	SELECT alias, room_id, creator, created_at_ms
	FROM roomserver_aliases WHERE alias = ?`

// Alias implements storage.RoomStore.
func (s *Rooms) Alias(ctx context.Context, alias string) (storage.RoomAlias, error) {
	var out storage.RoomAlias
	err := s.db.QueryRowContext(ctx, s.q(selectAliasSQL), alias).
		Scan(&out.Alias, &out.RoomID, &out.Creator, &out.CreatedAtMS)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return storage.RoomAlias{}, fmt.Errorf("alias %q: %w", alias, storage.ErrNoSuchAlias)
	case err != nil:
		return storage.RoomAlias{}, fmt.Errorf("resolve alias %q: %w", alias, err)
	}
	return out, nil
}

const selectRoomAliasesSQL = `
	SELECT alias FROM roomserver_aliases WHERE room_id = ? ORDER BY alias`

// AliasesForRoom implements storage.RoomStore.
func (s *Rooms) AliasesForRoom(ctx context.Context, roomID string) ([]string, error) {
	rows, err := s.db.QueryContext(ctx, s.q(selectRoomAliasesSQL), roomID)
	if err != nil {
		return nil, fmt.Errorf("read the aliases of %q: %w", roomID, err)
	}
	defer func() { _ = rows.Close() }()

	var out []string
	for rows.Next() {
		var alias string
		if err := rows.Scan(&alias); err != nil {
			return nil, fmt.Errorf("read the aliases of %q: %w", roomID, err)
		}
		out = append(out, alias)
	}
	return out, rows.Err()
}

const setPublishedSQL = `UPDATE roomserver_rooms SET is_published = ? WHERE room_id = ?`

// SetPublished implements storage.RoomStore.
func (s *Rooms) SetPublished(ctx context.Context, roomID string, published bool) error {
	result, err := s.db.ExecContext(ctx, s.q(setPublishedSQL), published, roomID)
	if err != nil {
		return fmt.Errorf("set the published flag of %q: %w", roomID, err)
	}
	return expectOneRow(result, roomID, storage.ErrNoSuchRoom)
}

// roomSummaryColumns is the select list a room listing shares, in the order
// scanRoomSummary expects.
const roomSummaryColumns = `
	room_id, room_version, creator, is_published, created_at_ms,
	name, topic, avatar_url, canonical_alias,
	join_rule, history_visibility, guest_access, room_type, joined_members`

func scanRoomSummary(scan func(...any) error) (storage.RoomSummary, error) {
	var r storage.RoomSummary
	err := scan(&r.ID, &r.Version, &r.Creator, &r.Published, &r.CreatedAtMS,
		&r.Name, &r.Topic, &r.AvatarURL, &r.CanonicalAlias,
		&r.JoinRule, &r.HistoryVisibility, &r.GuestAccess, &r.RoomType, &r.JoinedMembers)
	return r, err
}

// PublishedRooms implements storage.RoomStore.
func (s *Rooms) PublishedRooms(ctx context.Context, filter storage.RoomFilter) ([]storage.RoomSummary, int, error) {
	return s.listRooms(ctx, filter, true)
}

// ListRooms implements storage.RoomStore.
func (s *Rooms) ListRooms(ctx context.Context, filter storage.RoomFilter) ([]storage.RoomSummary, int, error) {
	return s.listRooms(ctx, filter, false)
}

// listRooms is the body of both listings: the same query with and without
// the published condition.
//
// Rooms are ordered by descending member count and then by room ID. The
// second key is not decoration: paging a listing whose order is not total
// can show the same room on two pages and skip another, and ties in member
// count are the common case in a directory of small rooms.
func (s *Rooms) listRooms(ctx context.Context, filter storage.RoomFilter, publishedOnly bool) ([]storage.RoomSummary, int, error) {
	where := ` WHERE 1 = 1`
	var args []any
	if publishedOnly {
		where += ` AND is_published = ?`
		args = append(args, true)
	}
	if search := strings.TrimSpace(filter.Search); search != "" {
		pattern := "%" + strings.ToLower(search) + "%"
		where += ` AND (LOWER(room_id) LIKE ? OR LOWER(name) LIKE ?
			OR LOWER(topic) LIKE ? OR LOWER(canonical_alias) LIKE ?)`
		args = append(args, pattern, pattern, pattern, pattern)
	}

	var total int
	if err := s.db.QueryRowContext(ctx,
		s.q(`SELECT COUNT(*) FROM roomserver_rooms`+where), args...).Scan(&total); err != nil {
		return nil, 0, fmt.Errorf("count rooms: %w", err)
	}

	limit, offset := pageBounds(filter.Limit, filter.Offset, defaultRoomPageSize, maxRoomPageSize)
	query := `SELECT ` + roomSummaryColumns + ` FROM roomserver_rooms` + where +
		` ORDER BY joined_members DESC, room_id LIMIT ? OFFSET ?`

	rows, err := s.db.QueryContext(ctx, s.q(query), append(args, limit, offset)...)
	if err != nil {
		return nil, 0, fmt.Errorf("list rooms: %w", err)
	}
	defer func() { _ = rows.Close() }()

	out := make([]storage.RoomSummary, 0, limit)
	for rows.Next() {
		r, err := scanRoomSummary(rows.Scan)
		if err != nil {
			return nil, 0, fmt.Errorf("list rooms: %w", err)
		}
		out = append(out, r)
	}
	if err := rows.Err(); err != nil {
		return nil, 0, fmt.Errorf("list rooms: %w", err)
	}
	return out, total, nil
}

var selectRoomSummarySQL = `SELECT ` + roomSummaryColumns +
	` FROM roomserver_rooms WHERE room_id = ?`

// RoomSummary implements storage.RoomStore.
func (s *Rooms) RoomSummary(ctx context.Context, roomID string) (storage.RoomSummary, error) {
	r, err := scanRoomSummary(s.db.QueryRowContext(ctx, s.q(selectRoomSummarySQL), roomID).Scan)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return storage.RoomSummary{}, fmt.Errorf("room %q: %w", roomID, storage.ErrNoSuchRoom)
	case err != nil:
		return storage.RoomSummary{}, fmt.Errorf("read room %q: %w", roomID, err)
	}
	return r, nil
}

// purgeStatements are run in order, children before parents.
//
// Both engines are given ON DELETE CASCADE on these tables, but SQLite
// enforces foreign keys only when the pragma is on, and the transactions
// table has no key at all because it outlives the room it refers to. The
// explicit order is what makes a purge mean the same thing on both.
var purgeStatements = []string{
	`DELETE FROM roomserver_state_entries WHERE snapshot_id IN (
		SELECT snapshot_id FROM roomserver_state_snapshots WHERE room_id = ?)`,
	`DELETE FROM roomserver_state_snapshots WHERE room_id = ?`,
	`DELETE FROM roomserver_current_state WHERE room_id = ?`,
	`DELETE FROM roomserver_extremities WHERE room_id = ?`,
	`DELETE FROM roomserver_transactions WHERE room_id = ?`,
	`DELETE FROM roomserver_aliases WHERE room_id = ?`,
	`DELETE FROM roomserver_invite_state WHERE room_id = ?`,
	`DELETE FROM roomserver_events WHERE room_id = ?`,
	`DELETE FROM roomserver_rooms WHERE room_id = ?`,
}

// PurgeRoom implements storage.RoomStore.
func (s *Rooms) PurgeRoom(ctx context.Context, roomID string) error {
	err := s.tx(ctx, func(tx *sql.Tx) error {
		for i, stmt := range purgeStatements {
			result, err := tx.ExecContext(ctx, s.q(stmt), roomID)
			if err != nil {
				return err
			}
			// The last statement is the room row itself: no row there means
			// there was no room, and the whole purge rolls back so that the
			// caller's 404 is the truth.
			if i == len(purgeStatements)-1 {
				if err := expectOneRow(result, roomID, storage.ErrNoSuchRoom); err != nil {
					return err
				}
			}
		}
		return nil
	})
	if err != nil {
		return fmt.Errorf("purge room %q: %w", roomID, err)
	}
	return nil
}

const roomStatsSQL = `
	SELECT
		(SELECT COUNT(*) FROM roomserver_rooms),
		(SELECT COUNT(*) FROM roomserver_rooms WHERE is_published = ?),
		(SELECT COUNT(*) FROM roomserver_events)`

// RoomStats implements storage.RoomStore.
func (s *Rooms) RoomStats(ctx context.Context) (storage.RoomStats, error) {
	var stats storage.RoomStats
	err := s.db.QueryRowContext(ctx, s.q(roomStatsSQL), true).
		Scan(&stats.Rooms, &stats.Published, &stats.Events)
	if err != nil {
		return storage.RoomStats{}, fmt.Errorf("read room statistics: %w", err)
	}
	return stats, nil
}

// placeholders renders n comma-separated `?` placeholders for an IN clause.
func placeholders(n int) string {
	return strings.TrimSuffix(strings.Repeat("?, ", n), ", ")
}

// pageBounds clamps a caller's limit and offset.
func pageBounds(limit, offset, defaultLimit, maxLimit int) (int, int) {
	if limit <= 0 {
		limit = defaultLimit
	}
	if limit > maxLimit {
		limit = maxLimit
	}
	if offset < 0 {
		offset = 0
	}
	return limit, offset
}

// expectOneRow turns "the UPDATE or DELETE matched nothing" into notFound.
func expectOneRow(result sql.Result, subject string, notFound error) error {
	n, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("count the affected rows: %w", err)
	}
	if n == 0 {
		return fmt.Errorf("%q: %w", subject, notFound)
	}
	return nil
}

var selectEventByIDSQL = `
	SELECT ` + eventColumns + `
	FROM roomserver_events e WHERE e.event_id = ?`

// EventByID implements storage.RoomStore.
func (s *Rooms) EventByID(ctx context.Context, eventID string) (storage.StoredEvent, error) {
	e, err := scanEvent(s.db.QueryRowContext(ctx, s.q(selectEventByIDSQL), eventID).Scan)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return storage.StoredEvent{}, fmt.Errorf("event %q: %w", eventID, storage.ErrNoSuchEvent)
	case err != nil:
		return storage.StoredEvent{}, fmt.Errorf("read event %q: %w", eventID, err)
	}
	return e, nil
}

const selectJoinedMemberOnSQL = `
	SELECT 1 FROM roomserver_current_state
	WHERE room_id = ? AND event_type = 'm.room.member' AND membership = 'join'
		AND state_key LIKE ? ESCAPE '\'
	LIMIT 1`

// HasJoinedMemberOn implements storage.RoomStore.
//
// A user ID's server is everything after its first colon, and a localpart
// cannot contain one, so "ends with ':' plus the server name" is exactly
// "belongs to that server".
func (s *Rooms) HasJoinedMemberOn(ctx context.Context, roomID, serverName string) (bool, error) {
	var one int
	err := s.db.QueryRowContext(ctx, s.q(selectJoinedMemberOnSQL), roomID, `%:`+escapeLike(serverName)).Scan(&one)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return false, nil
	case err != nil:
		return false, fmt.Errorf("look for members of %q in %q: %w", serverName, roomID, err)
	}
	return true, nil
}